transfer the svd infer into phi infermeta

paddle/fluid/operators/svd_op.cc

@@ -17,64 +17,20 @@
 #include <unordered_map>
 #include <vector>
 
+#include "paddle/fluid/framework/infershape_utils.h"
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/phi/core/ddim.h"
+#include "paddle/phi/infermeta/unary.h"
 #ifdef PADDLE_WITH_MKLDNN
 #include "paddle/fluid/platform/mkldnn_helper.h"
 #endif
 
 namespace paddle {
 namespace operators {
 
-using DDim = framework::DDim;
-static DDim UDDim(const DDim& x_dim, int k) {
-  // get x_dim and return the ddim of U
-  auto x_vec = vectorize(x_dim);
-  x_vec[x_vec.size() - 1] = k;
-  return phi::make_ddim(x_vec);
-}
-static DDim VHDDim(const DDim& x_dim, int k) {
-  // get x_dim and return the ddim of U
-  auto x_vec = vectorize(x_dim);
-  x_vec[x_vec.size() - 2] = k;
-  return phi::make_ddim(x_vec);
-}
-static DDim SDDim(const DDim& x_dim, int k) {
-  // get x_dim and return the ddim of U
-  auto x_vec = vectorize(x_dim);
-  x_vec[x_vec.size() - 2] = k;
-  x_vec.erase(x_vec.end() - 1);  // rank - 1
-  return phi::make_ddim(x_vec);
-}
-
 class SvdOp : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
-
-  void InferShape(framework::InferShapeContext* ctx) const override {
-    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "svd");
-    OP_INOUT_CHECK(ctx->HasOutput("U"), "Output", "U", "svd");
-    OP_INOUT_CHECK(ctx->HasOutput("VH"), "Output", "VH", "svd");
-    OP_INOUT_CHECK(ctx->HasOutput("S"), "Output", "S", "svd");
-
-    auto in_dims = ctx->GetInputDim("X");
-    int x_rank = in_dims.size();
-    PADDLE_ENFORCE_GE(in_dims.size(),
-                      2,
-                      platform::errors::InvalidArgument(
-                          "the rank of input must greater than 2"));
-    int m = in_dims[x_rank - 2];
-    int n = in_dims[x_rank - 1];
-    int k = std::min(m, n);
-    const bool full_uv = ctx->Attrs().Get<bool>("full_matrices");
-    ctx->SetOutputDim("U", !full_uv ? UDDim(in_dims, k) : UDDim(in_dims, m));
-    ctx->SetOutputDim("VH", !full_uv ? VHDDim(in_dims, k) : VHDDim(in_dims, n));
-    ctx->SetOutputDim("S", SDDim(in_dims, k));
-
-    ctx->ShareLoD("X", /*->*/ "U");
-    ctx->ShareLoD("X", /*->*/ "VH");
-    ctx->ShareLoD("X", /*->*/ "S");
-  }
 };
 
 class SvdOpMaker : public framework::OpProtoAndCheckerMaker {
@@ -159,10 +115,15 @@ class SvdGradMaker : public framework::SingleGradOpMaker<T> {
 
 namespace ops = paddle::operators;
 
+DECLARE_INFER_SHAPE_FUNCTOR(svd,
+                            SvdInferShapeFunctor,
+                            PD_INFER_META(phi::SvdInferMeta));
+
 REGISTER_OPERATOR(svd,
                   ops::SvdOp,
                   ops::SvdOpMaker,
                   ops::SvdGradMaker<paddle::framework::OpDesc>,
-                  ops::SvdGradMaker<paddle::imperative::OpBase>);
+                  ops::SvdGradMaker<paddle::imperative::OpBase>,
+                  SvdInferShapeFunctor);
 
 REGISTER_OPERATOR(svd_grad, ops::SvdGradOp);

paddle/phi/infermeta/unary.cc

@@ -26,6 +26,7 @@ limitations under the License. */
 #include "paddle/phi/kernels/funcs/pooling.h"
 #include "paddle/phi/kernels/funcs/slice_utils.h"
 #include "paddle/phi/kernels/funcs/strided_slice.h"
+#include "paddle/phi/kernels/funcs/svd.h"
 #include "paddle/phi/kernels/funcs/unfold_functor.h"
 #include "paddle/phi/kernels/funcs/unsqueeze.h"
 #include "paddle/phi/kernels/impl/einsum_impl.h"
@@ -2674,6 +2675,33 @@ void SumRawInferMeta(const MetaTensor& x,
   out->set_layout(x.layout());
 }
 
+void SvdInferMeta(const MetaTensor& x,
+                  bool full_matrices,
+                  MetaTensor* u,
+                  MetaTensor* s,
+                  MetaTensor* vh) {
+  auto in_dims = x.dims();
+  int x_rank = in_dims.size();
+  PADDLE_ENFORCE_GE(
+      in_dims.size(),
+      2,
+      phi::errors::InvalidArgument("the rank of input must greater than 2"));
+  int m = in_dims[x_rank - 2];
+  int n = in_dims[x_rank - 1];
+  int k = std::min(m, n);
+  u->set_dims(!full_matrices ? funcs::UDDim(in_dims, k)
+                             : funcs::UDDim(in_dims, m));
+  vh->set_dims(!full_matrices ? funcs::VHDDim(in_dims, k)
+                              : funcs::VHDDim(in_dims, n));
+  s->set_dims(funcs::SDDim(in_dims, k));
+  u->share_lod(x);
+  vh->share_lod(x);
+  s->share_lod(x);
+  u->set_dtype(x.dtype());
+  vh->set_dtype(x.dtype());
+  s->set_dtype(x.dtype());
+}
+
 void TemporalShiftInferMeta(const MetaTensor& x,
                             int seg_num,
                             float shift_ratio,

paddle/phi/infermeta/unary.h

@@ -383,6 +383,12 @@ void SumRawInferMeta(const MetaTensor& x,
                      DataType dtype,
                      MetaTensor* out);
 
+void SvdInferMeta(const MetaTensor& x,
+                  bool full_matrices,
+                  MetaTensor* u,
+                  MetaTensor* s,
+                  MetaTensor* vh);
+
 void TemporalShiftInferMeta(const MetaTensor& x,
                             int seg_num,
                             float shift_ratio,

paddle/phi/kernels/funcs/svd.h

@@ -0,0 +1,49 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/core/ddim.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+// TODO(paddle-dev): Remove this file when we can call related Kernel directly
+
+namespace phi {
+namespace funcs {
+
+using DDim = phi::DDim;
+static DDim UDDim(const DDim& x_dim, int k) {
+  // get x_dim and return the ddim of U
+  auto x_vec = vectorize(x_dim);
+  x_vec[x_vec.size() - 1] = k;
+  return phi::make_ddim(x_vec);
+}
+
+static DDim VHDDim(const DDim& x_dim, int k) {
+  // get x_dim and return the ddim of U
+  auto x_vec = vectorize(x_dim);
+  x_vec[x_vec.size() - 2] = k;
+  return phi::make_ddim(x_vec);
+}
+
+static DDim SDDim(const DDim& x_dim, int k) {
+  // get x_dim and return the ddim of U
+  auto x_vec = vectorize(x_dim);
+  x_vec[x_vec.size() - 2] = k;
+  x_vec.erase(x_vec.end() - 1);  // rank - 1
+  return phi::make_ddim(x_vec);
+}
+
+}  // namespace funcs
+}  // namespace phi