[PHI] Move lu to phi

paddle/fluid/operators/lu_op.cc

@@ -13,6 +13,11 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/lu_op.h"
+#include "paddle/fluid/framework/infershape_utils.h"
+#include "paddle/fluid/framework/op_registry.h"
+
+#include "paddle/phi/infermeta/backward.h"
+#include "paddle/phi/infermeta/unary.h"
 
 namespace paddle {
 namespace operators {
@@ -39,39 +44,6 @@ class LUOp : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
 
-  void InferShape(framework::InferShapeContext *context) const override {
-    OP_INOUT_CHECK(context->HasInput("X"), "Input", "X", "LU");
-    OP_INOUT_CHECK(context->HasOutput("Out"), "Output", "Out", "LU");
-    bool pivots = context->Attrs().Get<bool>("pivots");
-    auto x_dims = context->GetInputDim("X");
-    int x_rank = x_dims.size();
-    PADDLE_ENFORCE_GE(x_rank,
-                      2,
-                      platform::errors::InvalidArgument(
-                          "the rank of input must greater than 2"));
-    context->SetOutputDim("Out", x_dims);
-    int m = x_dims[x_rank - 1];
-    int n = x_dims[x_rank - 2];
-    int min_mn = std::min(m, n);
-    auto dims_vec = phi::vectorize(x_dims);
-    OP_INOUT_CHECK(context->HasOutput("Infos"), "Output", "Infos", "LU");
-    if (x_rank == 2) {
-      auto Infos_dim = std::vector<int>(1);
-      context->SetOutputDim("Infos", phi::make_ddim(Infos_dim));
-    } else {
-      auto Infos_dim =
-          std::vector<int>(dims_vec.begin(), dims_vec.begin() + x_rank - 2);
-      context->SetOutputDim("Infos", phi::make_ddim(Infos_dim));
-    }
-    if (pivots) {
-      OP_INOUT_CHECK(context->HasOutput("Pivots"), "Output", "Pivots", "LU");
-      auto Pivots_dim =
-          std::vector<int>(dims_vec.begin(), dims_vec.begin() + x_rank - 1);
-      Pivots_dim[x_rank - 2] = min_mn;
-      context->SetOutputDim("Pivots", phi::make_ddim(Pivots_dim));
-    }
-  }
-
  protected:
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext &ctx) const override {
@@ -99,57 +71,6 @@ class LUOpVarTypeInference : public framework::VarTypeInference {
   }
 };
 
-template <typename T>
-class LUKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const paddle::framework::ExecutionContext &ctx) const override {
-    auto pivots = ctx.Attr<bool>("pivots");
-    auto *xin = ctx.Input<framework::Tensor>("X");
-    auto *out = ctx.Output<framework::Tensor>("Out");
-    auto *IpivT = ctx.Output<framework::Tensor>("Pivots");
-    auto *InfoT = ctx.Output<framework::Tensor>("Infos");
-    PADDLE_ENFORCE_EQ(pivots,
-                      true,
-                      platform::errors::InvalidArgument(
-                          "lu without pivoting is not implemented on the CPU, "
-                          "but got pivots=False"));
-
-    math::DeviceIndependenceTensorOperations<phi::CPUContext, T> helper(ctx);
-    *out = helper.Transpose(*xin);
-
-    auto outdims = out->dims();
-    auto outrank = outdims.size();
-
-    int m = static_cast<int>(outdims[outrank - 1]);
-    int n = static_cast<int>(outdims[outrank - 2]);
-    int lda = std::max(1, m);
-
-    auto ipiv_dims = phi::slice_ddim(outdims, 0, outrank - 1);
-    ipiv_dims[outrank - 2] = std::min(m, n);
-    IpivT->Resize(ipiv_dims);
-    auto ipiv_data = IpivT->mutable_data<int>(ctx.GetPlace());
-
-    auto info_dims = phi::slice_ddim(outdims, 0, outrank - 2);
-    if (info_dims.size() == 0) {
-      info_dims = phi::make_ddim({1});
-    }
-    InfoT->Resize(info_dims);
-    auto info_data = InfoT->mutable_data<int>(ctx.GetPlace());
-
-    auto batchsize = product(info_dims);
-    batchsize = std::max(static_cast<int>(batchsize), 1);
-    auto out_data = out->mutable_data<T>(ctx.GetPlace());
-    for (int b = 0; b < batchsize; b++) {
-      auto out_data_item = &out_data[b * m * n];
-      int *info_data_item = &info_data[b];
-      int *ipiv_data_item = &ipiv_data[b * std::min(m, n)];
-      phi::funcs::lapackLu<T>(
-          m, n, out_data_item, lda, ipiv_data_item, info_data_item);
-    }
-    *out = helper.Transpose(*out);
-  }
-};
-
 template <typename T>
 class LUOpGradMaker : public framework::SingleGradOpMaker<T> {
  public:
@@ -184,23 +105,6 @@ class LUGradOp : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
 
-  void InferShape(framework::InferShapeContext *ctx) const override {
-    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "lu");
-    OP_INOUT_CHECK(ctx->HasInput("Out"), "Input", "Out", "lu");
-    OP_INOUT_CHECK(ctx->HasInput("Pivots"), "Input", "Pivots", "lu");
-    OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")),
-                   "Input",
-                   "Out@GRAD",
-                   "lu");
-
-    auto x_dims = ctx->GetInputDim("X");
-    auto x_grad_name = framework::GradVarName("X");
-
-    if (ctx->HasOutput(x_grad_name)) {
-      ctx->SetOutputDim(x_grad_name, x_dims);
-    }
-  }
-
  protected:
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext &ctx) const override {
@@ -219,19 +123,21 @@ DECLARE_INPLACE_OP_INFERER(LUGradOpInplaceInferer,
 namespace ops = paddle::operators;
 namespace plat = paddle::platform;
 
+DECLARE_INFER_SHAPE_FUNCTOR(lu,
+                            LUInferMetaFunctor,
+                            PD_INFER_META(phi::LUInferMeta));
+DECLARE_INFER_SHAPE_FUNCTOR(lu_grad,
+                            LUGradInferMetaFunctor,
+                            PD_INFER_META(phi::LUGradInferMeta));
+
 REGISTER_OPERATOR(lu,
                   ops::LUOp,
                   ops::LUOpMaker,
                   ops::LUOpVarTypeInference,
                   ops::LUOpGradMaker<paddle::framework::OpDesc>,
                   ops::LUOpGradMaker<paddle::imperative::OpBase>,
-                  ops::LUOpInplaceInferer);
+                  LUInferMetaFunctor);
 REGISTER_OPERATOR(lu_grad,
                   ops::LUGradOp,
                   ops::LUGradOpVarTypeInference,
-                  ops::LUGradOpInplaceInferer);
-
-REGISTER_OP_CPU_KERNEL(lu, ops::LUKernel<float>, ops::LUKernel<double>);
-REGISTER_OP_CPU_KERNEL(lu_grad,
-                       ops::LUGradKernel<phi::CPUContext, float>,
-                       ops::LUGradKernel<phi::CPUContext, double>);
+                  LUGradInferMetaFunctor);