[autodiff] Fix missing global load stmt in independent blocks

taichi/transforms/auto_diff.cpp

@@ -45,6 +45,22 @@ class IndependentBlocksJudger : public BasicStmtVisitor {
     }
   }
 
+  void visit(GlobalLoadStmt *stmt) override {
+    // We don't need to check the global load inside the range for-loops
+    // because
+    // 1. If the range for-loop is innermost, they will be captured by
+    // MakeAdjoint anyway
+    // 2. If the range for-loop is not innermost, they will be processed by
+    // another IndependentBlocksJudger
+    if (is_inside_loop_)
+      return;
+    // TODO: handle external ptr stmt after autodiff supporting ndarray
+    if (stmt->src->is<GlobalPtrStmt>() &&
+        stmt->src->as<GlobalPtrStmt>()->snode->has_adjoint()) {
+      qualified_glb_operations_ = true;
+    }
+  }
+
   void visit(RangeForStmt *stmt) override {
     inner_most_loop_ = false;
     is_inside_loop_ = true;

tests/python/test_ad_for.py

@@ -980,3 +980,33 @@ def compute_y():
 
     assert y[None] == 78.0
     assert x.grad[None] == 78.0
+
+
+@test_utils.test(require=ti.extension.adstack)
+def test_ib_global_load():
+    N = 10
+    a = ti.field(ti.f32, shape=N, needs_grad=True)
+    b = ti.field(ti.i32, shape=N)
+    p = ti.field(ti.f32, shape=N, needs_grad=True)
+
+    @ti.kernel
+    def compute():
+        for i in range(N):
+            val = a[i]
+            for j in range(b[i]):
+                p[i] += i
+            p[i] = val * i
+
+    for i in range(N):
+        a[i] = i
+        b[i] = 2
+
+    compute()
+
+    for i in range(N):
+        assert p[i] == i * i
+        p.grad[i] = 1
+
+    compute.grad()
+    for i in range(N):
+        assert a.grad[i] == i