[async] Add allocator async state

taichi-dev · Oct 18, 2020 · 8a54b46 · 8a54b46
1 parent 0298bca
commit 8a54b46
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 5 deletions.
diff --git a/taichi/program/async_utils.cpp b/taichi/program/async_utils.cpp
@@ -119,6 +119,17 @@ TaskMeta *get_task_meta(IRBank *ir_bank, const TaskLaunchRecord &t) {
       }
     }
 
+    if (auto *snode_op = stmt->cast<SNodeOpStmt>()) {
+      if (snode_op->op_type == SNodeOpType::activate ||
+          snode_op->op_type == SNodeOpType::deactivate) {
+        auto *sn = snode_op->snode;
+        if (is_gc_able(sn->type)) {
+          meta.input_states.emplace(sn, AsyncState::Type::allocator);
+          meta.output_states.emplace(sn, AsyncState::Type::allocator);
+        }
+      }
+    }
+
     if (auto ptr = stmt->cast<GlobalPtrStmt>()) {
       if (ptr->activate) {
         for (auto &snode : ptr->snodes.data) {
@@ -127,6 +138,10 @@ TaskMeta *get_task_meta(IRBank *ir_bank, const TaskLaunchRecord &t) {
             if (!s->is_path_all_dense) {
               meta.input_states.emplace(s, AsyncState::Type::mask);
               meta.output_states.emplace(s, AsyncState::Type::mask);
+              if (is_gc_able(s->type)) {
+                meta.input_states.emplace(s, AsyncState::Type::allocator);
+                meta.output_states.emplace(s, AsyncState::Type::allocator);
+              }
             }
             s = s->parent;
           }
@@ -170,6 +185,11 @@ TaskMeta *get_task_meta(IRBank *ir_bank, const TaskLaunchRecord &t) {
   } else if (root_stmt->task_type == OffloadedTaskType::struct_for) {
     meta.snode = root_stmt->snode;
     meta.input_states.emplace(root_stmt->snode, AsyncState::Type::list);
+  } else if ((root_stmt->task_type == OffloadedTaskType::gc) &&
+             (is_gc_able(root_stmt->snode->type))) {
+    meta.snode = root_stmt->snode;
+    meta.input_states.emplace(meta.snode, AsyncState::Type::allocator);
+    meta.output_states.emplace(meta.snode, AsyncState::Type::allocator);
   }
 
   meta_bank[t.ir_handle] = meta;

diff --git a/taichi/program/async_utils.h b/taichi/program/async_utils.h
@@ -78,7 +78,7 @@ class TaskLaunchRecord {
 };
 
 struct AsyncState {
-  enum class Type { mask, value, list };
+  enum class Type { mask, value, list, allocator };
 
   AsyncState(SNode *snode, Type type) : snode(snode), type(type) {
   }
@@ -106,6 +106,9 @@ struct AsyncState {
       case Type::list:
         type_name = "list";
         break;
+      case Type::allocator:
+        type_name = "allocator";
+        break;
     }
     return snode->get_node_type_name_hinted() + "_" + type_name;
   }

diff --git a/tests/python/test_gc.py b/tests/python/test_gc.py
@@ -2,7 +2,7 @@
 
 
 @ti.test(require=ti.extension.sparse)
-def test_block():
+def _test_block_gc():
     N = 100000
 
     dx = 1 / 128
@@ -38,14 +38,29 @@ def move():
             x[p] += ti.Vector([0.0, 0.1])
 
     assert grid.num_dynamically_allocated == 0
-    for i in range(100):
+    for _ in range(100):
         grid.deactivate_all()
         # Scatter the particles to the sparse grid
         build_grid()
         # Move the block of particles
         move()
-        # The block of particles can occupy at most two blocks on the sparse grid
-        assert 1 <= grid.num_dynamically_allocated <= 2
+
+    ti.sync()
+    # The block of particles can occupy at most two blocks on the sparse grid.
+    # It's fine to run 100 times and do just one final check, because
+    # num_dynamically_allocated stores the number of slots *ever* allocated.
+    assert 1 <= grid.num_dynamically_allocated <= 2, grid.num_dynamically_allocated
+
+
+@ti.test(require=ti.extension.sparse)
+def test_block():
+    _test_block_gc()
+
+
+@ti.test(require=[ti.extension.sparse, ti.extension.async_mode],
+         async_mode=True)
+def test_block_async():
+    _test_block_gc()
 
 
 @ti.test(require=ti.extension.sparse)