[phi]migrate increment addmm multinomial cholesky InferShapes to phi

paddle/fluid/operators/addmm_op.cc

@@ -16,7 +16,10 @@ limitations under the License. */
 #include <string>
 #include <unordered_map>
 #include <vector>
+#include "paddle/fluid/framework/infershape_utils.h"
 #include "paddle/fluid/framework/op_registry.h"
+#include "paddle/phi/core/infermeta_utils.h"
+#include "paddle/phi/infermeta/ternary.h"
 #ifdef PADDLE_WITH_MKLDNN
 #include "paddle/fluid/platform/mkldnn_helper.h"
 #endif
@@ -33,85 +36,6 @@ class AddMMOp : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
 
-  void InferShape(framework::InferShapeContext* ctx) const override {
-    PADDLE_ENFORCE_EQ(ctx->HasInput("Input"), true,
-                      platform::errors::NotFound(
-                          "Input(Input) of AddMMOp should not be null."));
-    PADDLE_ENFORCE_EQ(
-        ctx->HasInput("X"), true,
-        platform::errors::NotFound("Input(X) of AddMMOp should not be null."));
-    PADDLE_ENFORCE_EQ(
-        ctx->HasInput("Y"), true,
-        platform::errors::NotFound("Input(Y) of AddMMOp should not be null."));
-    PADDLE_ENFORCE_EQ(ctx->HasOutput("Out"), true,
-                      platform::errors::NotFound(
-                          "Output(Out) of AddMMOp should not be null."));
-
-    auto input_dims = ctx->GetInputDim("Input");
-    auto x_dims = ctx->GetInputDim("X");
-    auto y_dims = ctx->GetInputDim("Y");
-
-    auto ndim_input = input_dims.size();
-    auto ndim_x = x_dims.size();
-    auto ndim_y = y_dims.size();
-
-    float alpha = ctx->Attrs().Get<float>("Alpha");
-    float beta = ctx->Attrs().Get<float>("Beta");
-
-    VLOG(3) << "addmm operator input.shape=" << input_dims
-            << " x.shape=" << x_dims << " y.shape=" << y_dims
-            << " beta=" << beta << " alpha=" << alpha
-            << " ndim_input=" << ndim_input << " ndim_x=" << ndim_x
-            << " ndim_y=" << ndim_y;
-
-    PADDLE_ENFORCE_NE(phi::product(input_dims), 0,
-                      platform::errors::PreconditionNotMet(
-                          "The Input variable Input(%s) has not "
-                          "been initialized. You may need to confirm "
-                          "if you put exe.run(startup_program) "
-                          "after optimizer.minimize function.",
-                          ctx->Inputs("Input").front()));
-
-    PADDLE_ENFORCE_NE(phi::product(x_dims), 0,
-                      platform::errors::PreconditionNotMet(
-                          "The Input variable X(%s) has not "
-                          "been initialized. You may need to confirm "
-                          "if you put exe.run(startup_program) "
-                          "after optimizer.minimize function.",
-                          ctx->Inputs("X").front()));
-
-    PADDLE_ENFORCE_NE(phi::product(y_dims), 0,
-                      platform::errors::PreconditionNotMet(
-                          "The Input variable Y(%s) has not "
-                          "been initialized. You may need to confirm "
-                          "if you put exe.run(startup_program) "
-                          "after optimizer.minimize function.",
-                          ctx->Inputs("Y").front()));
-    // dim check
-    PADDLE_ENFORCE_EQ(ndim_input, 2,
-                      platform::errors::InvalidArgument(
-                          "The input tensor input's dimension must be 2. "
-                          "But received input's dimension = [%s].",
-                          ndim_input));
-    PADDLE_ENFORCE_EQ(ndim_x, 2,
-                      platform::errors::InvalidArgument(
-                          "The input tensor x's dimension must be 2. "
-                          "But received x's dimension = [%s].",
-                          ndim_x));
-    PADDLE_ENFORCE_EQ(ndim_y, 2,
-                      platform::errors::InvalidArgument(
-                          "The input tensor y's dimension must be 2. "
-                          "But received y's dimension = [%s].",
-                          ndim_y));
-
-    std::vector<int64_t> output_dims;
-    output_dims.push_back(x_dims[0]);
-    output_dims.push_back(y_dims[1]);
-
-    ctx->SetOutputDim("Out", phi::make_ddim(output_dims));
-    ctx->ShareLoD("Input", /*->*/ "Out");
-  }
-
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext& ctx) const {
     framework::LibraryType library = framework::LibraryType::kPlain;
@@ -223,9 +147,11 @@ class AddMMOpGradMaker : public framework::SingleGradOpMaker<T> {
 }  // namespace paddle
 
 namespace ops = paddle::operators;
-
+DELCARE_INFER_SHAPE_FUNCTOR(addmm, AddmmInferShapeFunctor,
+                            PT_INFER_META(phi::AddmmInferMeta));
 REGISTER_OPERATOR(addmm, ops::AddMMOp, ops::AddMMOpMaker,
                   ops::AddMMOpGradMaker<paddle::framework::OpDesc>,
-                  ops::AddMMOpGradMaker<paddle::imperative::OpBase>);
+                  ops::AddMMOpGradMaker<paddle::imperative::OpBase>,
+                  AddmmInferShapeFunctor);
 
 REGISTER_OPERATOR(addmm_grad, ops::AddMMGradOp);

paddle/fluid/operators/cholesky_op.cc

@@ -12,7 +12,10 @@ 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. */
 
+#include "paddle/fluid/framework/infershape_utils.h"
 #include "paddle/fluid/framework/op_registry.h"
+#include "paddle/phi/core/infermeta_utils.h"
+#include "paddle/phi/infermeta/unary.h"
 
 namespace paddle {
 namespace operators {
@@ -23,26 +26,6 @@ using framework::Tensor;
 class CholeskyOp : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
-
-  void InferShape(framework::InferShapeContext* ctx) const override {
-    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "Cholesky");
-    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "Cholesky");
-    auto dims = ctx->GetInputDim("X");
-    auto rank = dims.size();
-    PADDLE_ENFORCE_GE(rank, 2,
-                      platform::errors::InvalidArgument(
-                          "The Input(X) should have at least 2 dimensions. But "
-                          "received a %d dimension tensor.",
-                          rank));
-    PADDLE_ENFORCE_EQ(
-        dims[rank - 2], dims[rank - 1],
-        platform::errors::InvalidArgument(
-            "The inner-most 2 dimensions of Input(X) all should be symmetric "
-            "positive-definite matrices and have the same size. But received "
-            "X's shape[-2] = %d and shape[-1] = %d.",
-            dims[rank - 2], dims[rank - 1]));
-    ctx->SetOutputDim("Out", ctx->GetInputDim("X"));
-  }
 };
 
 class CholeskyOpMaker : public framework::OpProtoAndCheckerMaker {
@@ -107,7 +90,10 @@ class CholeskyGradOpMaker : public framework::SingleGradOpMaker<T> {
 }  // namespace paddle
 
 namespace ops = paddle::operators;
+DELCARE_INFER_SHAPE_FUNCTOR(cholesky, CholeskyInferShapeFunctor,
+                            PT_INFER_META(phi::CholeskyInferMeta));
 REGISTER_OPERATOR(cholesky, ops::CholeskyOp, ops::CholeskyOpMaker,
                   ops::CholeskyGradOpMaker<paddle::framework::OpDesc>,
-                  ops::CholeskyGradOpMaker<paddle::imperative::OpBase>);
+                  ops::CholeskyGradOpMaker<paddle::imperative::OpBase>,
+                  CholeskyInferShapeFunctor);
 REGISTER_OPERATOR(cholesky_grad, ops::CholeskyGradOp);

paddle/fluid/operators/increment_op.cc

@@ -12,7 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "paddle/fluid/framework/infershape_utils.h"
 #include "paddle/fluid/framework/op_registry.h"
+#include "paddle/phi/core/infermeta_utils.h"
+#include "paddle/phi/infermeta/unary.h"
 
 namespace paddle {
 namespace framework {
@@ -37,18 +40,6 @@ class IncrementOp : public framework::OperatorWithKernel {
               const framework::AttributeMap &attrs)
       : OperatorWithKernel(type, inputs, outputs, attrs) {}
 
-  void InferShape(framework::InferShapeContext *ctx) const override {
-    PADDLE_ENFORCE_EQ(phi::product(ctx->GetInputDim("X")), 1UL,
-                      platform::errors::InvalidArgument(
-                          "The number of elements in Input(X) should be 1."
-                          "Now the number is %d.",
-                          phi::product(ctx->GetInputDim("X"))));
-    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "increment");
-    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "increment");
-    ctx->SetOutputDim("Out", ctx->GetInputDim("X"));
-    ctx->ShareLoD("X", "Out");
-  }
-
  protected:
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext &ctx) const override {
@@ -96,6 +87,9 @@ class IncrementGradOpMaker : public framework::SingleGradOpMaker<T> {
 }  // namespace paddle
 
 namespace ops = paddle::operators;
+DELCARE_INFER_SHAPE_FUNCTOR(increment, IncrementInferShapeFunctor,
+                            PT_INFER_META(phi::IncrementInferMeta));
 REGISTER_OPERATOR(increment, ops::IncrementOp, ops::IncrementOpMaker,
                   ops::IncrementGradOpMaker<paddle::framework::OpDesc>,
-                  ops::IncrementGradOpMaker<paddle::imperative::OpBase>);
+                  ops::IncrementGradOpMaker<paddle::imperative::OpBase>,
+                  IncrementInferShapeFunctor);

paddle/fluid/operators/multinomial_op.cc

@@ -16,10 +16,11 @@ limitations under the License. */
 #include <string>
 #include <vector>
 
-#include "paddle/fluid/framework/generator.h"
+#include "paddle/fluid/framework/infershape_utils.h"
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/framework/operator.h"
-#include "paddle/fluid/operators/common_infer_shape_functions.h"
+#include "paddle/phi/core/infermeta_utils.h"
+#include "paddle/phi/infermeta/unary.h"
 
 namespace paddle {
 namespace operators {
@@ -45,46 +46,17 @@ This OP returns a Tensor filled with the sampled categoris according to Multinom
 class MultinomialOp : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
-
-  void InferShape(framework::InferShapeContext *ctx) const override {
-    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "Multinomial");
-    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "Multinomial");
-
-    auto x_dim = ctx->GetInputDim("X");
-    int64_t x_rank = x_dim.size();
-    PADDLE_ENFORCE_GT(x_rank, 0,
-                      platform::errors::InvalidArgument(
-                          "The number of dimensions of the input probability "
-                          "distribution should be > 0, but got %d.",
-                          x_rank));
-    PADDLE_ENFORCE_LE(x_rank, 2,
-                      platform::errors::InvalidArgument(
-                          "The number of dimensions of the input probability "
-                          "distribution should be <= 2, but got %d.",
-                          x_rank));
-
-    std::vector<int64_t> out_dims(x_rank);
-    for (int64_t i = 0; i < x_rank - 1; i++) {
-      out_dims[i] = x_dim[i];
-    }
-
-    int64_t num_samples = ctx->Attrs().Get<int>("num_samples");
-    PADDLE_ENFORCE_GT(
-        num_samples, 0,
-        platform::errors::InvalidArgument(
-            "The number of samples should be > 0, but got %d.", num_samples));
-    out_dims[x_rank - 1] = num_samples;
-
-    ctx->SetOutputDim("Out", phi::make_ddim(out_dims));
-  }
 };
 
 }  // namespace operators
 }  // namespace paddle
 
 namespace ops = paddle::operators;
 namespace plat = paddle::platform;
+DELCARE_INFER_SHAPE_FUNCTOR(multinomial, MultinomialInferShapeFunctor,
+                            PT_INFER_META(phi::MultinomialInferMeta));
 REGISTER_OPERATOR(
     multinomial, ops::MultinomialOp, ops::MultinomialOpMaker,
     paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
-    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
+    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>,
+    MultinomialInferShapeFunctor);

paddle/phi/infermeta/CMakeLists.txt

@@ -1,2 +1,2 @@
-cc_library(infermeta SRCS nullary.cc unary.cc binary.cc multiary.cc DEPS convert_utils meta_tensor infermeta_utils)
+cc_library(infermeta SRCS nullary.cc unary.cc binary.cc ternary.cc multiary.cc DEPS convert_utils meta_tensor infermeta_utils)
 cc_library(backward_infermeta SRCS backward.cc DEPS meta_tensor convert_utils)

paddle/phi/infermeta/ternary.cc

@@ -0,0 +1,92 @@
+/* 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. */
+
+#include "paddle/phi/infermeta/ternary.h"
+#include "paddle/phi/core/ddim.h"
+#include "paddle/phi/kernels/funcs/common_shape.h"
+
+namespace phi {
+
+void AddmmInferMeta(const MetaTensor& input,
+                    const MetaTensor& x,
+                    const MetaTensor& y,
+                    float alpha,
+                    float beta,
+                    MetaTensor* out) {
+  auto input_dims = input.dims();
+  auto x_dims = x.dims();
+  auto y_dims = y.dims();
+
+  auto ndim_input = input_dims.size();
+  auto ndim_x = x_dims.size();
+  auto ndim_y = y_dims.size();
+
+  VLOG(3) << "addmm operator input.shape=" << input_dims
+          << " x.shape=" << x_dims << " y.shape=" << y_dims << " beta=" << beta
+          << " alpha=" << alpha << " ndim_input=" << ndim_input
+          << " ndim_x=" << ndim_x << " ndim_y=" << ndim_y;
+
+  PADDLE_ENFORCE_NE(
+      product(input_dims),
+      0,
+      errors::PreconditionNotMet("The Input variable 'input' has not "
+                                 "been initialized. You may need to confirm "
+                                 "if you put exe.run(startup_program) "
+                                 "after optimizer.minimize function."));
+
+  PADDLE_ENFORCE_NE(
+      product(x_dims),
+      0,
+      errors::PreconditionNotMet("The Input variable 'x' has not "
+                                 "been initialized. You may need to confirm "
+                                 "if you put exe.run(startup_program) "
+                                 "after optimizer.minimize function."));
+
+  PADDLE_ENFORCE_NE(
+      product(y_dims),
+      0,
+      errors::PreconditionNotMet("The Input variable 'y' has not "
+                                 "been initialized. You may need to confirm "
+                                 "if you put exe.run(startup_program) "
+                                 "after optimizer.minimize function."));
+  // dim check
+  PADDLE_ENFORCE_EQ(
+      ndim_input,
+      2,
+      errors::InvalidArgument("The input tensor input's dimension must be 2. "
+                              "But received input's dimension = [%s].",
+                              ndim_input));
+  PADDLE_ENFORCE_EQ(
+      ndim_x,
+      2,
+      errors::InvalidArgument("The input tensor x's dimension must be 2. "
+                              "But received x's dimension = [%s].",
+                              ndim_x));
+  PADDLE_ENFORCE_EQ(
+      ndim_y,
+      2,
+      errors::InvalidArgument("The input tensor y's dimension must be 2. "
+                              "But received y's dimension = [%s].",
+                              ndim_y));
+
+  std::vector<int64_t> output_dims;
+  output_dims.push_back(x_dims[0]);
+  output_dims.push_back(y_dims[1]);
+
+  out->set_dims(make_ddim(output_dims));
+  out->share_lod(input);
+  out->set_dtype(input.dtype());
+}
+
+}  // namespace phi