migrate overlap_add and overlap_add_grad op

paddle/fluid/operators/overlap_add_op.cc

@@ -12,7 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/operators/overlap_add_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/backward.h"
+#include "paddle/phi/infermeta/unary.h"
 
 namespace paddle {
 namespace operators {
@@ -21,93 +25,6 @@ class OverlapAddOp : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
 
-  void InferShape(framework::InferShapeContext* ctx) const override {
-    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "overlap_add");
-    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "overlap_add");
-
-    const int hop_length = ctx->Attrs().Get<int>("hop_length");
-    const int axis = ctx->Attrs().Get<int>("axis");
-
-    const auto x_dims = ctx->GetInputDim("X");
-    const int x_rank = x_dims.size();
-
-    PADDLE_ENFORCE_GE(
-        x_rank,
-        2,
-        platform::errors::InvalidArgument(
-            "Input(X) of OverlapAddOp should be a tensor which contains "
-            "at least 2 dimensions, but got rank %s.",
-            x_rank));
-
-    PADDLE_ENFORCE_GT(
-        hop_length,
-        0,
-        platform::errors::InvalidArgument(
-            "Attribute(hop_length) of OverlapAddOp should be greater "
-            "than 0, but got %s.",
-            hop_length));
-
-    PADDLE_ENFORCE_EQ(
-        (axis == 0 || axis == -1),
-        true,
-        platform::errors::InvalidArgument(
-            "Attribute(axis) of OverlapAddOp should 0 or -1, but got %s.",
-            axis));
-
-    std::vector<int64_t> output_shape;
-    int n_frames;
-    int frame_length;
-    int seq_length;
-
-    int start_axis;
-    int end_axis;
-    if (axis == 0) {
-      n_frames = x_dims[0];
-      frame_length = x_dims[1];
-      start_axis = 2;
-      end_axis = x_rank - 1;
-    } else {
-      n_frames = x_dims[x_rank - 1];
-      frame_length = x_dims[x_rank - 2];
-      start_axis = 0;
-      end_axis = x_rank - 3;
-    }
-
-    bool contain_unknown_dim = phi::contain_unknown_dim(x_dims);
-    bool check = ctx->IsRuntime() || !contain_unknown_dim;
-    if (check) {
-      PADDLE_ENFORCE_LE(
-          hop_length,
-          frame_length,
-          platform::errors::InvalidArgument(
-              "Attribute(hop_length) of OverlapAddOp should be less or equal "
-              "than frame_length, but got hop_length(%s) > frame_length(%s).",
-              hop_length,
-              frame_length));
-    }
-
-    if (n_frames == -1) {
-      seq_length = -1;
-    } else {
-      seq_length = (n_frames - 1) * hop_length + frame_length;
-    }
-
-    // It won't go into for loop when x_rank == 2U.
-    for (int i = start_axis; i <= end_axis; i++) {
-      output_shape.push_back(x_dims[i]);
-    }
-
-    if (axis == 0) {
-      // (seq_length, ...)
-      output_shape.insert(output_shape.begin(), seq_length);
-    } else {
-      // (..., seq_length)
-      output_shape.push_back(seq_length);
-    }
-
-    ctx->SetOutputDim("Out", phi::make_ddim(output_shape));
-  }
-
  protected:
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext& ctx) const override {
@@ -137,17 +54,6 @@ class OverlapAddOpMaker : public framework::OpProtoAndCheckerMaker {
 class OverlapAddOpGrad : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
-  void InferShape(framework::InferShapeContext* ctx) const override {
-    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "overlap_add_grad");
-    OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")),
-                   "Input",
-                   "Out@GRAD",
-                   "overlap_add_grad");
-    const auto x_dims = ctx->GetInputDim("X");
-    if (ctx->HasOutput(framework::GradVarName("X"))) {
-      ctx->SetOutputDim(framework::GradVarName("X"), x_dims);
-    }
-  }
 
  protected:
   framework::OpKernelType GetExpectedKernelType(
@@ -176,30 +82,21 @@ class OverlapAddOpGradMaker : public framework::SingleGradOpMaker<T> {
 
 namespace ops = paddle::operators;
 
+DECLARE_INFER_SHAPE_FUNCTOR(overlap_add,
+                            OverlapAddInferShapeFunctor,
+                            PD_INFER_META(phi::OverlapAddInferMeta));
+
+DECLARE_INFER_SHAPE_FUNCTOR(overlap_add_grad,
+                            OverlapAddGradInferShapeFunctor,
+                            PD_INFER_META(phi::OverlapAddGradInferMeta));
+
 REGISTER_OPERATOR(overlap_add,
                   ops::OverlapAddOp,
                   ops::OverlapAddOpMaker,
                   ops::OverlapAddOpGradMaker<paddle::framework::OpDesc>,
-                  ops::OverlapAddOpGradMaker<paddle::imperative::OpBase>);
-
-REGISTER_OPERATOR(overlap_add_grad, ops::OverlapAddOpGrad);
-
-REGISTER_OP_CPU_KERNEL(
-    overlap_add,
-    ops::OverlapAddKernel<phi::CPUContext, int>,
-    ops::OverlapAddKernel<phi::CPUContext, int64_t>,
-    ops::OverlapAddKernel<phi::CPUContext, float>,
-    ops::OverlapAddKernel<phi::CPUContext, double>,
-    ops::OverlapAddKernel<phi::CPUContext, paddle::platform::complex<float>>,
-    ops::OverlapAddKernel<phi::CPUContext, paddle::platform::complex<double>>);
+                  ops::OverlapAddOpGradMaker<paddle::imperative::OpBase>,
+                  OverlapAddInferShapeFunctor);
 
-REGISTER_OP_CPU_KERNEL(
-    overlap_add_grad,
-    ops::OverlapAddGradKernel<phi::CPUContext, int>,
-    ops::OverlapAddGradKernel<phi::CPUContext, int64_t>,
-    ops::OverlapAddGradKernel<phi::CPUContext, float>,
-    ops::OverlapAddGradKernel<phi::CPUContext, double>,
-    ops::OverlapAddGradKernel<phi::CPUContext,
-                              paddle::platform::complex<float>>,
-    ops::OverlapAddGradKernel<phi::CPUContext,
-                              paddle::platform::complex<double>>);
+REGISTER_OPERATOR(overlap_add_grad,
+                  ops::OverlapAddOpGrad,
+                  OverlapAddGradInferShapeFunctor);