[opt] Add more strict alias analysis for ExternalPtrStmt

This is not likely not 100% safe, how about banning "passing the same ndarray as different arguments to a taichi kernel"?
taichi-dev · Jan 11, 2022 · 041f1bc · 041f1bc
1 parent 344b17f
commit 041f1bc
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/taichi/analysis/alias_analysis.cpp b/taichi/analysis/alias_analysis.cpp
@@ -99,8 +99,12 @@ AliasResult alias_analysis(Stmt *var1, Stmt *var2) {
       return AliasResult::different;
     auto ptr1 = var1->as<ExternalPtrStmt>();
     auto ptr2 = var2->as<ExternalPtrStmt>();
-    if (ptr1->base_ptrs[0] != ptr2->base_ptrs[0])
-      return AliasResult::uncertain;
+    if (ptr1->base_ptrs[0] != ptr2->base_ptrs[0]) {
+      auto base1 = ptr1->base_ptrs[0]->as<ArgLoadStmt>();
+      auto base2 = ptr2->base_ptrs[0]->as<ArgLoadStmt>();
+      return base1->arg_id != base2->arg_id ? AliasResult::different
+                                            : AliasResult::uncertain;
+    }
     TI_ASSERT(ptr1->indices.size() == ptr2->indices.size());
     bool uncertain = false;
     for (int i = 0; i < (int)ptr1->indices.size(); i++) {

diff --git a/tests/cpp/analysis/alias_analysis_test.cpp b/tests/cpp/analysis/alias_analysis_test.cpp
@@ -143,7 +143,7 @@ TEST(AliasAnalysis, ExternalPtr_DiffPtr) {
   auto *eptr2 = builder.create_external_ptr(arg2, indices);
 
   const auto aa = alias_analysis(eptr1, eptr2);
-  EXPECT_EQ(aa, AliasResult::uncertain);
+  EXPECT_EQ(aa, AliasResult::different);
 }
 
 }  // namespace analysis