Merge pull request #5 from taichi-dev/master

[SIMT] Add syncwarp warp intrinsics (taichi-dev#4917)
galeselee · May 6, 2022 · 93bbc0e · 93bbc0e
2 parents abd0136 + e3d58c6
commit 93bbc0e
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/python/taichi/lang/simt/warp.py b/python/taichi/lang/simt/warp.py
@@ -102,9 +102,10 @@ def active_mask():
     pass
 
 
-def sync():
-    # TODO
-    pass
+def sync(mask):
+    expr.Expr(
+        _ti_core.insert_internal_func_call("warp_barrier",
+                                           expr.make_expr_group(mask), False))
 
 
 __all__ = [

diff --git a/tests/python/test_simt.py b/tests/python/test_simt.py
@@ -257,8 +257,21 @@ def test_active_mask():
 
 @test_utils.test(arch=ti.cuda)
 def test_sync():
-    # TODO
-    pass
+    a = ti.field(dtype=ti.u32, shape=32)
+
+    @ti.kernel
+    def foo():
+        ti.loop_config(block_dim=32)
+        for i in range(32):
+            a[i] = i
+        ti.simt.warp.sync(ti.u32(0xFFFFFFFF))
+        for i in range(16):
+            a[i] = a[i + 16]
+
+    foo()
+
+    for i in range(32):
+        assert a[i] == i % 16 + 16
 
 
 # Higher level primitives test