[NewIR]Split python api and vjp

paddle/fluid/ir/dialect/op_generator/op_gen.py

@@ -172,7 +172,7 @@ class {op_name} : public ir::Op<{op_name}{interfaces}{traits}> {{
     'bool': 'ir::BoolAttribute',
 }
 
-_NO_NEED_GEN_OPS = {'add_n'}
+_NO_NEED_GEN_OPS = {'add_n', 'split_grad'}
 
 
 def to_phi_and_fluid_op_name(op_item):

paddle/fluid/ir/dialect/op_generator/vjp_interface_gen_op_list.py

@@ -29,5 +29,6 @@
     "sum",
     "add",
     "concat",
+    "split",
 ]
 vjp_interface_implementation_gen_op_list = ["tanh", "mean", "divide", "add"]

paddle/fluid/ir/dialect/paddle_dialect/ir/pd_dialect.cc

@@ -48,7 +48,7 @@ void PaddleDialect::initialize() {
 #define GET_OP_LIST
 #include "paddle/fluid/ir/dialect/paddle_dialect/ir/pd_op.h"  // NOLINT
       >();
-  RegisterOp<paddle::dialect::AddNOp>();
+  RegisterOps<paddle::dialect::AddNOp, paddle::dialect::SplitGradOp>();
 
   RegisterInterfaces<ParameterConvertInterface>();
 }

paddle/fluid/ir/dialect/paddle_dialect/ir/pd_manual_api.cc

@@ -18,5 +18,16 @@
 #include "paddle/ir/core/builtin_op.h"
 
 namespace paddle {
-namespace dialect {}  // namespace dialect
+namespace dialect {
+ir::OpResult split_grad(std::vector<ir::OpResult> out_grads,
+                        ir::OpResult axis) {
+  auto combine_op =
+      APIBuilder::Instance().GetBuilder()->Build<ir::CombineOp>(out_grads);
+  paddle::dialect::SplitGradOp split_grad_op =
+      APIBuilder::Instance().GetBuilder()->Build<paddle::dialect::SplitGradOp>(
+          combine_op.out(), axis);
+
+  return split_grad_op.x_grad();
+}
+}  // namespace dialect
 }  // namespace paddle

paddle/fluid/ir/dialect/paddle_dialect/ir/pd_manual_api.h

@@ -21,5 +21,9 @@
 #include "paddle/phi/common/place.h"
 
 namespace paddle {
-namespace dialect {}  // namespace dialect
+namespace dialect {
+
+ir::OpResult split_grad(std::vector<ir::OpResult> out_grads, ir::OpResult axis);
+
+}  // namespace dialect
 }  // namespace paddle

paddle/fluid/ir/dialect/paddle_dialect/ir/pd_manual_op.cc

@@ -14,6 +14,7 @@
 
 #include "paddle/fluid/ir/dialect/paddle_dialect/ir/pd_manual_op.h"
 #include "paddle/fluid/ir/dialect/paddle_dialect/ir/pd_attribute.h"
+#include "paddle/fluid/ir/dialect/paddle_dialect/ir/pd_op.h"
 #include "paddle/fluid/ir/dialect/paddle_dialect/ir/pd_type.h"
 #include "paddle/ir/core/builtin_attribute.h"
 #include "paddle/ir/core/builtin_op.h"
@@ -145,7 +146,221 @@ void AddNOp::InferMeta(phi::InferMetaContext *infer_meta) {
   fn(infer_meta);
 }
 
+const char *SplitGradOp::attributes_name[1] = {"axis"};
+
+OpInfoTuple SplitGradOp::GetOpInfo() {
+  std::vector<paddle::dialect::OpInputInfo> inputs = {
+      OpInputInfo("out_grad",
+                  "ir::VectorType<paddle::dialect::DenseTensorType>",
+                  false,
+                  false,
+                  false),
+      OpInputInfo(
+          "axis", "paddle::dialect::ScalarAttribute", false, false, true)};
+  std::vector<paddle::dialect::OpAttributeInfo> attributes = {};
+  std::vector<paddle::dialect::OpOutputInfo> outputs = {
+      OpOutputInfo("x_grad", "paddle::dialect::DenseTensorType", false, false)};
+  paddle::dialect::OpRunTimeInfo run_time_info =
+      OpRunTimeInfo("ConcatInferMeta",
+                    {"out_grad", "axis"},
+                    {"concat"},
+                    {"out_grad", "axis"},
+                    {"out_grad"},
+                    {},
+                    {},
+                    {});
+
+  return std::make_tuple(
+      inputs, attributes, outputs, run_time_info, "split_grad");
+}
+
+void SplitGradOp::Build(ir::Builder &builder,
+                        ir::OperationArgument &argument,
+                        ir::OpResult out_grad_,
+                        float axis) {
+  // Generate scalar mutable attribute: axis
+  paddle::dialect::FullOp full_axis_op = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{1}, axis, phi::DataType::FLOAT32, phi::CPUPlace());
+  ir::OpResult axis_ = full_axis_op->result(0);
+
+  VLOG(4) << "Builder construction inputs";
+  std::vector<ir::OpResult> argument_inputs = {out_grad_, axis_};
+  argument.AddOperands(argument_inputs.begin(), argument_inputs.end());
+
+  VLOG(4) << "Builder construction attributes";
+
+  VLOG(4) << "Builder construction outputs";
+  ir::VectorType out_grad = out_grad_.type().dyn_cast<ir::VectorType>();
+  std::vector<phi::DenseTensor> vec_dense_out_grad;
+  for (size_t i = 0; i < static_cast<size_t>(out_grad.size()); i++) {
+    vec_dense_out_grad.push_back(phi::DenseTensor(
+        std::make_unique<paddle::experimental::DefaultAllocator>(
+            paddle::platform::CPUPlace())
+            .get(),
+        phi::DenseTensorMeta(
+            paddle::dialect::TransToPhiDataType(
+                out_grad[i]
+                    .dyn_cast<paddle::dialect::DenseTensorType>()
+                    .dtype()),
+            out_grad[i].dyn_cast<paddle::dialect::DenseTensorType>().dims(),
+            out_grad[i]
+                .dyn_cast<paddle::dialect::DenseTensorType>()
+                .data_layout(),
+            out_grad[i].dyn_cast<paddle::dialect::DenseTensorType>().lod(),
+            out_grad[i]
+                .dyn_cast<paddle::dialect::DenseTensorType>()
+                .offset())));
+  }
+  std::vector<phi::MetaTensor> vec_meta_out_grad;
+  for (size_t i = 0; i < vec_dense_out_grad.size(); i++) {
+    vec_meta_out_grad.push_back(phi::MetaTensor(&vec_dense_out_grad[i]));
+  }
+
+  std::vector<const phi::MetaTensor *> meta_out_grad;
+  for (size_t i = 0; i < static_cast<size_t>(vec_meta_out_grad.size()); i++) {
+    meta_out_grad.push_back(&vec_meta_out_grad[i]);
+  }
+  phi::DenseTensor dense_x_grad;
+  phi::MetaTensor meta_x_grad(&dense_x_grad);
+
+  phi::ConcatInferMeta(meta_out_grad, axis, &meta_x_grad);
+
+  std::vector<ir::Type> argument_outputs;
+  ir::Type x_grad_dense_tensor_type = paddle::dialect::DenseTensorType::get(
+      ir::IrContext::Instance(),
+      paddle::dialect::TransToIrDataType(dense_x_grad.dtype()),
+      dense_x_grad.dims(),
+      dense_x_grad.layout(),
+      dense_x_grad.lod(),
+      dense_x_grad.offset());
+  argument_outputs.push_back(x_grad_dense_tensor_type);
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+}
+
+void SplitGradOp::Build(ir::Builder &builder,
+                        ir::OperationArgument &argument,
+                        ir::OpResult out_grad_,
+                        ir::OpResult axis_) {
+  VLOG(4) << "Builder construction inputs";
+  std::vector<ir::OpResult> argument_inputs = {out_grad_, axis_};
+  argument.AddOperands(argument_inputs.begin(), argument_inputs.end());
+
+  VLOG(4) << "Builder construction attributes";
+
+  VLOG(4) << "Builder construction outputs";
+  ir::VectorType out_grad = out_grad_.type().dyn_cast<ir::VectorType>();
+  int axis = axis_.owner()
+                 ->dyn_cast<paddle::dialect::FullOp>()
+                 .attributes()
+                 .at("value")
+                 .dyn_cast<paddle::dialect::ScalarAttribute>()
+                 .data()
+                 .to<int>();
+
+  std::vector<phi::DenseTensor> vec_dense_out_grad;
+  for (size_t i = 0; i < static_cast<size_t>(out_grad.size()); i++) {
+    vec_dense_out_grad.push_back(phi::DenseTensor(
+        std::make_unique<paddle::experimental::DefaultAllocator>(
+            paddle::platform::CPUPlace())
+            .get(),
+        phi::DenseTensorMeta(
+            TransToPhiDataType(out_grad[i]
+                                   .dyn_cast<paddle::dialect::DenseTensorType>()
+                                   .dtype()),
+            out_grad[i].dyn_cast<paddle::dialect::DenseTensorType>().dims(),
+            out_grad[i]
+                .dyn_cast<paddle::dialect::DenseTensorType>()
+                .data_layout(),
+            out_grad[i].dyn_cast<paddle::dialect::DenseTensorType>().lod(),
+            out_grad[i]
+                .dyn_cast<paddle::dialect::DenseTensorType>()
+                .offset())));
+  }
+  std::vector<phi::MetaTensor> vec_meta_out_grad;
+  for (size_t i = 0; i < vec_dense_out_grad.size(); i++) {
+    vec_meta_out_grad.push_back(phi::MetaTensor(&vec_dense_out_grad[i]));
+  }
+
+  std::vector<const phi::MetaTensor *> meta_out_grad;
+  for (size_t i = 0; i < static_cast<size_t>(vec_meta_out_grad.size()); i++) {
+    meta_out_grad.push_back(&vec_meta_out_grad[i]);
+  }
+  phi::DenseTensor dense_x_grad;
+  phi::MetaTensor meta_x_grad(&dense_x_grad);
+
+  phi::ConcatInferMeta(meta_out_grad, axis, &meta_x_grad);
+
+  std::vector<ir::Type> argument_outputs;
+  ir::Type x_grad_dense_tensor_type = paddle::dialect::DenseTensorType::get(
+      ir::IrContext::Instance(),
+      TransToIrDataType(dense_x_grad.dtype()),
+      dense_x_grad.dims(),
+      dense_x_grad.layout(),
+      dense_x_grad.lod(),
+      dense_x_grad.offset());
+  argument_outputs.push_back(x_grad_dense_tensor_type);
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+}
+
+void SplitGradOp::Verify() {
+  VLOG(4) << "Start Verifying inputs, outputs and attributes for: SplitGradOp.";
+  VLOG(4) << "Verifying inputs:";
+  {
+    auto input_size = num_operands();
+    PADDLE_ENFORCE_EQ(
+        input_size,
+        2u,
+        phi::errors::PreconditionNotMet(
+            "The size %d of inputs must be equal to 2.", input_size));
+    if (auto vec_type =
+            (*this)->operand_source(0).type().dyn_cast<ir::VectorType>()) {
+      for (size_t i = 0; i < vec_type.size(); ++i) {
+        PADDLE_ENFORCE(vec_type[i].isa<paddle::dialect::DenseTensorType>(),
+                       phi::errors::PreconditionNotMet(
+                           "Type validation failed for the 0th input."));
+      }
+    } else {
+      PADDLE_ENFORCE((*this)
+                         ->operand_source(0)
+                         .type()
+                         .isa<paddle::dialect::DenseTensorType>(),
+                     phi::errors::PreconditionNotMet(
+                         "Type validation failed for the 0th input."));
+    }
+    PADDLE_ENFORCE((*this)
+                       ->operand_source(1)
+                       .type()
+                       .isa<paddle::dialect::DenseTensorType>(),
+                   phi::errors::PreconditionNotMet(
+                       "Type validation failed for the 1th input."));
+  }
+  VLOG(4) << "Verifying attributes:";
+  {
+    // Attributes num is 0, not need to check attributes type.
+  }
+  VLOG(4) << "Verifying outputs:";
+  {
+    auto output_size = num_results();
+    PADDLE_ENFORCE_EQ(
+        output_size,
+        1u,
+        phi::errors::PreconditionNotMet(
+            "The size %d of outputs must be equal to 1.", output_size));
+    PADDLE_ENFORCE(
+        (*this)->result(0).type().isa<paddle::dialect::DenseTensorType>(),
+        phi::errors::PreconditionNotMet(
+            "Type validation failed for the 0th output."));
+  }
+  VLOG(4) << "End Verifying for: SplitGradOp.";
+}
+
+void SplitGradOp::InferMeta(phi::InferMetaContext *infer_meta) {
+  auto fn = PD_INFER_META(phi::ConcatInferMeta);
+  fn(infer_meta);
+}
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AddNOp)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SplitGradOp)

paddle/fluid/ir/dialect/paddle_dialect/ir/pd_manual_op.h

@@ -14,7 +14,7 @@
 
 #ifdef GET_MANUAL_OP_LIST
 #undef GET_MANUAL_OP_LIST
-paddle::dialect::AddNOp
+paddle::dialect::AddNOp, paddle::dialect::SplitGradOp
 
 #else
 
@@ -51,9 +51,33 @@ class AddNOp : public ir::Op<AddNOp, OpYamlInfoInterface> {
   static void InferMeta(phi::InferMetaContext *infer_meta);
 };
 
+class SplitGradOp : public ir::Op<SplitGradOp, OpYamlInfoInterface> {
+ public:
+  using Op::Op;
+  static const char *name() { return "pd.split_grad"; }
+  static const char *attributes_name[1];
+  static constexpr uint32_t attributes_num = 1;
+  static OpInfoTuple GetOpInfo();
+  static void Build(ir::Builder &builder,             // NOLINT
+                    ir::OperationArgument &argument,  // NOLINT
+                    ir::OpResult x_,
+                    float axis = 0);
+  static void Build(ir::Builder &builder,             // NOLINT
+                    ir::OperationArgument &argument,  // NOLINT
+                    ir::OpResult out_grad_,
+                    ir::OpResult axis_);
+
+  void Verify();
+  ir::Value out_grad() { return operand_source(0); }
+  ir::Value axis() { return operand_source(1); }
+  ir::OpResult x_grad() { return result(0); }
+  static void InferMeta(phi::InferMetaContext *infer_meta);
+};
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AddNOp)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SplitGradOp)
 
 #endif

paddle/fluid/ir/dialect/paddle_dialect/ir/pd_op_vjp_manual.cc

@@ -87,5 +87,39 @@ std::vector<std::vector<ir::OpResult>> SumOp::Vjp(
   }
   return res;
 }
+
+std::vector<std::vector<ir::OpResult>> SplitOp::Vjp(
+    ir::Operation* op,
+    const std::vector<std::vector<ir::OpResult>>& out_grads,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  SplitOp op_obj = op->dyn_cast<SplitOp>();
+
+  Tensor axis(std::make_shared<primitive::LazyTensor>(op_obj.axis()));
+  std::vector<Tensor> out_grads_;
+  for (size_t idx = 0; idx < out_grads[0].size(); idx++) {
+    out_grads_.emplace_back(
+        std::make_shared<primitive::LazyTensor>(out_grads[0][idx]));
+  }
+
+  std::vector<std::vector<Tensor>> tensor_res =
+      primitive::split_vjp(out_grads_, axis, stop_gradients);
+
+  std::vector<std::vector<ir::OpResult>> res(tensor_res.size(),
+                                             std::vector<ir::OpResult>());
+
+  for (uint64_t i = 0; i < tensor_res.size(); i++) {
+    res[i].resize(tensor_res[i].size());
+    for (uint64_t j = 0; j < tensor_res[i].size(); j++) {
+      if (tensor_res[i][j].defined()) {
+        res[i][j] = std::static_pointer_cast<primitive::LazyTensor>(
+                        tensor_res[i][j].impl())
+                        ->getValue()
+                        .dyn_cast<ir::OpResult>();
+      }
+    }
+  }
+  return res;
+}
+
 }  // namespace dialect
 }  // namespace paddle

paddle/fluid/primitive/backend/manual/manual_backend.h

@@ -33,6 +33,9 @@ std::vector<Tensor> concat_grad(const std::vector<Tensor>& x,
                                 const Tensor& out_grad,
                                 const Tensor& axis);
 
+template <typename T>
+Tensor split_grad(const std::vector<Tensor>& out_grads, const Tensor& axis);
+
 }  // namespace backend
 }  // namespace primitive
 }  // namespace paddle