support auto-gen cumsum (#54948)

paddle/fluid/operators/cum_op.cc

-/* Copyright (c) 2018 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/fluid/framework/infershape_utils.h"
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/fluid/framework/op_version_registry.h"
-#include "paddle/fluid/prim/api/composite_backward/composite_backward_api.h"
-#include "paddle/fluid/prim/utils/static/composite_grad_desc_maker.h"
-#include "paddle/fluid/prim/utils/static/desc_tensor.h"
-#include "paddle/phi/core/infermeta_utils.h"
-#include "paddle/phi/infermeta/unary.h"
-
-namespace paddle {
-namespace operators {
-
-class CumOp : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const override {
-    auto input_data_type =
-        framework::OperatorWithKernel::IndicateVarDataType(ctx, "X");
-    return phi::KernelKey(input_data_type, ctx.GetPlace());
-  }
-};
-
-class CumGradOp : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
-  void InferShape(framework::InferShapeContext* ctx) const override {
-    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "cumsum");
-    OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")),
-                   "Input",
-                   "Out@GRAD",
-                   "cumsum");
-    ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
-  }
-
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const override {
-    auto input_data_type =
-        framework::OperatorWithKernel::IndicateVarDataType(ctx, "X");
-    return phi::KernelKey(input_data_type, ctx.GetPlace());
-  }
-};
-
-class CumsumOpMaker : public framework::OpProtoAndCheckerMaker {
- public:
-  void Make() override {
-    AddInput("X", "Input of cumsum operator");
-    AddOutput("Out", "Output of cumsum operator");
-    AddAttr<int>("axis",
-                 "The dimension to accumulate along. -1 means the last "
-                 "dimension [default -1].")
-        .SetDefault(-1)
-        .SupportTensor();
-    AddAttr<bool>("flatten",
-                  "Whether to compute the cumsum over the flattened array. "
-                  "[default false].")
-        .SetDefault(false);
-    AddAttr<bool>("exclusive",
-                  "Whether to perform exclusive cumsum. [default false].")
-        .SetDefault(false);
-    AddAttr<bool>("reverse",
-                  "If true, the cumsum is performed in the reversed direction. "
-                  "[default false].")
-        .SetDefault(false);
-    AddComment(R"DOC(
-The cumulative sum of the elements along a given axis.
-By default, the first element of the result is the same of the first element of
-the input. If exclusive is true, the first element of the result is 0.
-)DOC");
-  }
-};
-
-template <typename T>
-class CumsumGradMaker : public framework::SingleGradOpMaker<T> {
- public:
-  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
-
- protected:
-  void Apply(GradOpPtr<T> grad_op) const override {
-    grad_op->SetType("cumsum_grad");
-    grad_op->SetInput("X", this->Input("X"));
-    grad_op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
-    grad_op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
-    grad_op->SetAttrMap(this->Attrs());
-    grad_op->SetAttr("reverse",
-                     PADDLE_GET_CONST(bool, this->GetAttr("reverse")));
-  }
-};
-
-class CumsumCompositeGradOpMaker : public prim::CompositeGradOpMakerBase {
-  using prim::CompositeGradOpMakerBase::CompositeGradOpMakerBase;
-
- public:
-  void Apply() override {
-    paddle::Tensor x = this->GetSingleForwardInput("X");
-    paddle::Tensor out_grad = this->GetSingleOutputGrad("Out");
-    paddle::Tensor dx = this->GetSingleInputGrad("X");
-    auto* dx_ptr = this->GetOutputPtr(&dx);
-    std::string dx_name = this->GetOutputName(dx);
-    int axis = static_cast<int>(this->Attr<int>("axis"));
-    bool flatten = static_cast<bool>(this->Attr<bool>("flatten"));
-    bool exclusive = static_cast<bool>(this->Attr<bool>("exclusive"));
-    bool reverse = static_cast<bool>(this->Attr<bool>("reverse"));
-    VLOG(6) << "Runing cumsum composite func";
-    prim::cumsum_grad<prim::DescTensor>(
-        x, out_grad, axis, flatten, exclusive, reverse, dx_ptr);
-    this->RecoverOutputName(dx, dx_name);
-  }
-};
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-using CPU = phi::CPUContext;
-DECLARE_INFER_SHAPE_FUNCTOR(cumsum,
-                            CumsumInferShapeFunctor,
-                            PD_INFER_META(phi::CumScalarAxisInferMeta));
-
-REGISTER_OPERATOR(cumsum,
-                  ops::CumOp,
-                  ops::CumsumOpMaker,
-                  ops::CumsumCompositeGradOpMaker,
-                  ops::CumsumGradMaker<paddle::framework::OpDesc>,
-                  ops::CumsumGradMaker<paddle::imperative::OpBase>,
-                  CumsumInferShapeFunctor);
-REGISTER_OPERATOR(cumsum_grad, ops::CumGradOp);
-
-REGISTER_OP_VERSION(cumsum).AddCheckpoint(
-    R"ROC(
-      Upgrade cumsum add a new attribute [flatten].
-    )ROC",
-    paddle::framework::compatible::OpVersionDesc().NewAttr(
-        "flatten",
-        "In order to compute the cumsum over the flattened array when the "
-        "argument `axis` in python API is None.",
-        false));

paddle/phi/api/yaml/backward.yaml

@@ -493,6 +493,18 @@
   kernel :
     func : cumprod_grad
 
+- backward_op : cumsum_grad
+  forward : cumsum(Tensor x, Scalar axis=-1, bool flatten=false, bool exclusive=false, bool reverse=false) -> Tensor(out)
+  args : (Tensor x, Tensor out_grad, Scalar axis, bool flatten, bool exclusive, bool reverse)
+  output : Tensor(x_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param: [x]
+  kernel :
+    func : cumsum_grad
+    data_type: x
+  composite: cumsum_grad(x, out_grad, axis, flatten, exclusive, reverse, x_grad)
+
 - backward_op : depthwise_conv2d_double_grad
   forward : depthwise_conv2d_grad (Tensor input, Tensor filter, Tensor grad_out, int[] strides, int[] paddings, str padding_algorithm, int groups, int[] dilations, str data_format) -> Tensor(grad_input), Tensor(grad_filter)
   args : (Tensor input, Tensor filter, Tensor grad_out, Tensor grad_input_grad, Tensor grad_filter_grad, int[] strides, int[] paddings, str padding_algorithm, int groups, int[] dilations, str data_format)

paddle/phi/api/yaml/legacy_backward.yaml

@@ -159,18 +159,6 @@
     data_type : x
   backward : conv2d_transpose_double_grad
 
-- backward_op : cumsum_grad
-  forward : cumsum(Tensor x, Scalar axis, bool flatten, bool exclusive, bool reverse) -> Tensor(out)
-  args : (Tensor x, Tensor out_grad, Scalar axis, bool flatten, bool exclusive, bool reverse)
-  output : Tensor(x_grad)
-  infer_meta :
-    func : UnchangedInferMeta
-    param: [x]
-  kernel :
-    func : cumsum_grad
-    data_type: x
-  composite: cumsum_grad(x, out_grad, axis, flatten, exclusive, reverse, x_grad)
-
 - backward_op : deformable_conv_grad
   forward : deformable_conv(Tensor x, Tensor offset, Tensor filter, Tensor mask, int[] strides, int[] paddings, int[] dilations, int deformable_groups, int groups, int im2col_step) -> Tensor(out)
   args : (Tensor x, Tensor offset, Tensor filter, Tensor mask, Tensor out_grad, int[] strides, int[] paddings, int[] dilations, int deformable_groups, int groups, int im2col_step)

paddle/phi/api/yaml/legacy_ops.yaml

@@ -166,15 +166,6 @@
   output : Tensor(out)
   invoke : copy_to_impl(x, place, blocking)
 
-- op : cumsum
-  args : (Tensor x, Scalar axis, bool flatten, bool exclusive, bool reverse)
-  output : Tensor(out)
-  infer_meta :
-    func : CumScalarAxisInferMeta
-  kernel :
-    func : cumsum
-  backward : cumsum_grad
-
 - op : decode_jpeg
   args : (Tensor x, str mode, Place place)
   output : Tensor(out)

paddle/phi/api/yaml/op_compat.yaml

@@ -663,7 +663,7 @@
   scalar:
     axis:
       data_type : int
-      tensor_name: AxisTensor
+      support_tensor : true
 
 - op : data_norm
   backward : data_norm_grad

paddle/phi/api/yaml/op_version.yaml

@@ -140,6 +140,14 @@
           comment : In order to add additional size to one side of each dimension in the output.
           default : "std::vector<int>{}"
 
+- op : cumsum
+  version :
+    - checkpoint : Upgrade cumsum add a new attribute [flatten].
+      action :
+        - add_attr : flatten
+          comment : In order to compute the cumsum over the flattened array when the argument `axis` in python API is None.
+          default : "false"
+
 - op : depthwise_conv2d
   version :
     - checkpoint : Upgrade depthwise_conv2d, add a new attribute [use_addto].

paddle/phi/api/yaml/ops.yaml

@@ -588,6 +588,16 @@
     func : cumprod
   backward : cumprod_grad
 
+- op : cumsum
+  args : (Tensor x, Scalar axis=-1, bool flatten=false, bool exclusive=false, bool reverse=false)
+  output : Tensor(out)
+  infer_meta :
+    func : CumScalarAxisInferMeta
+  kernel :
+    func : cumsum
+    data_type : x
+  backward : cumsum_grad
+
 - op : depthwise_conv2d
   args : (Tensor input, Tensor filter, int[] strides={1, 1}, int[] paddings={0, 0}, str padding_algorithm="EXPLICIT", int groups=1, int[] dilations={1, 1}, str data_format="NCHW")
   output : Tensor(out)

paddle/phi/ops/compat/cumsum_sig.cc

-// Copyright (c) 2023 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/core/compat/op_utils.h"
-
-namespace phi {
-
-KernelSignature CumsumOpArgumentMapping(
-    const ArgumentMappingContext& ctx UNUSED) {
-  return KernelSignature("cumsum_grad",
-                         {"X", "Out@GRAD"},
-                         {"axis", "flatten", "exclusive", "reverse"},
-                         {"X@GRAD"});
-}
-
-}  // namespace phi
-
-PD_REGISTER_ARG_MAPPING_FN(cumsum_grad, phi::CumsumOpArgumentMapping);