[SIMT] Add ballot_sync warp intrinsics (#4641)

taichi-dev · Mar 29, 2022 · 22a099d · 22a099d
1 parent c3f7431
commit 22a099d
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 5 deletions.
diff --git a/python/taichi/lang/simt/warp.py b/python/taichi/lang/simt/warp.py
@@ -17,9 +17,11 @@ def unique():
     pass
 
 
-def ballot():
-    # TODO
-    pass
+def ballot(predicate):
+    return expr.Expr(
+        _ti_core.insert_internal_func_call("cuda_ballot_i32",
+                                           expr.make_expr_group(predicate),
+                                           False))
 
 
 def shfl_i32(mask, val, offset):

diff --git a/taichi/runtime/llvm/runtime.cpp b/taichi/runtime/llvm/runtime.cpp
@@ -1040,6 +1040,14 @@ int32 cuda_ballot_sync(int32 mask, bool bit) {
   return 0;
 }
 
+int32 cuda_ballot_i32(int32 predicate) {
+  return cuda_ballot_sync(UINT32_MAX, (bool)predicate);
+}
+
+int32 cuda_ballot_sync_i32(u32 mask, int32 predicate) {
+  return cuda_ballot_sync(mask, (bool)predicate);
+}
+
 i32 cuda_match_any_sync_i32(i32 mask, i32 value) {
   return 0;
 }

diff --git a/tests/python/test_simt.py b/tests/python/test_simt.py
@@ -1,3 +1,5 @@
+import random
+
 from pytest import approx
 
 import taichi as ti
@@ -24,8 +26,24 @@ def test_unique():
 
 @test_utils.test(arch=ti.cuda)
 def test_ballot():
-    # TODO
-    pass
+    a = ti.field(dtype=ti.u32, shape=32)
+    b = ti.field(dtype=ti.u32, shape=32)
+
+    @ti.kernel
+    def foo():
+        ti.loop_config(block_dim=32)
+        for i in range(32):
+            a[i] = ti.simt.warp.ballot(b[i])
+
+    key = 0
+    for i in range(32):
+        b[i] = i % 2
+        key += b[i] * pow(2, i)
+
+    foo()
+
+    for i in range(32):
+        assert a[i] == key
 
 
 @test_utils.test(arch=ti.cuda)