Repository: taichi-dev/taichi (language: python)

Issue:
Thread Local Storage Optimisation Bug
There issue provides a minimal repro example for the bug in #7255, and discuss why this is likely a bug with our TLS optimization.

Example:
```py
import taichi as ti
ti.init(ti.cpu, print_ir = True, advanced_optimization = True)

N = 1
p = ti.field(float, N)

@ti.kernel
def k():
    for n in range(N):
        for t in range(2):
            print("outside", p[n])
            for m in range(1):
                print("inside", p[n])
            p[n] = p[n] + 1.0

k()
```
The result printed is
```
outside 0.000000
inside 0.000000
outside 1.000000
inside 0.000000
```
while the expected result is
```
outside 0.000000
inside 0.000000
outside 1.000000
inside 1.000000
```

I suspect this is due to a bug in the TLS optimisation. If we look at the IR, we can see that the reference to `p[n]` in `print("outside", p[n])` and in `p[n] = p[n] + 1.0` are optimised into local storages, while the reference to `p[n]` in `print("inside", p[n])` is still a global read:
```
kernel {
$0 = offloaded range_for(0, 1) grid_dim=0 block_dim=32 
body {
  <f32> $1 = const 1.0
  <i32> $2 = const 1
  <i32> $3 = const 0
  <i32> $4 = loop $0 index 0
  <i32> $5 = const 2
  <f32> $6 = alloca
  <*gen> $7 = get root [S0root][root]
  <*gen> $8 = [S0root][root]::lookup($7, $3) activate = false
  <*gen> $9 = get child [S0root->S1dense] $8
  <*gen> $10 = [S1dense][dense]::lookup($9, $4) activate = false
  <*f32> $11 = get child [S1dense->S2place<f32>] $10
  <f32> $12 = global load $11
  <f32> $13 : local store [$6 <- $12]
  $14 : for in range($3, $5) block_dim=adaptive {
    <f32> $15 = local load [$6]
    print "outside ", $15, "\n"
    $17 : for in range($3, $2) block_dim=adaptive {
      print "inside ", $12, "\n"
    }
    <f32> $19 = add $15 $1
    <f32> $20 : local store [$6 <- $19]
  }
  <f32> $21 = local load [$6]
  $22 : global store [$11 <- $21]
}
}
```

Also, I verified that if `advanced_optimization` is disabled, the program behaves correctly.

cc @strongoier 


Discussion / hints:
A similar bug report from https://forum.taichi-lang.cn/t/topic/3853/4:

```python
import taichi as ti
ti.init()
vec3 = ti.math.vec3

@ti.dataclass
class Particle:
    r: vec3

particle = Particle.field(shape=1)

@ti.kernel
def simulate():
    for i in range(1):
        for _ in range(5):
            for __ in range(1):
                print('inside', particle[i].r)
            particle[i].r += vec3(1)
            print('outside', particle[i].r)

simulate()
```

The output is 

```
inside [0.000000, 0.000000, 0.000000]
outside [1.000000, 1.000000, 1.000000]
inside [0.000000, 0.000000, 0.000000]
outside [2.000000, 2.000000, 2.000000]
inside [0.000000, 0.000000, 0.000000]
outside [3.000000, 3.000000, 3.000000]
inside [0.000000, 0.000000, 0.000000]
outside [4.000000, 4.000000, 4.000000]
inside [0.000000, 0.000000, 0.000000]
outside [5.000000, 5.000000, 5.000000]
```

As can be seen, `inside` prints wrong results and `outside` prints correct results.

---

Another (possibly) related bug is, if I use the struct directly instead of putting it into a field, the `atomic_add` operation is not allowed in the nested for loop (but works OK if  the line `for __ in range(1):` is removed ):

```python
import taichi as ti
ti.init()
vec3 = ti.math.vec3

@ti.dataclass
class Particle:
    r: vec3

particle = Particle()

@ti.kernel
def simulate():
    for i in range(1):
        for _ in range(5):
            for __ in range(1):
                print('inside', particle.r)
            particle.r += vec3(1)
            print('outside', particle[i].r)

simulate()
```

The error message is

```

[Taichi] Starting on arch=x64
Traceback (most recent call last):
  File "<string>", line 19, in __PYTHON_EL_eval
  File "/home/zhao/workspace/test.py", line 20, in <module>
    simulate()
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/kernel_impl.py", line 947, in wrapped
    raise type(e)('\n' + str(e)) from None
taichi.lang.exception.TaichiCompilationError: 
File "/home/zhao/workspace/test.py", line 17, in simulate:
            particle.r += vec3(1)
            ^^^^^^^^^^^^^^^^^^^^^
Traceback (most recent call last):
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/ast/ast_transformer_utils.py", line 26, in __call__
    return method(ctx, node)
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/ast/ast_transformer.py", line 911, in build_AugAssign
    node.ptr = node.target.ptr._augassign(node.value.ptr,
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/common_ops.py", line 272, in _augassign
    self += x
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/common_ops.py", line 184, in __iadd__
    self._atomic_add(other)
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/common_ops.py", line 138, in _atomic_add
    return ops.atomic_add(self, other)
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/ops.py", line 113, in wrapped
    return a._element_wise_writeback_binary(imp_foo, b)
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/matrix.py", line 554, in _element_wise_writeback_binary
    entries = [[foo(self(i, j), other(i, j)) for j in range(self.m)]
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/matrix.py", line 554, in <listcomp>
    entries = [[foo(self(i, j), other(i, j)) for j in range(self.m)]
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/matrix.py", line 554, in <listcomp>
    entries = [[foo(self(i, j), other(i, j)) for j in range(self.m)]
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/ops.py", line 106, in imp_foo
    return foo(x, wrap_if_not_expr(y))
  File "/home/zhao/.local/lib/python3.10/site-packages/taichi/lang/ops.py", line 1243, in atomic_add
    expr.Expr(_ti_core.expr_atomic_add(x.ptr, y.ptr), tb=stack_info()))
AttributeError: 'float' object has no attribute 'ptr'
```
Thanks for reporting this! The bug seems to get introduced from the `CacheLoopInvariantGlobalVars` pass. cc @lin-hitonami 

```
[I 02/02/23 14:04:33.173 3850646] [compile_to_offloads.cpp:operator()@22] [k_c74_0] Atomics demoted I:
kernel {
$0 = offloaded range_for(0, 1) grid_dim=0 block_dim=32 
body {
  <f32> $1 = const 1.0
  <i32> $2 = const 1
  <i32> $3 = const 0
  <i32> $4 = loop $0 index 0
  <*f32> $5 = global ptr [S2place<f32>], index [$4] activate=true
  <i32> $6 = const 2
  $7 : for in range($3, $6) block_dim=adaptive {
    <f32> $8 = global load $5
    print "outside ", $8, "\n"
    $10 : for in range($3, $2) block_dim=adaptive {
      <f32> $11 = global load $5
      print "inside ", $11, "\n"
    }
    <f32> $13 = add $8 $1
    $14 : global store [$5 <- $13]
  }
}
}
[I 02/02/23 14:04:33.173 3850646] [compile_to_offloads.cpp:operator()@22] [k_c74_0] Cache loop-invariant global vars:
kernel {
$0 = offloaded range_for(0, 1) grid_dim=0 block_dim=32 
body {
  <f32> $1 = const 1.0
  <i32> $2 = const 1
  <i32> $3 = const 0
  <i32> $4 = loop $0 index 0
  <*f32> $5 = global ptr [S2place<f32>], index [$4] activate=true
  <i32> $6 = const 2
  <f32> $7 = alloca
  $8 = global load $5
  $9 : local store [$7 <- $8]
  <f32> $10 = alloca
  $11 = global load $5
  $12 : local store [$10 <- $11]
  $13 : for in range($3, $6) block_dim=adaptive {
    $14 = local load [$7]
    print "outside ", $14, "\n"
    <f32> $16 = alloca
    $17 = local load [$10]
    $18 : local store [$16 <- $17]
    $19 : for in range($3, $2) block_dim=adaptive {
      $20 = local load [$16]
      print "inside ", $20, "\n"
    }
    <f32> $22 = add $14 $1
    $23 : local store [$7 <- $22]
  }
  $24 = local load [$7]
  $25 : global store [$5 <- $24]
}
}
```

Fix patch:
diff --git a/taichi/transforms/cache_loop_invariant_global_vars.cpp b/taichi/transforms/cache_loop_invariant_global_vars.cpp
index 2db7e69bf13e5..3966bd4d9d181 100644
--- a/taichi/transforms/cache_loop_invariant_global_vars.cpp
+++ b/taichi/transforms/cache_loop_invariant_global_vars.cpp
@@ -16,7 +16,7 @@ class CacheLoopInvariantGlobalVars : public LoopInvariantDetector {
 
   typedef std::unordered_map<Stmt *, std::pair<CacheStatus, AllocaStmt *>>
       CacheMap;
-  std::stack<CacheMap> cached_maps;
+  std::vector<CacheMap> cached_maps;
 
   DelayedIRModifier modifier;
   std::unordered_map<const SNode *, GlobalPtrStmt *> loop_unique_ptr_;
@@ -87,33 +87,33 @@ class CacheLoopInvariantGlobalVars : public LoopInvariantDetector {
   }
 
   void visit_loop(Block *body) override {
-    cached_maps.emplace();
+    cached_maps.emplace_back();
     LoopInvariantDetector::visit_loop(body);
-    cached_maps.pop();
+    cached_maps.pop_back();
   }
 
-  void add_writeback(AllocaStmt *alloca_stmt, Stmt *global_var) {
+  void add_writeback(AllocaStmt *alloca_stmt, Stmt *global_var, int depth) {
     auto final_value = std::make_unique<LocalLoadStmt>(alloca_stmt);
     auto global_store =
         std::make_unique<GlobalStoreStmt>(global_var, final_value.get());
-    modifier.insert_after(current_loop_stmt(), std::move(global_store));
-    modifier.insert_after(current_loop_stmt(), std::move(final_value));
+    modifier.insert_after(get_loop_stmt(depth), std::move(global_store));
+    modifier.insert_after(get_loop_stmt(depth), std::move(final_value));
   }
 
-  void set_init_value(AllocaStmt *alloca_stmt, Stmt *global_var) {
+  void set_init_value(AllocaStmt *alloca_stmt, Stmt *global_var, int depth) {
     auto new_global_load = std::make_unique<GlobalLoadStmt>(global_var);
     auto local_store =
         std::make_unique<LocalStoreStmt>(alloca_stmt, new_global_load.get());
-    modifier.insert_before(current_loop_stmt(), std::move(new_global_load));
-    modifier.insert_before(current_loop_stmt(), std::move(local_store));
+    modifier.insert_before(get_loop_stmt(depth), std::move(new_global_load));
+    modifier.insert_before(get_loop_stmt(depth), std::move(local_store));
   }
 
-  AllocaStmt *cache_global_to_local(Stmt *dest, CacheStatus status) {
-    if (auto &[cached_status, alloca_stmt] = cached_maps.top()[dest];
+  AllocaStmt *cache_global_to_local(Stmt *dest, CacheStatus status, int depth) {
+    if (auto &[cached_status, alloca_stmt] = cached_maps[depth][dest];
         cached_status != CacheStatus::None) {
       // The global variable has already been cached.
       if (cached_status == CacheStatus::Read && status == CacheStatus::Write) {
-        add_writeback(alloca_stmt, dest);
+        add_writeback(alloca_stmt, dest, depth);
         cached_status = CacheStatus::ReadWrite;
       }
       return alloca_stmt;
@@ -121,19 +121,35 @@ class CacheLoopInvariantGlobalVars : public LoopInvariantDetector {
     auto alloca_unique =
         std::make_unique<AllocaStmt>(dest->ret_type.ptr_removed());
     auto alloca_stmt = alloca_unique.get();
-    modifier.insert_before(current_loop_stmt(), std::move(alloca_unique));
-    set_init_value(alloca_stmt, dest);
+    modifier.insert_before(get_loop_stmt(depth), std::move(alloca_unique));
+    set_init_value(alloca_stmt, dest, depth);
     if (status == CacheStatus::Write) {
-      add_writeback(alloca_stmt, dest);
+      add_writeback(alloca_stmt, dest, depth);
     }
-    cached_maps.top()[dest] = {status, alloca_stmt};
+    cached_maps[depth][dest] = {status, alloca_stmt};
     return alloca_stmt;
   }
 
+  std::optional<int> find_cache_depth_if_cacheable(Stmt *operand,
+                                                   Block *current_scope) {
+    if (!is_offload_unique(operand)) {
+      return std::nullopt;
+    }
+    std::optional<int> depth;
+    for (int n = loop_blocks.size() - 1; n > 0; n--) {
+      if (is_operand_loop_invariant(operand, current_scope, n)) {
+        depth = n;
+      } else {
+        break;
+      }
+    }
+    return depth;
+  }
+
   void visit(GlobalLoadStmt *stmt) override {
-    if (is_offload_unique(stmt->src) &&
-        is_operand_loop_invariant(stmt->src, stmt->parent)) {
-      auto alloca_stmt = cache_global_to_local(stmt->src, CacheStatus::Read);
+    if (auto depth = find_cache_depth_if_cacheable(stmt->src, stmt->parent)) {
+      auto alloca_stmt =
+          cache_global_to_local(stmt->src, CacheStatus::Read, depth.value());
       auto local_load = std::make_unique<LocalLoadStmt>(alloca_stmt);
       stmt->replace_usages_with(local_load.get());
       modifier.insert_before(stmt, std::move(local_load));
@@ -142,9 +158,9 @@ class CacheLoopInvariantGlobalVars : public LoopInvariantDetector {
   }
 
   void visit(GlobalStoreStmt *stmt) override {
-    if (is_offload_unique(stmt->dest) &&
-        is_operand_loop_invariant(stmt->dest, stmt->parent)) {
-      auto alloca_stmt = cache_global_to_local(stmt->dest, CacheStatus::Write);
+    if (auto depth = find_cache_depth_if_cacheable(stmt->dest, stmt->parent)) {
+      auto alloca_stmt =
+          cache_global_to_local(stmt->dest, CacheStatus::Write, depth.value());
       auto local_store =
           std::make_unique<LocalStoreStmt>(alloca_stmt, stmt->val);
       stmt->replace_usages_with(local_store.get());
diff --git a/taichi/transforms/loop_invariant_detector.h b/taichi/transforms/loop_invariant_detector.h
index 3913b7bd46733..65b83aa6d669e 100644
--- a/taichi/transforms/loop_invariant_detector.h
+++ b/taichi/transforms/loop_invariant_detector.h
@@ -12,7 +12,7 @@ class LoopInvariantDetector : public BasicStmtVisitor {
  public:
   using BasicStmtVisitor::visit;
 
-  std::stack<Block *> loop_blocks;
+  std::vector<Block *> loop_blocks;
 
   const CompileConfig &config;
 
@@ -20,13 +20,18 @@ class LoopInvariantDetector : public BasicStmtVisitor {
     allow_undefined_visitor = true;
   }
 
-  bool is_operand_loop_invariant_impl(Stmt *operand, Block *current_scope) {
+  bool is_operand_loop_invariant_impl(Stmt *operand,
+                                      Block *current_scope,
+                                      Block *loop_block = nullptr) {
+    if (!loop_block) {
+      loop_block = loop_blocks.back();
+    }
     if (operand->parent == current_scope) {
       // This statement has an operand that is in the current scope,
       // so it can not be moved out of the scope.
       return false;
     }
-    if (current_scope != loop_blocks.top()) {
+    if (current_scope != loop_block) {
       // If we enable moving code from a nested if block, we need to check
       // visibility. Example:
       // for i in range(10):
@@ -43,7 +48,7 @@ class LoopInvariantDetector : public BasicStmtVisitor {
         // If the one of the current_scope of the operand is the top loop
         // scope Then it will not be visible if we move it outside the top
         // loop scope
-        if (operand_parent == loop_blocks.top()->parent_stmt) {
+        if (operand_parent == loop_block->parent_stmt) {
           return false;
         }
       }
@@ -51,15 +56,21 @@ class LoopInvariantDetector : public BasicStmtVisitor {
     return true;
   }
 
-  bool is_operand_loop_invariant(Stmt *operand, Block *current_scope) {
-    if (loop_blocks.size() <= 1)
+  bool is_operand_loop_invariant(Stmt *operand,
+                                 Block *current_scope,
+                                 int depth = -1) {
+    if (depth == -1) {
+      depth = loop_blocks.size() - 1;
+    }
+    if (depth <= 0)
       return false;
-    return is_operand_loop_invariant_impl(operand, current_scope);
+    return is_operand_loop_invariant_impl(operand, current_scope,
+                                          loop_blocks[depth]);
   }
 
   bool is_loop_invariant(Stmt *stmt, Block *current_scope) {
     if (loop_blocks.size() <= 1 || (!config.move_loop_invariant_outside_if &&
-                                    current_scope != loop_blocks.top()))
+                                    current_scope != loop_blocks.back()))
       return false;
 
     bool is_invariant = true;
@@ -73,21 +84,24 @@ class LoopInvariantDetector : public BasicStmtVisitor {
     return is_invariant;
   }
 
-  Stmt *current_loop_stmt() {
-    return loop_blocks.top()->parent_stmt;
+  Stmt *get_loop_stmt(int depth) {
+    return loop_blocks[depth]->parent_stmt;
   }
 
+  Stmt *current_loop_stmt() {
+    return get_loop_stmt(loop_blocks.size() - 1);
+  }
   void visit(Block *stmt_list) override {
     for (auto &stmt : stmt_list->statements)
       stmt->accept(this);
   }
 
   virtual void visit_loop(Block *body) {
-    loop_blocks.push(body);
+    loop_blocks.push_back(body);
 
     body->accept(this);
 
-    loop_blocks.pop();
+    loop_blocks.pop_back();
   }
 
   void visit(RangeForStmt *stmt) override {


Test patch:
diff --git a/tests/python/test_loops.py b/tests/python/test_loops.py
index 81b0456df8cd3..dd58d698e8a08 100644
--- a/tests/python/test_loops.py
+++ b/tests/python/test_loops.py
@@ -187,3 +187,19 @@ def foo() -> ti.i32:
         return a
 
     assert foo() == 100
+
+
+@test_utils.test()
+def test_cache_loop_invariant_global_var_in_nested_loops():
+    p = ti.field(float, 1)
+
+    @ti.kernel
+    def k():
+        for n in range(1):
+            for t in range(2):
+                for m in range(1):
+                    p[n] = p[n] + 1.0
+                p[n] = p[n] + 1.0
+
+    k()
+    assert p[0] == 4.0