[Bug] Fix Erroneous handling of ndarray in real function in CFG

taichi/analysis/alias_analysis.cpp

@@ -14,6 +14,30 @@ AliasResult alias_analysis(Stmt *var1, Stmt *var2) {
   if (!var1 || !var2)
     return AliasResult::different;
 
+  if (var1->is<ExternalTensorBasePtrStmt>() ||
+      var2->is<ExternalTensorBasePtrStmt>()) {
+    auto *base = var1->cast<ExternalTensorBasePtrStmt>();
+    Stmt *other = var2;
+    if (!base) {
+      base = var2->cast<ExternalTensorBasePtrStmt>();
+      other = var1;
+    }
+    auto *external_ptr = other->cast<ExternalPtrStmt>();
+    if (!external_ptr) {
+      if (auto *matrix_ptr = other->cast<MatrixPtrStmt>()) {
+        external_ptr = matrix_ptr->origin->cast<ExternalPtrStmt>();
+      }
+      if (!external_ptr)
+        return AliasResult::different;
+    }
+    if (base->is_grad != external_ptr->is_grad)
+      return AliasResult::different;
+    if (base->arg_id == external_ptr->base_ptr->as<ArgLoadStmt>()->arg_id) {
+      return AliasResult::uncertain;
+    }
+    return AliasResult::different;
+  }
+
   // TODO: further optimize with offset inside MatrixPtrStmt
   // If at least one of var1 and var2 is local, they will be treated here.
   auto retrieve_local = [&](Stmt *var) {

taichi/analysis/gather_func_store_dests.cpp

@@ -62,7 +62,17 @@ class GatherFuncStoreDests : public BasicStmtVisitor {
       return;
     }
     auto result = irpass::analysis::get_store_destination(stmt);
-    results_.insert(result.begin(), result.end());
+    for (const auto &dest : result) {
+      if (dest->is<ExternalPtrStmt>()) {
+        continue;
+      }
+      if (auto matrix_ptr = dest->cast<MatrixPtrStmt>()) {
+        if (matrix_ptr->origin->is<ExternalPtrStmt>()) {
+          continue;
+        }
+      }
+      results_.insert(dest);
+    }
   }
 
   void visit(FuncCallStmt *stmt) override {

tests/python/test_ndarray.py

@@ -1129,3 +1129,22 @@ def test(x: ti.types.ndarray(ndim=1)) -> vec3:
     x = ti.Vector.ndarray(3, ti.f32, shape=(1))
     x[0] = vec3(1, 2, 3)
     assert (test(x) == vec3(1, 2, 3)).all()
+
+
+@test_utils.test(arch=[ti.cpu, ti.cuda])
+def test_real_func_write_ndarray_cfg():
+    @ti.experimental.real_func
+    def bar(a: ti.types.ndarray(ndim=1)):
+        a[0] = vec3(1)
+
+    @ti.kernel
+    def foo(
+        a: ti.types.ndarray(ndim=1),
+    ):
+        a[0] = vec3(3)
+        bar(a)
+        a[0] = vec3(3)
+
+    a = ti.Vector.ndarray(3, float, shape=(2,))
+    foo(a)
+    assert (a[0] == vec3(3)).all()