[PIR] Insert backward ops into right insertion point (#70595)

SigureMo · web-flow · commit cb57b17df0f2 · 2025-01-03T13:07:14.000+08:00
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -1113,7 +1113,7 @@ def calc_gradient_helper(
     grad_outputs: Value | Sequence[Value | None] | None = None,
     no_grad_set: set[Value] | None = None,
 ) -> ValueDict:
-    block = outputs[0].get_defining_op().get_parent_block()
+    block = paddle.base.libpaddle.pir.get_current_insertion_point().block()
     state = State(block)
     if all_stop_gradient_true(block):
         logging.warning(
diff --git a/python/paddle/pir/core.py b/python/paddle/pir/core.py
@@ -217,6 +217,8 @@ def switch_main_program(program, insertion_point=None):
     prev_program = _main_program_
     prev_insertion_point = get_current_insertion_point()
     _main_program_ = program
+    if program == prev_program and insertion_point is None:
+        insertion_point = prev_insertion_point
     if insertion_point is None:
         set_insertion_point_to_block_end(_main_program_.global_block())
     else:
diff --git a/test/dygraph_to_static/test_grad.py b/test/dygraph_to_static/test_grad.py
@@ -200,5 +200,24 @@ def test_run(self):
         np.testing.assert_array_equal(out.stop_gradient, True)
 
 
+def grad_with_if_case(x):
+    y = paddle.tanh(x)
+    if x.numel() > 0:
+        return paddle.grad([y], [x])[0]
+    return paddle.ones_like(x, dtype='float32')
+
+
+class TestGradWithIf(Dy2StTestBase):
+    @test_pir_only
+    def test_grad_with_if(self):
+        fn = grad_with_if_case
+        static_fn = paddle.jit.to_static(fn)
+        x = paddle.randn([2, 2])
+        x.stop_gradient = False
+        dx = fn(x)
+        dx_st = static_fn(x)
+        np.testing.assert_allclose(dx, dx_st)
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/ir/pir/test_ir_backward.py b/test/ir/pir/test_ir_backward.py
@@ -296,7 +296,7 @@ def test_skip_vjp(self):
             z = paddle.nn.functional.relu(y)
             loss = paddle.mean(z)
 
-        paddle.autograd.ir_backward.append_backward(loss)
+            paddle.autograd.ir_backward.append_backward(loss)
         relu_grad_number = 0
         for op in program.global_block().ops:
             if op.name() == "pd_op.relu_grad":
