[PIR] Append backward in sub-block (#59982)

changeyoung98 · web-flow · commit 289282fa217b · 2023-12-15T14:44:09.000+08:00
* deal with subblock op backward

* add test

* fix CI bud

* fix bug
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -308,6 +308,10 @@ void BindBlock(py::module *m) {
           "program",
           [](Block &self) { return self.GetParentOp()->GetParentProgram(); },
           return_value_policy::reference)
+      .def_property_readonly(
+          "get_parent",
+          [](Block &self) { return self.GetParentOp()->GetParent(); },
+          return_value_policy::reference)
       .def_property_readonly("ops",
                              [](Block &self) -> py::list {
                                py::list op_list;
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -246,7 +246,7 @@ def prune_ops(total_ops, inputs_set, outputs_set, no_grad_set):
 
 
 def update_no_grad_set_after_prune(
-    block, effective_forward_ops, no_grad_set, inputs, outputs
+    total_ops, effective_forward_ops, no_grad_set, inputs, outputs
 ):
     '''
     update no_grad_set after forward prune
@@ -256,7 +256,7 @@ def update_no_grad_set_after_prune(
     '''
     inputs_set = set(inputs)
     if inputs_set:
-        for op in block.ops:
+        for op in total_ops:
             if some_in_set(get_real_op_inputs(op), inputs_set):
                 for value in op.results():
                     if value not in no_grad_set:
@@ -783,8 +783,6 @@ def calc_gradient_helper(outputs, inputs, grad_outputs, no_grad_set):
     block = outputs[0].get_defining_op().get_parent_block()
     state = State(block)
 
-    # check all inputs and outputs in the same block
-    check_all_puts(block, inputs, outputs)
     # update no_grad_set if some value stop_gradient=True
     update_no_grad_set_by_stopgradient(block, no_grad_set)
     complete_outputs, _, backward_ops = prepare_grad_outputs(
@@ -793,11 +791,16 @@ def calc_gradient_helper(outputs, inputs, grad_outputs, no_grad_set):
 
     inputs_set = set(inputs)
     outputs_set = set(complete_outputs)
+    total_ops = []
+    if block.get_parent is not None:
+        total_ops += block.get_parent.ops
+    total_ops += block.ops
+
     effective_forward_ops, _ = prune_ops(
-        block.ops, inputs_set, outputs_set, no_grad_set
+        total_ops, inputs_set, outputs_set, no_grad_set
     )
     update_no_grad_set_after_prune(
-        block, effective_forward_ops, no_grad_set, inputs, complete_outputs
+        total_ops, effective_forward_ops, no_grad_set, inputs, complete_outputs
     )
 
     outputs_fwd_set, inputs_fwd_set = prepare_backward_prune_set(
diff --git a/test/ir/pir/test_ir_backward.py b/test/ir/pir/test_ir_backward.py
@@ -14,6 +14,8 @@
 
 import unittest
 
+import numpy as np
+
 import paddle
 from paddle import pir
 from paddle.autograd.ir_backward import grad
@@ -249,5 +251,48 @@ def test_basic_network(self):
             input_grad = grad(res, x)
 
 
+class TestBackward_4(unittest.TestCase):
+    def tearDown(self) -> None:
+        paddle.framework.set_flags({"FLAGS_enable_pir_api": False})
+
+    def test_basic_network(self):
+        if not paddle.framework.in_pir_mode():
+            return
+        program = paddle.static.default_main_program()
+        with paddle.static.program_guard(program):
+            x = paddle.randn((2, 2))
+            x.stop_gradient = False
+            b = paddle.to_tensor([12])
+            grad_x = 0
+            double_x = x * 2
+            pred = b > 0
+
+            def true_func():
+                y = double_x * 3
+                out = grad(y, x)
+                filted_dx = [dxi for dxi in out if dxi is not None]
+                grad_x = filted_dx
+                return grad_x
+
+            def false_func():
+                y = double_x * 4
+                out = grad(y, x)
+                filted_dx = [dxi for dxi in out if dxi is not None]
+                grad_x = filted_dx
+                return grad_x
+
+            out = paddle.static.nn.cond(pred, true_func, false_func)
+
+            place = (
+                paddle.base.CUDAPlace(0)
+                if paddle.base.core.is_compiled_with_cuda()
+                else paddle.base.CPUPlace()
+            )
+            exe = paddle.static.Executor(place)
+            (grad_x,) = exe.run(program, fetch_list=[out])
+            res = np.full([2, 2], 6.0, dtype='float32')
+            self.assertEqual((grad_x == res).all(), True)
+
+
 if __name__ == "__main__":
     unittest.main()
