From 70619b0dfa5326cee6bc22db190c66a869f2f1b5 Mon Sep 17 00:00:00 2001
From: xumingkuan <xumingkuan0721@126.com>
Date: Wed, 1 Jul 2020 16:47:02 -0400
Subject: [PATCH 1/2] [bug] [opt] Fix CFG ignoring local atomics after
 lower_access

---
 taichi/analysis/build_cfg.cpp    | 16 +++++++++++-----
 taichi/ir/control_flow_graph.cpp |  3 +--
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/taichi/analysis/build_cfg.cpp b/taichi/analysis/build_cfg.cpp
index c5f314b4023ea..80ed6c3543562 100644
--- a/taichi/analysis/build_cfg.cpp
+++ b/taichi/analysis/build_cfg.cpp
@@ -15,7 +15,7 @@ class CFGBuilder : public IRVisitor {
   int current_stmt_id;
   int begin_location;
   std::vector<CFGNode *> prev_nodes;
-  bool in_offloaded_for;
+  bool in_parallel_for;
 
  public:
   CFGBuilder()
@@ -23,7 +23,7 @@ class CFGBuilder : public IRVisitor {
         last_node_in_current_block(nullptr),
         current_stmt_id(-1),
         begin_location(-1),
-        in_offloaded_for(false) {
+        in_parallel_for(false) {
     allow_undefined_visitor = true;
     invoke_default_visitor = true;
     graph = std::make_unique<ControlFlowGraph>();
@@ -40,7 +40,7 @@ class CFGBuilder : public IRVisitor {
 
   CFGNode *new_node(int next_begin_location) {
     auto node = graph->push_back(current_block, begin_location, current_stmt_id,
-                                 in_offloaded_for, last_node_in_current_block);
+                                 in_parallel_for, last_node_in_current_block);
     for (auto &prev_node : prev_nodes) {
       CFGNode::add_edge(prev_node, node);
     }
@@ -125,11 +125,17 @@ class CFGBuilder : public IRVisitor {
   }
 
   void visit(RangeForStmt *stmt) override {
+    auto old_in_parallel_for = in_parallel_for;
+    in_parallel_for = true;
     visit_loop(stmt->body.get(), new_node(-1), false);
+    in_parallel_for = old_in_parallel_for;
   }
 
   void visit(StructForStmt *stmt) override {
+    auto old_in_parallel_for = in_parallel_for;
+    in_parallel_for = true;
     visit_loop(stmt->body.get(), new_node(-1), false);
+    in_parallel_for = old_in_parallel_for;
   }
 
   void visit(OffloadedStmt *stmt) override {
@@ -149,10 +155,10 @@ class CFGBuilder : public IRVisitor {
       auto block_begin_index = graph->size();
       if (stmt->task_type == OffloadedStmt::TaskType::range_for ||
           stmt->task_type == OffloadedStmt::TaskType::struct_for) {
-        in_offloaded_for = true;
+        in_parallel_for = true;
       }
       stmt->body->accept(this);
-      in_offloaded_for = false;
+      in_parallel_for = false;
       prev_nodes.push_back(graph->back());
       // Container statements don't belong to any CFGNodes.
       begin_location = offload_stmt_id + 1;
diff --git a/taichi/ir/control_flow_graph.cpp b/taichi/ir/control_flow_graph.cpp
index f0a726d50a808..cb92bb04030ef 100644
--- a/taichi/ir/control_flow_graph.cpp
+++ b/taichi/ir/control_flow_graph.cpp
@@ -101,8 +101,7 @@ void CFGNode::reaching_definition_analysis(bool after_lower_access) {
     auto data_source_ptr = irpass::analysis::get_store_destination(stmt);
     if (data_source_ptr) {
       // stmt provides a data source
-      if (after_lower_access &&
-          !(stmt->is<AllocaStmt>() || stmt->is<LocalStoreStmt>())) {
+      if (after_lower_access && !(data_source_ptr->is<AllocaStmt>())) {
         // After lower_access, we only analyze local variables.
         continue;
       }

From b65206043577425b303d6b053db743b73fd3369b Mon Sep 17 00:00:00 2001
From: xumingkuan <xumingkuan0721@126.com>
Date: Wed, 1 Jul 2020 16:56:10 -0400
Subject: [PATCH 2/2] refine in_parallel_for

---
 taichi/analysis/build_cfg.cpp | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/taichi/analysis/build_cfg.cpp b/taichi/analysis/build_cfg.cpp
index 80ed6c3543562..b876356a09095 100644
--- a/taichi/analysis/build_cfg.cpp
+++ b/taichi/analysis/build_cfg.cpp
@@ -15,6 +15,7 @@ class CFGBuilder : public IRVisitor {
   int current_stmt_id;
   int begin_location;
   std::vector<CFGNode *> prev_nodes;
+  OffloadedStmt *current_offload;
   bool in_parallel_for;
 
  public:
@@ -23,6 +24,7 @@ class CFGBuilder : public IRVisitor {
         last_node_in_current_block(nullptr),
         current_stmt_id(-1),
         begin_location(-1),
+        current_offload(nullptr),
         in_parallel_for(false) {
     allow_undefined_visitor = true;
     invoke_default_visitor = true;
@@ -126,19 +128,22 @@ class CFGBuilder : public IRVisitor {
 
   void visit(RangeForStmt *stmt) override {
     auto old_in_parallel_for = in_parallel_for;
-    in_parallel_for = true;
+    if (!current_offload)
+      in_parallel_for = true;
     visit_loop(stmt->body.get(), new_node(-1), false);
     in_parallel_for = old_in_parallel_for;
   }
 
   void visit(StructForStmt *stmt) override {
     auto old_in_parallel_for = in_parallel_for;
-    in_parallel_for = true;
+    if (!current_offload)
+      in_parallel_for = true;
     visit_loop(stmt->body.get(), new_node(-1), false);
     in_parallel_for = old_in_parallel_for;
   }
 
   void visit(OffloadedStmt *stmt) override {
+    current_offload = stmt;
     if (stmt->prologue) {
       auto before_offload = new_node(-1);
       int offload_stmt_id = current_stmt_id;
@@ -174,6 +179,7 @@ class CFGBuilder : public IRVisitor {
       begin_location = offload_stmt_id + 1;
       CFGNode::add_edge(before_offload, graph->nodes[block_begin_index].get());
     }
+    current_offload = nullptr;
   }
 
   void visit(Block *block) override {