【pir】modify ir Backward for prune

python/paddle/autograd/ir_backward.py

@@ -58,9 +58,10 @@ def update_no_grad_set_by_stopgradient(block, no_grad_set):
                 no_grad_set.add(value)
 
 
-def update_bwdop_structure(backward_ops, op_to_opgrad_list, grad_op):
-    backward_ops.append(grad_op)
-    op_to_opgrad_list.append(grad_op)
+def update_bwdop_structure(backward_ops, op_to_opgrad_list, grad_op_list):
+    for grad_op in grad_op_list:
+        backward_ops.append(grad_op)
+        op_to_opgrad_list.append(grad_op)
 
 
 def prepare_grad_outputs(grad_outputs, outputs, state):
@@ -87,18 +88,19 @@ def prepare_grad_outputs(grad_outputs, outputs, state):
     for i, grad in enumerate(grad_outputs):
         output = outputs[i]
         # fwd : op1 -> op2 -> op3 -> output
-        # bwd : op1G <- op2G <- op3G <- outputG <- fillop/feedop
+        # bwd : op1G <- op2G <- op3G <- outputG <- full_likeop/feedop
         if grad is None:
             output_grad = paddle.full_like(
                 output,
                 1.0,
                 dtype=output.dtype,
             )
-            fillop = output_grad.get_defining_op()
+            full_likeop = output_grad.get_defining_op()
+            fullop = full_likeop.operand_source(1).get_defining_op()
             update_bwdop_structure(
                 backward_ops,
                 state.op_to_opgrad[output.get_defining_op()],
-                fillop,
+                [full_likeop, fullop],
             )
             state.value_to_valuegrad[output] = [[output_grad]]
         else:
@@ -116,7 +118,7 @@ def prepare_grad_outputs(grad_outputs, outputs, state):
             update_bwdop_structure(
                 backward_ops,
                 state.op_to_opgrad[output.get_defining_op()],
-                feedop,
+                [feedop],
             )
             state.value_to_valuegrad[output] = [[grad]]
 
@@ -138,12 +140,13 @@ def prepare_grad_outputs(grad_outputs, outputs, state):
                     0.0,
                     opresult.dtype,
                 )
-                fillop = grad_value.get_defining_op()
+                full_likeop = grad_value.get_defining_op()
+                fullop = full_likeop.operand_source(1).get_defining_op()
 
                 update_bwdop_structure(
                     backward_ops,
                     state.op_to_opgrad[opresult.get_defining_op()],
-                    fillop,
+                    [full_likeop, fullop],
                 )
                 state.value_to_valuegrad[opresult] = [[grad_value]]
 
@@ -383,11 +386,9 @@ def make_output_with_output_grad(op):
                 combineop = bwd_block.ops[len(bwd_block.ops) - 2]
                 sumop = bwd_block.ops[len(bwd_block.ops) - 1]
                 update_bwdop_structure(
-                    backward_ops, state.op_to_opgrad[op], combineop
-                )
-                update_bwdop_structure(
-                    backward_ops, state.op_to_opgrad[op], sumop
+                    backward_ops, state.op_to_opgrad[op], [combineop, sumop]
                 )
+
                 state.value_to_valuegrad[value] = [[sumop.result(0)]]
                 state.value_to_sumvaluegrad[value] = state.value_to_valuegrad[
                     value
@@ -426,10 +427,13 @@ def make_output_with_output_grad(op):
                         0.0,
                         dtype=value.dtype,
                     )
-                    fillop = grad_value.get_defining_op()
+                    full_likeop = grad_value.get_defining_op()
+                    fullop = full_likeop.operand_source(1).get_defining_op()
 
                     update_bwdop_structure(
-                        backward_ops, state.op_to_opgrad[op], fillop
+                        backward_ops,
+                        state.op_to_opgrad[op],
+                        [full_likeop, fullop],
                     )
                     zero_flag[i] = True
 
@@ -548,10 +552,12 @@ def update_input_grad_map(op, input_grads):
             after_ops_num = len(bwd_block.ops)
 
             # update grad_op structure
-            for i in range(before_ops_num, after_ops_num):
-                update_bwdop_structure(
-                    backward_ops, state.op_to_opgrad[op], bwd_block.ops[i]
-                )
+            bwd_ops = [
+                bwd_block.ops[i] for i in range(before_ops_num, after_ops_num)
+            ]
+            update_bwdop_structure(
+                backward_ops, state.op_to_opgrad[op], bwd_ops
+            )
 
             # update input_grad map
             update_input_grad_map(op, input_grads)
@@ -570,10 +576,9 @@ def update_input_grad_map(op, input_grads):
                         combineop = bwd_block.ops[len(bwd_block.ops) - 2]
                         sumop = bwd_block.ops[len(bwd_block.ops) - 1]
                         update_bwdop_structure(
-                            backward_ops, state.op_to_opgrad[op], combineop
-                        )
-                        update_bwdop_structure(
-                            backward_ops, state.op_to_opgrad[op], sumop
+                            backward_ops,
+                            state.op_to_opgrad[op],
+                            [combineop, sumop],
                         )
                         state.value_to_valuegrad[value] = [[sumop.result(0)]]
                         state.value_to_sumvaluegrad[
@@ -585,20 +590,35 @@ def update_input_grad_map(op, input_grads):
                 state.op_to_opgrad[op] = []
 
 
-def create_backward_prune_set(inputs, outputs, no_grad_set, state):
-    outputs_set = set()
+def prepare_backward_prune_set(inputs, outputs):
+    outputs_fwd_set = set()
     for input_ in inputs:
         if not input_.use_empty():
             for item in input_.first_use().owner().operands_source():
-                if state.value_to_valuegrad[item] != []:
-                    outputs_set.add(state.value_to_valuegrad[item][0][0])
+                outputs_fwd_set.add(item)
         else:
             logging.warning("input privided by inputs has no use")
 
-    inputs_set = set()
+    inputs_fwd_set = set()
     for output in outputs:
-        if state.value_to_valuegrad[output] != []:
-            inputs_set.add(state.value_to_valuegrad[output][0][0])
+        inputs_fwd_set.add(output)
+
+    return outputs_fwd_set, inputs_fwd_set
+
+
+def create_backward_prune_set(
+    outputs_fwd_set, inputs_fwd_set, no_grad_set, state
+):
+    outputs_set = set()
+    for item in outputs_fwd_set:
+        if state.value_to_valuegrad[item] != []:
+            outputs_set.add(state.value_to_valuegrad[item][0][0])
+
+    inputs_set = set()
+    for item in inputs_fwd_set:
+        if state.value_to_valuegrad[item] != []:
+            inputs_set.add(state.value_to_valuegrad[item][0][0])
+
     inputs_set_tmp = set()
     for out_grad in inputs_set:
         if not out_grad.use_empty():
@@ -660,13 +680,19 @@ def calc_gradient_helper(outputs, inputs, grad_outputs, no_grad_set):
         block, effective_forward_ops, no_grad_set, inputs, complete_outputs
     )
 
+    outputs_fwd_set, inputs_fwd_set = prepare_backward_prune_set(
+        inputs, complete_outputs
+    )
+
     append_backward_ops(
         block, block, effective_forward_ops, no_grad_set, backward_ops, state
     )
+
     # now value_to_valuegrad should be value <-> value (add sum op for the same values's gradvalue)
     outputs_set, inputs_set, no_gradvar_set = create_backward_prune_set(
-        inputs, complete_outputs, no_grad_set, state
+        outputs_fwd_set, inputs_fwd_set, no_grad_set, state
     )
+
     _, remove_ops = prune_ops(
         backward_ops, inputs_set, outputs_set, no_gradvar_set
     )

test/ir/pir/test_ir_backward.py

@@ -101,7 +101,7 @@ def test_no_grad_set(self):
             out = paddle.mean(tanh_out)
             input_grad = grad(out, input, no_grad_vars=[input])
             self.assertEqual(
-                pir_program.global_block().ops[-1].name(), "pd_op.full"
+                pir_program.global_block().ops[-1].name(), "pd_op.mean"
             )
 
     def test_split(self):

test/legacy_test/test_segment_ops.py

@@ -127,7 +127,7 @@ def test_check_output(self):
         self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(["X"], "Out")
+        self.check_grad(["X"], "Out", check_pir=True)
 
     def convert_bf16(self):
         if self.dtype == np.uint16:
@@ -277,7 +277,7 @@ def test_check_output(self):
         self.check_output_with_place(self.place, check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad_with_place(self.place, ["X"], "Out")
+        self.check_grad_with_place(self.place, ["X"], "Out", check_pir=True)
 
 
 @unittest.skipIf(
@@ -300,6 +300,7 @@ def test_check_grad(self):
             ["X"],
             "Out",
             user_defined_grads=[self.gradient],
+            check_pir=True,
         )
 
 
@@ -323,6 +324,7 @@ def test_check_grad(self):
             ["X"],
             "Out",
             user_defined_grads=[self.gradient],
+            check_pir=True,
         )
 
 
@@ -341,7 +343,7 @@ def test_check_output(self):
         self.check_output_with_place(self.place, check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad_with_place(self.place, ["X"], "Out")
+        self.check_grad_with_place(self.place, ["X"], "Out", check_pir=True)
 
 
 class API_SegmentOpsTest(unittest.TestCase):