[autodiff] Skip store forwarding to keep the GlobalLoadStmt alive

taichi/ir/control_flow_graph.cpp

@@ -275,7 +275,21 @@ bool CFGNode::store_to_load_forwarding(bool after_lower_access) {
       }
     } else if (auto global_load = stmt->cast<GlobalLoadStmt>()) {
       if (!after_lower_access) {
-        result = get_store_forwarding_data(global_load->src, i);
+        bool store_forwarding = true;
+        if (auto global_ptr = global_load->src->cast<GlobalPtrStmt>()) {
+          TI_ASSERT(global_ptr->width() == 1);
+          auto &snodes = global_ptr->snodes;
+          if (snodes[0]->has_adjoint()) {
+            // Has adjoint SNode. Skip the store forwarding
+            // to keep the global load chain,
+            // so that the grad of intermidiate variable can be computed
+            // by GlobalLoadStmt
+            store_forwarding = false;
+          }
+        }
+        if (store_forwarding) {
+          result = get_store_forwarding_data(global_load->src, i);
+        }
       }
     }
     if (result) {

taichi/ir/snode.cpp

@@ -302,7 +302,7 @@ void SNode::end_shared_exp_placement() {
 }
 
 bool SNode::is_primal() const {
-  return grad_info->is_primal();
+  return grad_info && grad_info->is_primal();
 }
 
 bool SNode::has_adjoint() const {

tests/python/test_ad_basics.py

@@ -351,3 +351,31 @@ def calc_loss(input_field: ti.template(), loss: ti.template()):
     expected = np.modf(randoms)[0] * 2
     for i in range(n):
         assert grads[i] == test_utils.approx(expected[i], rel=1e-4)
+
+
+@test_utils.test()
+def test_ad_global_store_forwarding():
+    x = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
+    a = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
+    b = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
+    c = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
+    d = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
+    e = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
+
+    @ti.kernel
+    def func():
+        a[None] = x[None]
+        b[None] = a[None] * 2
+        c[None] = b[None] * 3
+        d[None] = c[None] * 4
+        e[None] = d[None] * 5
+
+    x[None] = 1
+
+    with ti.ad.Tape(loss=e):
+        func()
+    assert x.grad[None] == 120.0
+    assert a.grad[None] == 120.0
+    assert b.grad[None] == 60.0
+    assert c.grad[None] == 20.0
+    assert d.grad[None] == 5.0