[Bug] [sparse] Insert gc tasks only for pointer and dynamic SNodes (

#781)

Co-authored-by: Taichi Gardener <taichigardener@gmail.com>

README.md

@@ -9,7 +9,7 @@
 python3 -m pip install taichi
 ```
 **Supported OS**: Windows, Linux, Mac OS X; **Python**: 3.6, 3.7, 3.8; **Backends**: x64 CPUs, CUDA, Apple Metal.
- 
+
 Please build from source for other configurations (e.g., you need the experimental OpenGL backend or your CPU is ARM).
 
 **Note:**

docs/contributor_guide.rst

@@ -107,6 +107,7 @@ Existing tags:
 - ``[Lang]``: frontend language features, including syntax sugars;
 - ``[Std]``: standard library, e.g. `ti.Matrix` and `ti.Vector`;
 - ``[IR]``: intermediate representation;
+- ``[Sparse]``: sparse computation, dynamic memory allocator, and garbage collection;
 - ``[Opt]``: IR optimization passes;
 - ``[Async]``: asynchronous execution engine;
 - ``[Type]``: type system;
@@ -120,6 +121,7 @@ Existing tags:
 - ``[Test]``: adding or improving tests under ``tests/``;
 - ``[PyPI]``: PyPI package release;
 - ``[Misc]``: something that doesn't belong to any category, such as version bump, reformatting;
+- ``[Bug]``: bug fixes;
 - **When introducing a new tag, please update the list here in the first PR with that tag, so that people can follow.**
 
 .. note::

taichi/codegen/codegen_llvm.cpp

@@ -943,7 +943,8 @@ void CodeGenLLVM::visit(SNodeOpStmt *stmt) {
     llvm_val[stmt] =
         call(snode, llvm_val[stmt->ptr], "is_active", {llvm_val[stmt->val]});
   } else if (stmt->op_type == SNodeOpType::deactivate) {
-    if (snode->type == SNodeType::pointer || snode->type == SNodeType::hash) {
+    if (snode->type == SNodeType::pointer || snode->type == SNodeType::hash ||
+        snode->type == SNodeType::bitmasked) {
       llvm_val[stmt] =
           call(snode, llvm_val[stmt->ptr], "deactivate", {llvm_val[stmt->val]});
     } else if (snode->type == SNodeType::dynamic) {

taichi/lang_util.cpp

@@ -130,6 +130,10 @@ std::string snode_type_name(SNodeType t) {
   return type_names[t];
 }
 
+bool is_gc_able(SNodeType t) {
+  return (t == SNodeType::pointer || t == SNodeType::dynamic);
+}
+
 std::string unary_op_type_name(UnaryOpType type) {
   static std::map<UnaryOpType, std::string> type_names;
   if (type_names.empty()) {

taichi/lang_util.h

@@ -66,6 +66,8 @@ enum class SNodeType {
 
 std::string snode_type_name(SNodeType t);
 
+bool is_gc_able(SNodeType t);
+
 enum class UnaryOpType : int {
 #define PER_UNARY_OP(x) x,
 #include "taichi/inc/unary_op.inc.h"

taichi/program/program.cpp

@@ -211,8 +211,7 @@ void Program::initialize_runtime_system(StructCompiler *scomp) {
                                   (int)snodes.size());
 
   for (int i = 0; i < (int)snodes.size(); i++) {
-    if (snodes[i]->type == SNodeType::pointer ||
-        snodes[i]->type == SNodeType::dynamic) {
+    if (is_gc_able(snodes[i]->type)) {
       std::size_t node_size;
       if (snodes[i]->type == SNodeType::pointer)
         node_size = tlctx->get_type_size(

taichi/runtime/llvm/node_bitmasked.h

@@ -20,6 +20,15 @@ void Bitmasked_activate(Ptr meta, Ptr node, int i) {
   atomic_or_u32(&mask_begin[i / 32], 1UL << (i % 32));
 }
 
+void Bitmasked_deactivate(Ptr meta, Ptr node, int i) {
+  auto smeta = (StructMeta *)meta;
+  auto element_size = StructMeta_get_element_size(smeta);
+  auto num_elements = Bitmasked_get_num_elements(meta, node);
+  auto data_section_size = element_size * num_elements;
+  auto mask_begin = (u32 *)(node + data_section_size);
+  atomic_and_u32(&mask_begin[i / 32], ~(1UL << (i % 32)));
+}
+
 i32 Bitmasked_is_active(Ptr meta, Ptr node, int i) {
   auto smeta = (StructMeta *)meta;
   auto element_size = StructMeta_get_element_size(smeta);

taichi/transforms/offload.cpp

@@ -487,10 +487,12 @@ void insert_gc(IRNode *root) {
 
   for (int i = (int)b->statements.size() - 1; i >= 0; i--) {
     auto snodes = gc_statements[i].second;
-    for (auto j = 0; j < snodes.size(); j++) {
-      b->statements.insert(
-          b->statements.begin() + i + 1,
-          Stmt::make<OffloadedStmt>(OffloadedStmt::TaskType::gc, snodes[j]));
+    for (auto *snode : snodes) {
+      if (is_gc_able(snode->type)) {
+        b->statements.insert(
+            b->statements.begin() + i + 1,
+            Stmt::make<OffloadedStmt>(OffloadedStmt::TaskType::gc, snode));
+      }
     }
   }
 }

tests/python/test_bitmasked.py

@@ -140,3 +140,34 @@ def count():
     func()
     count()
     assert c[None] == n
+
+
+@archs_support_bitmasked
+def test_deactivate():
+    # https://github.com/taichi-dev/taichi/issues/778
+    a = ti.var(ti.i32)
+    a_a = ti.root.bitmasked(ti.i, 4)
+    a_b = a_a.dense(ti.i, 4)
+    a_b.place(a)
+    c = ti.var(ti.i32)
+    ti.root.place(c)
+
+    @ti.kernel
+    def run():
+        a[0] = 123
+
+    @ti.kernel
+    def is_active():
+        c[None] = ti.is_active(a_a, [0])
+
+    @ti.kernel
+    def deactivate():
+        ti.deactivate(a_a, [0])
+
+    run()
+    is_active()
+    assert c[None] == 1
+
+    deactivate()
+    is_active()
+    assert c[None] == 0