try lock dynamic

zxzhaixiang · Dec 23, 2019 · 8362d6d · 8362d6d
1 parent 26385b1
commit 8362d6d
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 21 deletions.
diff --git a/taichi/backends/struct_llvm.cpp b/taichi/backends/struct_llvm.cpp
@@ -70,7 +70,7 @@ void StructCompilerLLVM::generate_types(SNode &snode) {
   } else if (type == SNodeType::dynamic) {
     body_type = llvm::PointerType::getInt8PtrTy(*ctx);
     // TODO: maybe load a struct from runtime?
-    // n (number of elements), and the mutex
+    // mutex and n (number of elements)
     aux_type =
         llvm::StructType::get(*ctx, {llvm::PointerType::getInt32Ty(*ctx),
                                      llvm::PointerType::getInt32Ty(*ctx)});

diff --git a/taichi/runtime/node_dynamic.h b/taichi/runtime/node_dynamic.h
@@ -1,9 +1,9 @@
 #pragma once
 
 struct DynamicNode {
+  Ptr ptr;
   i32 lock;
   i32 n;
-  Ptr ptr;
 };
 
 // Specialized Attributes and functions
@@ -14,27 +14,28 @@ struct DynamicMeta : public StructMeta {
 STRUCT_FIELD(DynamicMeta, chunk_size);
 
 void Dynamic_activate(Ptr meta_, Ptr node_, int i) {
-  // TODO: lock
   auto meta = (DynamicMeta *)(meta_);
   auto node = (DynamicNode *)(node_);
-  if (i < node->n)
-    return;
-  node->n = i + 1;
-  int chunk_start = 0;
-  auto p_chunk_ptr = &node->ptr;
-  auto chunk_size = meta->chunk_size;
-  auto rt = (Runtime *)meta->context->runtime;
-  auto alloc = rt->node_allocators[meta->snode_id];
-  while (true) {
-    if (*p_chunk_ptr == nullptr) {
-      *p_chunk_ptr = NodeAllocator_allocate(alloc);
-    }
-    if (i < chunk_start + chunk_size) {
+  locked_task(Ptr(&node->lock), [&] {
+    if (i < node->n)
       return;
+    node->n = i + 1;
+    int chunk_start = 0;
+    auto p_chunk_ptr = &node->ptr;
+    auto chunk_size = meta->chunk_size;
+    auto rt = (Runtime *)meta->context->runtime;
+    auto alloc = rt->node_allocators[meta->snode_id];
+    while (true) {
+      if (*p_chunk_ptr == nullptr) {
+        *p_chunk_ptr = NodeAllocator_allocate(alloc);
+      }
+      if (i < chunk_start + chunk_size) {
+        return;
+      }
+      p_chunk_ptr = (Ptr *)*p_chunk_ptr;
+      chunk_start += chunk_size;
     }
-    p_chunk_ptr = (Ptr *)*p_chunk_ptr;
-    chunk_start += chunk_size;
-  }
+  });
 }
 
 bool Dynamic_is_active(Ptr meta_, Ptr node_, int i) {

diff --git a/tests/python/test_dynamic.py b/tests/python/test_dynamic.py
@@ -46,7 +46,6 @@ def func():
 
 @ti.all_archs
 def test_dynamic_matrix():
-  return
   if ti.cfg.arch == ti.cuda:
     return
 
@@ -59,7 +58,6 @@ def place():
 
   @ti.kernel
   def func():
-    ti.serialize()
     for i in range(n // 4):
       x[i * 4][1, 0] = i