From 2a26e68d0c9edda899dd4e8c3f75c77fa493df69 Mon Sep 17 00:00:00 2001
From: winter-wang <1030748926@qq.com>
Date: Fri, 29 Mar 2024 17:43:10 +0800
Subject: [PATCH 1/2] fix auto backward bug

---
 python/paddle/autograd/backward_utils.py |  8 ++++++++
 python/paddle/autograd/ir_backward.py    |  4 ++++
 test/ir/pir/test_ir_backward.py          | 25 ++++++++++++++++++++++++
 3 files changed, 37 insertions(+)

diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
index ff6c42613d06bf..0c82ce1aeaf213 100644
--- a/python/paddle/autograd/backward_utils.py
+++ b/python/paddle/autograd/backward_utils.py
@@ -444,6 +444,14 @@ def all_stop_gradient_true(block):
     return True
 
 
+def all_input_stop_gradient_true(list_of_list):
+    for list_ in list_of_list:
+        for stop_gradient in list_:
+            if stop_gradient is False:
+                return False
+    return True
+
+
 def all_output_grad_none(list_of_list):
     for list_ in list_of_list:
         for value in list_:
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
index 551e55a18b942a..bd8741598d45bc 100644
--- a/python/paddle/autograd/ir_backward.py
+++ b/python/paddle/autograd/ir_backward.py
@@ -22,6 +22,7 @@
     ValueDict,
     ValueSet,
     _as_list,
+    all_input_stop_gradient_true,
     all_output_grad_none,
     all_stop_gradient_true,
     argument_to_value,
@@ -642,6 +643,9 @@ def append_yield(
                         len(output_grads) == 0
                         or all(zero_flag)
                         or all_output_grad_none(output_grads)
+                        or all_input_stop_gradient_true(
+                            input_grad_stopgradients
+                        )
                     ) and op.name() not in [
                         "pd_op.while",
                         "pd_op.if",
diff --git a/test/ir/pir/test_ir_backward.py b/test/ir/pir/test_ir_backward.py
index 5e4f5386a1cdac..3f8a77eed354fe 100644
--- a/test/ir/pir/test_ir_backward.py
+++ b/test/ir/pir/test_ir_backward.py
@@ -292,6 +292,31 @@ def false_func():
             self.assertEqual((grad_x == res).all(), True)
 
 
+class TestBackward_5(unittest.TestCase):
+    def tearDown(self) -> None:
+        paddle.framework.set_flags({"FLAGS_enable_pir_api": False})
+
+    def test_skip_vjp(self):
+        if not paddle.framework.in_pir_mode():
+            return
+        program = paddle.static.Program()
+        with paddle.static.program_guard(program):
+            x = paddle.static.data('x', [4, 4], 'float32')
+            x.stop_gradient = True
+            y = paddle.nn.functional.relu(x)
+            y.stop_gradient = False
+            z = paddle.nn.functional.relu(y)
+            loss = paddle.mean(z)
+
+        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":
+                relu_grad_number += 1
+
+        self.assertEqual(relu_grad_number, 1)
+
+
 class TestValueSet(unittest.TestCase):
     def setUp(self) -> None:
         with paddle.pir_utils.IrGuard():

From 21a1eb4b0983849c605865d2b74f33c0d0b3fd73 Mon Sep 17 00:00:00 2001
From: winter-wang <1030748926@qq.com>
Date: Fri, 29 Mar 2024 19:56:45 +0800
Subject: [PATCH 2/2] fix ci error.

---
 python/paddle/autograd/ir_backward.py                | 11 ++++++++---
 test/auto_parallel/pir/test_to_static_pir_program.py |  2 --
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
index bd8741598d45bc..8b72bb35a04cc5 100644
--- a/python/paddle/autograd/ir_backward.py
+++ b/python/paddle/autograd/ir_backward.py
@@ -643,9 +643,6 @@ def append_yield(
                         len(output_grads) == 0
                         or all(zero_flag)
                         or all_output_grad_none(output_grads)
-                        or all_input_stop_gradient_true(
-                            input_grad_stopgradients
-                        )
                     ) and op.name() not in [
                         "pd_op.while",
                         "pd_op.if",
@@ -653,6 +650,14 @@ def append_yield(
                     ]:
                         continue
 
+                    if all_input_stop_gradient_true(
+                        input_grad_stopgradients
+                    ) and op.name() not in [
+                        "pd_op.array_read",
+                        "pd_op.array_write_",
+                        "pd_op.increment_",
+                    ]:
+                        continue
                     if op.name() == "pd_op.if":
                         origin_inputs = get_real_op_inputs(op)
                         for sub_block in op.blocks():
diff --git a/test/auto_parallel/pir/test_to_static_pir_program.py b/test/auto_parallel/pir/test_to_static_pir_program.py
index 3085e3a726de0b..486011ad0e77b5 100644
--- a/test/auto_parallel/pir/test_to_static_pir_program.py
+++ b/test/auto_parallel/pir/test_to_static_pir_program.py
@@ -139,8 +139,6 @@ def test_to_static_program(self):
         backward_op_list = [
             "pd_op.sgd_",
             "pd_op.sgd_",
-            "pd_op.relu_grad",
-            "pd_op.c_allreduce_sum_",
             "pd_op.matmul_grad",
             "pd_op.relu_grad",
             "pd_op.matmul_grad",