[PHI] move diag_embed to phi.

paddle/fluid/operators/diag_embed_op.cc

@@ -12,89 +12,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/operators/diag_embed_op.h"
+#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 {
 
 class DiagEmbedOp : 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 of DiagEmbedOp is not found."));
-
-    PADDLE_ENFORCE_EQ(
-        ctx->HasOutput("Out"),
-        true,
-        platform::errors::NotFound("Output of DiagEmbedOp is not found."));
-
-    int offset = ctx->Attrs().Get<int>("offset");
-    int dim1 = ctx->Attrs().Get<int>("dim1");
-    int dim2 = ctx->Attrs().Get<int>("dim2");
-
-    auto x_dims = ctx->GetInputDim("Input");
-
-    PADDLE_ENFORCE_GE(
-        dim1,
-        -(x_dims.size() + 1),
-        platform::errors::OutOfRange(
-            "Dim1 is out of range (expected to be in range of [%ld, "
-            "%ld], but got %ld).",
-            -(x_dims.size() + 1),
-            x_dims.size(),
-            dim1));
-    PADDLE_ENFORCE_LE(
-        dim1,
-        x_dims.size(),
-        platform::errors::OutOfRange(
-            "Dim1 is out of range (expected to be in range of [%ld, "
-            "%ld], but got %ld).",
-            -(x_dims.size() + 1),
-            x_dims.size(),
-            dim1));
-
-    PADDLE_ENFORCE_GE(
-        dim2,
-        -(x_dims.size() + 1),
-        platform::errors::OutOfRange(
-            "Dim2 is out of range (expected to be in range of [%ld, "
-            "%ld], but got %ld).",
-            -(x_dims.size() + 1),
-            x_dims.size(),
-            dim2));
-    PADDLE_ENFORCE_LE(
-        dim2,
-        x_dims.size(),
-        platform::errors::OutOfRange(
-            "Dim2 is out of range (expected to be in range of [%ld, "
-            "%ld], but got %ld).",
-            -(x_dims.size() + 1),
-            x_dims.size(),
-            dim2));
-
-    int dim1_ = dim1 < 0 ? x_dims.size() + dim1 + 1 : dim1;
-    int dim2_ = dim2 < 0 ? x_dims.size() + dim2 + 1 : dim2;
-    int offset_ = std::abs(offset);
-
-    PADDLE_ENFORCE_NE(dim1_,
-                      dim2_,
-                      platform::errors::InvalidArgument(
-                          "diagonal dimensions should not be identical "
-                          "%ld vs %ld.",
-                          dim1,
-                          dim2));
-
-    int new_dim_len = offset_ + x_dims[x_dims.size() - 1];
-    auto sizes = vectorize(x_dims);
-    sizes.pop_back();
-    sizes.insert(sizes.begin() + std::min(dim1_, dim2_), new_dim_len);
-    sizes.insert(sizes.begin() + std::max(dim1_, dim2_), new_dim_len);
-    ctx->SetOutputDim("Out", phi::make_ddim(sizes));
-  }
 };
 
 class DiagEmbedOpMaker : public framework::OpProtoAndCheckerMaker {
@@ -131,15 +59,14 @@ class DiagEmbedOpMaker : public framework::OpProtoAndCheckerMaker {
 }  // namespace paddle
 
 namespace ops = paddle::operators;
-namespace platform = paddle::platform;
+DECLARE_INFER_SHAPE_FUNCTOR(diag_embed,
+                            DiagEmbedInferShapeFunctor,
+                            PD_INFER_META(phi::DiagEmbedInferMeta));
+
 REGISTER_OPERATOR(
     diag_embed,
     ops::DiagEmbedOp,
     ops::DiagEmbedOpMaker,
     paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
-    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
-REGISTER_OP_CPU_KERNEL(diag_embed,
-                       ops::DiagEmbedKernel<phi::CPUContext, int>,
-                       ops::DiagEmbedKernel<phi::CPUContext, float>,
-                       ops::DiagEmbedKernel<phi::CPUContext, double>,
-                       ops::DiagEmbedKernel<phi::CPUContext, int64_t>);
+    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>,
+    DiagEmbedInferShapeFunctor);

paddle/phi/api/yaml/legacy_api.yaml

@@ -507,6 +507,14 @@
     func : determinant
   backward : det_grad
 
+- api : diag_embed
+  args : (Tensor x, int offset, int dim1, int dim2)
+  output : Tensor
+  infer_meta :
+    func : DiagEmbedInferMeta
+  kernel :
+    func : diag_embed
+
 - api : divide
   args : (Tensor x, Tensor y)
   output : Tensor

paddle/phi/infermeta/unary.cc

@@ -3378,6 +3378,69 @@ void IdentityLossInferMeta(const MetaTensor& x,
   }
 }
 
+void DiagEmbedInferMeta(
+    const MetaTensor& x, int offset, int dim1, int dim2, MetaTensor* out) {
+  auto x_dims = x.dims();
+
+  PADDLE_ENFORCE_GE(
+      dim1,
+      -(x_dims.size() + 1),
+      phi::errors::OutOfRange(
+          "Dim1 is out of range (expected to be in range of [%ld, "
+          "%ld], but got %ld).",
+          -(x_dims.size() + 1),
+          x_dims.size(),
+          dim1));
+  PADDLE_ENFORCE_LE(
+      dim1,
+      x_dims.size(),
+      phi::errors::OutOfRange(
+          "Dim1 is out of range (expected to be in range of [%ld, "
+          "%ld], but got %ld).",
+          -(x_dims.size() + 1),
+          x_dims.size(),
+          dim1));
+
+  PADDLE_ENFORCE_GE(
+      dim2,
+      -(x_dims.size() + 1),
+      phi::errors::OutOfRange(
+          "Dim2 is out of range (expected to be in range of [%ld, "
+          "%ld], but got %ld).",
+          -(x_dims.size() + 1),
+          x_dims.size(),
+          dim2));
+  PADDLE_ENFORCE_LE(
+      dim2,
+      x_dims.size(),
+      phi::errors::OutOfRange(
+          "Dim2 is out of range (expected to be in range of [%ld, "
+          "%ld], but got %ld).",
+          -(x_dims.size() + 1),
+          x_dims.size(),
+          dim2));
+
+  int dim1_ = dim1 < 0 ? x_dims.size() + dim1 + 1 : dim1;
+  int dim2_ = dim2 < 0 ? x_dims.size() + dim2 + 1 : dim2;
+  int offset_ = std::abs(offset);
+
+  PADDLE_ENFORCE_NE(dim1_,
+                    dim2_,
+                    phi::errors::InvalidArgument(
+                        "diagonal dimensions should not be identical "
+                        "%ld vs %ld.",
+                        dim1,
+                        dim2));
+
+  int new_dim_len = offset_ + x_dims[x_dims.size() - 1];
+  auto sizes = vectorize(x_dims);
+  sizes.pop_back();
+  sizes.insert(sizes.begin() + std::min(dim1_, dim2_), new_dim_len);
+  sizes.insert(sizes.begin() + std::max(dim1_, dim2_), new_dim_len);
+  out->set_dims(phi::make_ddim(sizes));
+  out->set_dtype(x.dtype());
+}
+
 }  // namespace phi
 
 PD_REGISTER_INFER_META_FN(flatten, phi::FlattenInferMeta);

paddle/phi/infermeta/unary.h

@@ -486,4 +486,7 @@ void ChannelShuffleInferMeta(const MetaTensor& x,
 
 void IdentityLossInferMeta(const MetaTensor& x, int reduction, MetaTensor* out);
 
+void DiagEmbedInferMeta(
+    const MetaTensor& x, int offset, int dim1, int dim2, MetaTensor* out);
+
 }  // namespace phi

paddle/phi/kernels/cpu/diag_embed_kernel.cc

@@ -0,0 +1,28 @@
+// 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/kernels/diag_embed_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/diag_embed_impl.h"
+
+PD_REGISTER_KERNEL(diag_embed,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::DiagEmbedKernel,
+                   int,
+                   int64_t,
+                   float,
+                   double) {}