PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/op_generator/op_gen.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/pir/dialect/op_generator/op_gen.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
Copy file name to clipboard
+7-7Lines changed: 7 additions & 7 deletions b/‎paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
Copy file name to clipboard
+7-7Lines changed: 7 additions & 7 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+163-1Lines changed: 163 additions & 1 deletion b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+163-1Lines changed: 163 additions & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+42Lines changed: 42 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+42Lines changed: 42 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op_vjp.cc
Copy file name to clipboard
+41Lines changed: 41 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op_vjp.cc
Copy file name to clipboard
+41Lines changed: 41 additions & 0 deletions
@@ -286,6 +286,7 @@ class {TEST_API} {op_name} : public pir::Op<{op_name}{interfaces}{traits}> {{
     'expand',
     'increment',
     'increment_',
+    'assign_out_',
 }
 
 attr_types_map = {
 
@@ -282,20 +282,20 @@ std::vector<std::vector<pir::Value>> IfOp::Vjp(
   VLOG(6) << "Prepare outputs for if_grad";
 
   std::vector<pir::Type> output_types;
-  for (size_t i = 0; i < inputs_.size(); ++i) {
-    if (!stop_gradients[i][0]) {
+  for (size_t i = 1; i < inputs_.size(); ++i) {
+    if (!stop_gradients[i - 1][0]) {
       output_types.push_back(inputs_[i][0].type());
     }
   }
 
   auto if_grad = ApiBuilder::Instance().GetBuilder()->Build<IfOp>(
       cond_val, std::move(output_types));
 
-  std::vector<std::vector<pir::Value>> res{inputs_.size()};
-  for (size_t i = 0, j = 0; i < inputs_.size(); ++i) {
-    res[i].resize(1);
-    if (!stop_gradients[i][0]) {
-      res[i][0] = if_grad->result(j++);
+  std::vector<std::vector<pir::Value>> res{inputs_.size() - 1};
+  for (size_t i = 1, j = 0; i < inputs_.size(); ++i) {
+    res[i - 1].resize(1);
+    if (!stop_gradients[i - 1][0]) {
+      res[i - 1][0] = if_grad->result(j++);
     }
   }
   return res;
 
@@ -15,7 +15,7 @@
 #undef GET_OP_LIST
 paddle::dialect::AddNOp, paddle::dialect::AddN_Op,
     paddle::dialect::AddNWithKernelOp, paddle::dialect::AddNArrayOp,
-    paddle::dialect::FusedGemmEpilogueOp,
+    paddle::dialect::FusedGemmEpilogueOp, paddle::dialect::AssignOut_Op,
     paddle::dialect::FusedGemmEpilogueGradOp, paddle::dialect::SplitGradOp,
     paddle::dialect::ExpandOp, paddle::dialect::CreateArrayOp,
     paddle::dialect::CreateArrayLikeOp, paddle::dialect::ArrayLengthOp,
@@ -3976,6 +3976,167 @@ phi::DataType Increment_Op::GetKernelTypeForVar(
   return expected_kernel_dtype;
 }
 
+OpInfoTuple AssignOut_Op::GetOpInfo() {
+  std::vector<paddle::dialect::OpInputInfo> inputs = {
+      paddle::dialect::OpInputInfo(
+          "x", "paddle::dialect::DenseTensorType", false, false, false, true),
+      paddle::dialect::OpInputInfo("output",
+                                   "paddle::dialect::DenseTensorType",
+                                   false,
+                                   false,
+                                   false,
+                                   true)};
+  std::vector<paddle::dialect::OpAttributeInfo> attributes = {};
+  std::vector<paddle::dialect::OpOutputInfo> outputs = {
+      paddle::dialect::OpOutputInfo(
+          "out", "paddle::dialect::DenseTensorType", false, false)};
+  paddle::dialect::OpRunTimeInfo run_time_info =
+      paddle::dialect::OpRunTimeInfo("UnchangedInferMeta",
+                                     {"x"},
+                                     "assign",
+                                     {"x"},
+                                     {},
+                                     {},
+                                     {{"out", "output"}},
+                                     {});
+  return std::make_tuple(
+      inputs, attributes, outputs, run_time_info, "assign_out_");
+}
+
+void AssignOut_Op::Build(pir::Builder &builder,
+                         pir::OperationArgument &argument,
+                         pir::Value x_,
+                         pir::Value output_) {
+  VLOG(4) << "Start build AssignOut_Op";
+
+  VLOG(4) << "Builder construction inputs";
+  std::vector<pir::Value> argument_inputs = {x_, output_};
+  argument.AddInputs(argument_inputs);
+
+  VLOG(4) << "Builder construction attributes";
+  pir::AttributeMap argument_attributes = {};
+
+  std::vector<pir::Type> argument_outputs =
+      AssignOut_Op::InferMeta(argument_inputs, argument_attributes);
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  constexpr char kStopGradientAttrName[] = "stop_gradient";
+  auto stop_gradient0 =
+      argument.inputs[0].attribute<pir::BoolAttribute>(kStopGradientAttrName);
+  argument.inputs[1].set_attribute(kStopGradientAttrName,
+                                   builder.bool_attr(stop_gradient0.data()));
+}
+
+void AssignOut_Op::VerifySig() {
+  VLOG(4)
+      << "Start Verifying inputs, outputs and attributes for: AssignOut_Op.";
+  VLOG(4) << "Verifying inputs:";
+  {
+    auto input_size = num_operands();
+    IR_ENFORCE(input_size == 2u,
+               "The size %d of inputs must be equal to 2.",
+               input_size);
+    IR_ENFORCE((*this)
+                   ->operand_source(0)
+                   .type()
+                   .isa<paddle::dialect::DenseTensorType>(),
+               "Type validation failed for the 0th input, got %s.",
+               (*this)->operand_source(0).type());
+    IR_ENFORCE((*this)
+                   ->operand_source(1)
+                   .type()
+                   .isa<paddle::dialect::DenseTensorType>(),
+               "Type validation failed for the 1th input, got %s.",
+               (*this)->operand_source(1).type());
+  }
+  VLOG(4) << "Verifying attributes:";
+  {
+    // Attributes num is 0, not need to check attributes type.
+  }
+  VLOG(4) << "Verifying outputs:";
+  {
+    auto output_size = num_results();
+    IR_ENFORCE(output_size == 1u,
+               "The size %d of outputs must be equal to 1.",
+               output_size);
+    IR_ENFORCE(
+        (*this)->result(0).type().isa<paddle::dialect::DenseTensorType>(),
+        "Type validation failed for the 0th output.");
+  }
+  VLOG(4) << "End Verifying for: AssignOut_Op.";
+}
+
+void AssignOut_Op::InferMeta(phi::InferMetaContext *infer_meta) {
+  auto fn = PD_INFER_META(phi::UnchangedInferMeta);
+  fn(infer_meta);
+}
+
+std::vector<pir::Type> AssignOut_Op::InferMeta(
+    const std::vector<pir::Value> &input_values,
+    const pir::AttributeMap &attributes) {
+  IR_ENFORCE(input_values.size() == 2,
+             "Num of inputs is expected to be 2 but got %d.",
+             input_values.size());
+
+  pir::Value x_ = input_values[0];
+  (void)x_;
+  VLOG(4) << "Builder construction outputs";
+
+  paddle::dialect::DenseTensorType x;
+  if (x_.type().isa<paddle::dialect::DenseTensorType>()) {
+    x = x_.type().dyn_cast<paddle::dialect::DenseTensorType>();
+    (void)x;
+  } else if (x_.type().isa<paddle::dialect::AllocatedDenseTensorType>()) {
+    paddle::dialect::AllocatedDenseTensorType allocated_x =
+        x_.type().dyn_cast<paddle::dialect::AllocatedDenseTensorType>();
+    x = paddle::dialect::DenseTensorType::get(pir::IrContext::Instance(),
+                                              allocated_x.dtype(),
+                                              allocated_x.dims(),
+                                              allocated_x.data_layout(),
+                                              allocated_x.lod(),
+                                              allocated_x.offset());
+    (void)x;
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "Only support paddle::dialect::DenseTensorType or "
+        "paddle::dialect::AllocatedDenseTensorType"));
+  }
+
+  VLOG(4) << "Builder construction  dense_x";
+  paddle::dialect::IrTensor ir_tensor_x(
+      paddle::dialect::TransToPhiDataType(x.dtype()),
+      x.dims(),
+      x.data_layout(),
+      x.lod(),
+      x.offset());
+  VLOG(4) << "Builder construction  meta_x";
+  paddle::dialect::IrMetaTensor meta_x(&ir_tensor_x);
+  paddle::dialect::IrTensor dense_out;
+  paddle::dialect::IrMetaTensor meta_out(&dense_out);
+
+  phi::UnchangedInferMeta(meta_x, &meta_out);
+
+  std::vector<pir::Type> argument_outputs;
+  pir::Type out_dense_tensor_type = paddle::dialect::DenseTensorType::get(
+      pir::IrContext::Instance(),
+      paddle::dialect::TransToIrDataType(dense_out.dtype()),
+      dense_out.dims(),
+      dense_out.layout(),
+      dense_out.lod(),
+      dense_out.offset());
+  argument_outputs.push_back(out_dense_tensor_type);
+
+  return argument_outputs;
+}
+
+phi::DataType AssignOut_Op::GetKernelTypeForVar(
+    const std::string &var_name,
+    const phi::DataType &tensor_dtype,
+    const phi::DataType &expected_kernel_dtype) {
+  VLOG(4) << "Get KernelType for Var of op: AssignOut_Op";
+
+  return expected_kernel_dtype;
+}
+
 void ShapeBroadcastOp::Build(pir::Builder &builder,
                              pir::OperationArgument &argument,
                              pir::Value x_,
@@ -4282,6 +4443,7 @@ IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SplitGradOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AddN_Op)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AddNWithKernelOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AddNArrayOp)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AssignOut_Op)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::FusedGemmEpilogueOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::FusedGemmEpilogueGradOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::CreateArrayOp)
 
@@ -673,6 +673,47 @@ class Increment_Op
       const std::vector<std::vector<bool>> &stop_gradients);
 };
 
+class AssignOut_Op
+    : public pir::Op<AssignOut_Op,
+                     paddle::dialect::OpYamlInfoInterface,
+                     paddle::dialect::InferMetaInterface,
+                     paddle::dialect::VjpInterface,
+                     paddle::dialect::GetKernelTypeForVarInterface,
+                     paddle::dialect::InplaceTrait> {
+ public:
+  using Op::Op;
+  static const char *name() { return "pd_op.assign_out_"; }
+  static constexpr const char **attributes_name = nullptr;
+  static constexpr uint32_t attributes_num = 0;
+  static OpInfoTuple GetOpInfo();
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value x_,
+                    pir::Value output_);
+
+  void VerifySig();
+
+  static phi::DataType GetKernelTypeForVar(
+      const std::string &var_name,
+      const phi::DataType &tensor_dtype,
+      const phi::DataType &expected_kernel_dtype);
+
+  pir::Value x() { return operand_source(0); }
+  pir::Value output() { return operand_source(1); }
+  pir::Value out() { return result(0); }
+
+  static void InferMeta(phi::InferMetaContext *infer_meta);
+  static std::vector<pir::Type> InferMeta(
+      const std::vector<pir::Value> &input_values,
+      const pir::AttributeMap &attributes);
+  static std::vector<std::vector<pir::Value>> Vjp(
+      pir::Operation *op,
+      const std::vector<std::vector<pir::Value>> &inputs_,
+      const std::vector<std::vector<pir::Value>> &outputs,
+      const std::vector<std::vector<pir::Value>> &out_grads,
+      const std::vector<std::vector<bool>> &stop_gradients);
+};
+
 class MemcpyD2hMultiIoOp
     : public pir::Op<MemcpyD2hMultiIoOp,
                      paddle::dialect::OpYamlInfoInterface,
@@ -746,6 +787,7 @@ IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SplitGradOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AddN_Op)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AddNWithKernelOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AddNArrayOp)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AssignOut_Op)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::FusedGemmEpilogueOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::FusedGemmEpilogueGradOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::CreateArrayOp)
 
@@ -187,6 +187,47 @@ std::vector<std::vector<pir::Value>> Increment_Op::Vjp(
   return res;
 }
 
+std::vector<std::vector<pir::Value>> AssignOut_Op::Vjp(
+    pir::Operation* op,
+    const std::vector<std::vector<pir::Value>>& inputs_,
+    const std::vector<std::vector<pir::Value>>& outputs,
+    const std::vector<std::vector<pir::Value>>& out_grads,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  PADDLE_ENFORCE_EQ(
+      inputs_.size(),
+      2,
+      platform::errors::InvalidArgument(
+          "assign_out_ op's inputs size should be 2, but now is %d.",
+          inputs_.size()));
+  PADDLE_ENFORCE_EQ(
+      outputs.size(),
+      1,
+      platform::errors::InvalidArgument(
+          "assign_out_ op's outputs size should be 1, but now is %d.",
+          outputs.size()));
+
+  VLOG(6) << "Prepare inputs of assign_out__grad";
+
+  Tensor out_grad(std::make_shared<primitive::LazyTensor>(out_grads[0][0]));
+
+  VLOG(6) << "Vjp prepare Prepare attributes of assign_out__grad";
+
+  VLOG(6) << "Vjp prepare call assign_out_'s vjp interface";
+
+  VLOG(6) << "Vjp prepare stop gradient of assign_out__grad";
+
+  std::vector<std::vector<pir::Value>> res(2);
+  res[0].resize(1);
+  if (!stop_gradients[0][0]) {
+    res[0][0] = out_grads[0][0];
+  }
+  res[1].resize(1);
+  if (!stop_gradients[1][0]) {
+    res[1][0] = paddle::dialect::zeros_like(inputs_[1][0]);
+  }
+  return res;
+}
+
 std::vector<std::vector<pir::Value>> ArrayWrite_Op::Vjp(
     pir::Operation* op,
     const std::vector<std::vector<pir::Value>>& inputs_,
Original file line number	Diff line number	Diff line change
`@@ -286,6 +286,7 @@ class {TEST_API} {op_name} : public pir::Op<{op_name}{interfaces}{traits}> {{`
`286`	`286`	`'expand',`
`287`	`287`	`'increment',`
`288`	`288`	`'increment_',`
	`289`	`+ 'assign_out_',`
`289`	`290`	`}`
`290`	`291`
`291`	`292`	`attr_types_map = {`