[opt] Properly handle FuncCallStmt in CFG and simplify passes

taichi/ir/control_flow_graph.cpp

@@ -217,6 +217,9 @@ Stmt *CFGNode::get_store_forwarding_data(Stmt *var, int position) const {
 
   // Check if store_stmt will ever influence the value of var
   auto may_contain_address = [](Stmt *store_stmt, Stmt *var) {
+    if (store_stmt->is<FuncCallStmt>()) {
+      return true;
+    }
     for (auto store_ptr : irpass::analysis::get_store_destination(store_stmt)) {
       if (var->is<MatrixPtrStmt>() && !store_ptr->is<MatrixPtrStmt>()) {
         // check for aliased address with var
@@ -976,7 +979,8 @@ void ControlFlowGraph::reaching_definition_analysis(bool after_lower_access) {
   for (int i = 0; i < num_nodes; i++) {
     for (int j = nodes[i]->begin_location; j < nodes[i]->end_location; j++) {
       auto stmt = nodes[i]->block->statements[j].get();
-      if ((stmt->is<MatrixPtrStmt>() &&
+      if (stmt->is<FuncCallStmt>() ||
+          (stmt->is<MatrixPtrStmt>() &&
            stmt->as<MatrixPtrStmt>()->origin->is<AllocaStmt>()) ||
           (!after_lower_access &&
            (stmt->is<GlobalPtrStmt>() || stmt->is<ExternalPtrStmt>() ||

taichi/transforms/simplify.cpp

@@ -95,9 +95,6 @@ class BasicBlockSimplify : public IRVisitor {
                 }
                 continue;
               }
-              if (block->statements[j]->is<FuncCallStmt>()) {
-                has_store = true;
-              }
               if (!irpass::analysis::gather_statements(
                        block->statements[j].get(),
                        [&](Stmt *s) {
@@ -107,6 +104,8 @@ class BasicBlockSimplify : public IRVisitor {
                          else if (auto atomic = s->cast<AtomicOpStmt>())
                            return irpass::analysis::maybe_same_address(
                                atomic->dest, stmt->src);
+                         else if (auto func_call = s->cast<FuncCallStmt>())
+                           return true;
                          else
                            return false;
                        })

tests/python/test_function.py

@@ -155,13 +155,15 @@ def inc(x: ti.template()):
         x[None] += 1
 
     @ti.kernel
-    def run_func():
+    def run_func(a: ti.u1):
         x[None] = 10
         y[None] = 20
-        inc(x)
+        if a:
+            inc(x)
         answer[0] = x[None]
         answer[1] = y[None]
-        inc(y)
+        if a:
+            inc(y)
         answer[2] = x[None]
         answer[3] = y[None]
 
@@ -172,7 +174,7 @@ def verify():
         assert answer[3] == 21
 
     run_kernel()
-    run_func()
+    run_func(True)
     verify()
 
 
@@ -324,28 +326,22 @@ def bar(a: ti.i32) -> ti.i32:
 
 @test_utils.test(arch=[ti.cpu, ti.cuda], debug=True)
 def test_ref():
-    # FIXME:
-    #  Failed test if we put assert inside kernel. But test passes with assertion inside kernel if we insert print
-    #  statement after assert.
     @ti.experimental.real_func
     def foo(a: ti.ref(ti.f32)):
         a = 7
 
     @ti.kernel
-    def bar() -> ti.f32:
+    def bar():
         a = 5.0
         foo(a)
-        return a
+        assert a == 7
 
-    assert bar() == 7
+    bar()
 
 
 @test_utils.test(arch=[ti.cpu, ti.cuda], debug=True)
 def test_ref_atomic():
-    # FIXME:
-    #  1. Failed test on Pascal (and potentially older) architecture
-    #  2. Failed test if we put assert inside kernel. But test passes with assertion inside kernel if we insert print
-    #     statement after assert.
+    # FIXME: failed test on Pascal (and potentially older) architecture.
     # Please remove this guardiance when you fix this issue
     cur_arch = ti.lang.impl.get_runtime().prog.config().arch
     if cur_arch == ti.cuda and ti.lang.impl.get_cuda_compute_capability() < 70:
@@ -358,12 +354,12 @@ def foo(a: ti.ref(ti.f32)):
         a += a
 
     @ti.kernel
-    def bar() -> ti.f32:
+    def bar():
         a = 5.0
         foo(a)
-        return a
+        assert a == 10.0
 
-    assert bar() == 10.0
+    bar()
 
 
 @test_utils.test(arch=[ti.cpu, ti.cuda], debug=True)