[PIR]modify vjp interface gen for control_flow (#58355)

xiaoguoguo626807 · web-flow · commit a4a59c48342b · 2023-10-26T09:15:41.000+08:00
* modify vjp interface

* modify vjp interface

* modify test

* modify test
diff --git a/paddle/fluid/pir/dialect/op_generator/op_interface_gen.py b/paddle/fluid/pir/dialect/op_generator/op_interface_gen.py
@@ -21,65 +21,44 @@
   fn(infer_meta);
 }}
 """
+CHECK_INPUT_TEMPLATE = """
+    PADDLE_ENFORCE_EQ(
+      inputs_.size(),
+      {inputs_size},
+      platform::errors::InvalidArgument("{op_name} op's inputs size should be {inputs_size}, but now is %d.", inputs_.size()));
+    PADDLE_ENFORCE_EQ(
+      outputs.size(),
+      {outputs_size},
+      platform::errors::InvalidArgument("{op_name} op's outputs size should be {outputs_size}, but now is %d.", outputs.size()));
+"""
 
 OP_VJP_FORWARD_INPUT_OR_OUTPUT_TEMPLATE = """
-    {input_type} {input_name}(std::make_shared<primitive::LazyTensor>(op_obj.{input_name}()));"""
+    {input_type} {input_name}(std::make_shared<primitive::LazyTensor>({vjp_param_name}[{input_idx}][0]));"""
 
 OP_VJP_FORWARD_MULTI_INPUT_TEMPLATE = """
-    pir::CombineOp combine_op_obj_{input_name} =
-      op_obj.{input_name}().dyn_cast<pir::OpResult>().owner()->dyn_cast<pir::CombineOp>();
     std::vector<Tensor> {input_name};
-    for (size_t idx = 0; idx < combine_op_obj_{input_name}.inputs().size(); idx++) {{
+    for (size_t idx = 0; idx < {vjp_param_name}[{input_idx}].size(); idx++) {{
         {input_name}.emplace_back(
-            std::make_shared<primitive::LazyTensor>(combine_op_obj_{input_name}.inputs()[idx]));
+            std::make_shared<primitive::LazyTensor>({vjp_param_name}[{input_idx}][idx]));
     }}"""
 
 OP_VJP_FORWARD_OPTIONAL_INPUT_TEMPLATE = """
     paddle::optional<Tensor> {input_name};
-    if (!IsEmptyValue(op_obj.{input_name}())){{
-        {input_name} = paddle::make_optional<Tensor>(Tensor(std::make_shared<primitive::LazyTensor>(op_obj.{input_name}())));
+    if (!IsEmptyValue({vjp_param_name}[{input_idx}][0])){{
+        {input_name} = paddle::make_optional<Tensor>(Tensor(std::make_shared<primitive::LazyTensor>({vjp_param_name}[{input_idx}][0])));
     }}"""
 
 OP_VJP_FORWARD_OPTIONAL_VECTOR_INPUT_TEMPLATE = """
     paddle::optional<std::vector<Tensor>> {input_name};
-    if (!IsEmptyValue(op_obj.{input_name}())){{
-        pir::CombineOp combine_op_obj =
-            op_obj.{input_name}().dyn_cast<pir::OpResult>().owner()->dyn_cast<pir::CombineOp>();
-        std::vector<Tensor> optional_{input_name};
-        for (size_t idx = 0; idx < combine_op_obj.inputs().size(); idx++) {{
+    std::vector<Tensor> optional_{input_name};
+    if (!IsEmptyValue({vjp_param_name}[{input_idx}][0])){{
+        for (size_t idx = 0; idx < {vjp_param_name}[{input_idx}].size(); idx++) {{
             optional_{input_name}.emplace_back(
-                std::make_shared<primitive::LazyTensor>(combine_op_obj.inputs()[idx]));
+                std::make_shared<primitive::LazyTensor>({vjp_param_name}[{input_idx}][idx]));
         }}
         {input_name} = paddle::make_optional<std::vector<Tensor>>(optional_{input_name});
     }}"""
 
-OP_VJP_FORWARD_OUTPUT_GRAD_TEMPLATE = """
-    Tensor {output_grad_name}(std::make_shared<primitive::LazyTensor>(out_grads[{idx1}][{idx2}]));"""
-
-OP_VJP_FORWARD_OUTPUT_GRAD_LIST_TEMPLATE = """
-    std::vector<Tensor> {output_grad_name};
-    for (size_t idx = 0; idx < out_grads[{index}].size(); idx++) {{
-        {output_grad_name}.emplace_back(
-            std::make_shared<primitive::LazyTensor>(out_grads[{index}][idx]));
-    }}"""
-
-OP_VJP_FORWARD_OPTIONAL_OUTPUT_GRAD_TEMPLATE = """
-    paddle::optional<Tensor> {output_grad_name};
-    if (!IsEmptyValue(out_grads[{idx1}][{idx2}])){{
-        {output_grad_name} = paddle::make_optional<Tensor>(Tensor(std::make_shared<primitive::LazyTensor>(out_grads[{idx1}][{idx2}])));
-    }}"""
-
-OP_VJP_FORWARD_OPTIONAL_VECTOR_OUTPUT_GRAD_TEMPLATE = """
-    paddle::optional<std::vector<Tensor>> {output_grad_name};
-    std::vector<Tensor> optional_{output_grad_name};
-    if (!IsEmptyValue(out_grads[{index}])){{
-        for (size_t idx = 0; idx < out_grads[{index}].size(); idx++) {{
-            optional_{output_grad_name}.emplace_back(
-                std::make_shared<primitive::LazyTensor>(out_grads[{index}][idx]));
-        }}
-        {output_grad_name} = paddle::make_optional<std::vector<Tensor>>(optional_{output_grad_name});
-    }}"""
-
 OP_VJP_ATTRIBUTE_TEMPLATE = """
     {attr_type} {attr_name} = op->attribute("{attr_name}").dyn_cast<{attr_parse_type}>().{func}();"""
 
@@ -110,11 +89,9 @@
 
 OP_VJP_DEFINE_TEMPLATE = """
 std::vector<std::vector<pir::OpResult>> {op_class_name}::Vjp(pir::Operation* op, const std::vector<std::vector<pir::Value>>& inputs_, const std::vector<std::vector<pir::OpResult>>& outputs, const std::vector<std::vector<pir::Value>>& out_grads, const std::vector<std::vector<bool>>& stop_gradients){{
-    {op_class_name} op_obj = op->dyn_cast<{op_class_name}>(); (void)op_obj;
-
+{check_param}
     VLOG(6) << "Prepare inputs of {op_grad_name}";
-{forward_input_output_code}
-{forward_output_grad_code}
+{backward_input_code}
 
     VLOG(6) << "Vjp prepare Prepare attributes of {op_grad_name}";
 {attribute_code}
@@ -142,74 +119,66 @@ def gen_op_vjp_str(
     op_grad_info,
 ):
     bw_input_list = op_grad_info.input_name_list
-    forward_input_output_code = ''
-    forward_output_grad_code = ''
+    fwd_input_and_mutable_attr_name_list = (
+        op_info.input_name_list + op_info.mutable_attribute_name_list
+    )
+
+    backward_input_code = ''
     build_args_str = ''
     grad_idx = -1
     for idx in range(len(bw_input_list)):
-        build_args_str += bw_input_list[idx] + ", "
+        bw_input_name = bw_input_list[idx]
+        build_args_str += bw_input_name + ", "
         input_type = input_types_map[op_grad_info.input_type_list[idx]]
-        if (
-            bw_input_list[idx] in op_info.input_name_list
-            or bw_input_list[idx] in op_info.output_name_list
-        ):
-            if op_grad_info.input_optional_list[idx] == 'true':
-                if input_type == 'Tensor':
-                    forward_input_output_code += (
-                        OP_VJP_FORWARD_OPTIONAL_INPUT_TEMPLATE.format(
-                            input_name=bw_input_list[idx],
-                        )
-                    )
-                else:
-                    forward_input_output_code += (
-                        OP_VJP_FORWARD_OPTIONAL_VECTOR_INPUT_TEMPLATE.format(
-                            input_name=bw_input_list[idx],
-                        )
+
+        vjp_param_name = ''
+        index_0 = -1
+        if bw_input_name in fwd_input_and_mutable_attr_name_list:
+            vjp_param_name = 'inputs_'
+            index_0 = fwd_input_and_mutable_attr_name_list.index(bw_input_name)
+        elif bw_input_name in op_info.output_name_list:
+            vjp_param_name = 'outputs'
+            index_0 = op_info.output_name_list.index(bw_input_name)
+        else:
+            vjp_param_name = 'out_grads'
+            grad_idx += 1
+            index_0 = grad_idx
+        if op_grad_info.input_optional_list[idx] == 'true':
+            if input_type == 'Tensor':
+                backward_input_code += (
+                    OP_VJP_FORWARD_OPTIONAL_INPUT_TEMPLATE.format(
+                        vjp_param_name=vjp_param_name,
+                        input_name=bw_input_name,
+                        input_idx=index_0,
                     )
+                )
             else:
-                if input_type == 'Tensor':
-                    forward_input_output_code += (
-                        OP_VJP_FORWARD_INPUT_OR_OUTPUT_TEMPLATE.format(
-                            input_type=input_type,
-                            input_name=bw_input_list[idx],
-                        )
-                    )
-                else:
-                    forward_input_output_code += (
-                        OP_VJP_FORWARD_MULTI_INPUT_TEMPLATE.format(
-                            input_name=bw_input_list[idx],
-                        )
+                backward_input_code += (
+                    OP_VJP_FORWARD_OPTIONAL_VECTOR_INPUT_TEMPLATE.format(
+                        vjp_param_name=vjp_param_name,
+                        input_name=bw_input_name,
+                        input_idx=index_0,
                     )
+                )
         else:
-            grad_idx += 1
-            if op_grad_info.input_optional_list[idx] == 'true':
-                if input_type == 'Tensor':
-                    forward_input_output_code += (
-                        OP_VJP_FORWARD_OPTIONAL_OUTPUT_GRAD_TEMPLATE.format(
-                            output_grad_name=bw_input_list[idx],
-                            idx1=grad_idx,
-                            idx2=0,
-                        )
-                    )
-                else:
-                    forward_input_output_code += OP_VJP_FORWARD_OPTIONAL_VECTOR_OUTPUT_GRAD_TEMPLATE.format(
-                        output_grad_name=bw_input_list[idx], index=grad_idx
+            if input_type == 'Tensor':
+                backward_input_code += (
+                    OP_VJP_FORWARD_INPUT_OR_OUTPUT_TEMPLATE.format(
+                        vjp_param_name=vjp_param_name,
+                        input_type=input_type,
+                        input_name=bw_input_name,
+                        input_idx=index_0,
                     )
+                )
             else:
-                if input_type == 'Tensor':
-                    forward_output_grad_code += (
-                        OP_VJP_FORWARD_OUTPUT_GRAD_TEMPLATE.format(
-                            output_grad_name=bw_input_list[idx],
-                            idx1=grad_idx,
-                            idx2=0,
-                        )
-                    )
-                else:
-                    forward_input_output_code += (
-                        OP_VJP_FORWARD_OUTPUT_GRAD_LIST_TEMPLATE.format(
-                            output_grad_name=bw_input_list[idx], index=grad_idx
-                        )
+                backward_input_code += (
+                    OP_VJP_FORWARD_MULTI_INPUT_TEMPLATE.format(
+                        vjp_param_name=vjp_param_name,
+                        input_name=bw_input_name,
+                        input_idx=index_0,
                     )
+                )
+
     op_attribute_list = op_grad_info.attribute_name_list
     attribute_code = ''
     build_attr_str = ''
@@ -219,8 +188,12 @@ def gen_op_vjp_str(
             if op_attribute_list[idx] in op_info.mutable_attribute_name_list:
                 attribute_code += (
                     OP_VJP_FORWARD_INPUT_OR_OUTPUT_TEMPLATE.format(
+                        vjp_param_name='inputs_',
                         input_type="Tensor",
                         input_name=op_attribute_list[idx],
+                        input_idx=fwd_input_and_mutable_attr_name_list.index(
+                            op_attribute_list[idx]
+                        ),
                     )
                 )
                 build_args_str += op_attribute_list[idx] + ", "
@@ -270,14 +243,19 @@ def gen_op_vjp_str(
         inputs_list=build_args_str,
     )
     stop_gradient_input_grad_code = OP_VJP_STOPGRADIENT_TEMPLATE
+    check_param = CHECK_INPUT_TEMPLATE.format(
+        op_name=op_phi_name_format,
+        inputs_size=len(fwd_input_and_mutable_attr_name_list),
+        outputs_size=len(op_info.output_name_list),
+        out_grads_size=grad_idx + 1,
+    )
 
     str = OP_VJP_DEFINE_TEMPLATE.format(
+        check_param=check_param,
         op_class_name=op_class_name,
         op_grad_name=op_grad_name,
         op_phi_name=op_phi_name,
-        res_size=len(op_info.input_name_list),
-        forward_input_output_code=forward_input_output_code,
-        forward_output_grad_code=forward_output_grad_code,
+        backward_input_code=backward_input_code,
         attribute_code=attribute_code,
         call_vjp_code=call_vjp_code,
         stop_gradient_input_grad_code=stop_gradient_input_grad_code,
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op_vjp.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op_vjp.cc
@@ -33,20 +33,24 @@ std::vector<std::vector<pir::OpResult>> AddNOp::Vjp(
     const std::vector<std::vector<pir::OpResult>>& outputs,
     const std::vector<std::vector<pir::Value>>& out_grads,
     const std::vector<std::vector<bool>>& stop_gradients) {
-  AddNOp op_obj = op->dyn_cast<AddNOp>();
-
   VLOG(6) << "Prepare inputs of add_n_grad";
+  PADDLE_ENFORCE_EQ(
+      inputs_.size(),
+      1u,
+      platform::errors::InvalidArgument(
+          "addn op's inputs size should be 1 but now is %d", inputs_.size()));
+  PADDLE_ENFORCE_EQ(
+      outputs.size(),
+      1u,
+      platform::errors::InvalidArgument(
+          "addn op's outputs size should be 1 but now is %d", outputs.size()));
   PADDLE_ENFORCE(
-      op_obj.inputs() != nullptr,
-      paddle::platform::errors::Fatal("addn op's inputs can't be null"));
-  pir::CombineOp combine_op_obj = op_obj.inputs()
-                                      .dyn_cast<pir::OpResult>()
-                                      .owner()
-                                      ->dyn_cast<pir::CombineOp>();
+      inputs_[0].size() != 0,
+      paddle::platform::errors::Fatal("addn op's inputs[0] can't be null"));
   std::vector<Tensor> inputs;
-  for (size_t idx = 0; idx < combine_op_obj.inputs().size(); idx++) {
+  for (size_t idx = 0; idx < inputs_[0].size(); idx++) {
     inputs.emplace_back(
-        std::make_shared<primitive::LazyTensor>(combine_op_obj.inputs()[idx]));
+        std::make_shared<primitive::LazyTensor>(inputs_[0][idx]));
   }
 
   Tensor out_grad(std::make_shared<primitive::LazyTensor>(out_grads[0][0]));
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -363,6 +363,7 @@ def make_output_with_output_grad(op):
         outputs = []
         output_grads = []
         for i, value in enumerate(op.results()):
+            new_value = [value]
             if (
                 value in state.value_to_valuegrad
                 and len(state.value_to_valuegrad[value]) > 1
@@ -405,7 +406,7 @@ def make_output_with_output_grad(op):
                     zero_flag[i] = all(split_zero_flag)
                     grad_values = [value[0] for value in split_output_grad]
                     state.value_to_valuegrad[value] = [grad_values]
-                    outputs.append([info[0] for info in split_outputs])
+                    new_value = [info[0] for info in split_outputs]
                 else:
                     # first case:
                     # this fwd_op's output didn't used by other fwd_op,
@@ -428,7 +429,7 @@ def make_output_with_output_grad(op):
 
                     state.value_to_valuegrad[value] = [[grad_value]]
 
-                    outputs.append([value])
+            outputs.append(new_value)
             output_grads.append(state.value_to_valuegrad[value][0])
 
         return zero_flag, outputs, output_grads
@@ -445,6 +446,7 @@ def make_input_with_input_stopgradient(op):
             op.operands_source(), grad_semantic_info
         ):
             if not grad_semantic:
+                inputs.append([input])
                 continue
             if (
                 input.get_defining_op() is not None
diff --git a/test/cpp/prim/test_vjp.cc b/test/cpp/prim/test_vjp.cc
@@ -236,7 +236,8 @@ TEST(VJP, ConcatBackwardTest) {
   paddle::dialect::FullOp op4 = builder->Build<paddle::dialect::FullOp>(
       std::vector<int64_t>{2, 2}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
   std::vector<std::vector<bool>> stop_gradients{{false, false}};
-  std::vector<std::vector<pir::Value>> inputs{{op1.out(), op1.out()}};
+  std::vector<std::vector<pir::Value>> inputs{{op1.out(), op1.out()},
+                                              {op3.axis()}};
   std::vector<std::vector<pir::OpResult>> outputs{{op3.out()}};
   std::vector<std::vector<pir::Value>> out_grads{{op4.out()}};
   pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.concat");
@@ -303,7 +304,7 @@ TEST(VJP, AddBackwardTest) {
       std::vector<int64_t>{1}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
 
   std::vector<std::vector<bool>> stop_gradients{{false}, {false}};
-  std::vector<std::vector<pir::Value>> inputs{{op1.out(), op2.out()}};
+  std::vector<std::vector<pir::Value>> inputs{{op1.out()}, {op2.out()}};
   std::vector<std::vector<pir::OpResult>> outputs{{op3.out()}};
   std::vector<std::vector<pir::Value>> out_grads{{op4.out()}};
 
@@ -371,7 +372,7 @@ TEST(VJP, Add_BackwardTest) {
       std::vector<int64_t>{1}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
 
   std::vector<std::vector<bool>> stop_gradients{{false}, {false}};
-  std::vector<std::vector<pir::Value>> inputs{{op1.out(), op2.out()}};
+  std::vector<std::vector<pir::Value>> inputs{{op1.out()}, {op2.out()}};
   std::vector<std::vector<pir::OpResult>> outputs{{op3.out()}};
   std::vector<std::vector<pir::Value>> out_grads{{op4.out()}};
 
@@ -439,7 +440,8 @@ TEST(VJP, SplitBackwardTest) {
       std::vector<int64_t>{1, 2}, 1.0, phi::DataType::FLOAT32, phi::CPUPlace());
 
   std::vector<std::vector<bool>> stop_gradients{{false}};
-  std::vector<std::vector<pir::Value>> inputs{{op1.out()}};
+  std::vector<std::vector<pir::Value>> inputs{
+      {op2.x()}, {op2.sections()}, {op2.axis()}};
   std::vector<std::vector<pir::OpResult>> outputs{{op3.outputs()}};
   std::vector<std::vector<pir::Value>> out_grads{{op3.result(0), op4.out()}};
   pir::OpInfo op2_info = ctx->GetRegisteredOpInfo("pd_op.split");
diff --git a/test/prim/pir_prim/test_vjp_prim.py b/test/prim/pir_prim/test_vjp_prim.py
