mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
from tilelang import language as T
|
|
|
|
|
|
def _compile_cuda(func):
|
|
tilelang.disable_cache()
|
|
try:
|
|
return tilelang.compile(func, target="cuda", execution_backend="tvm_ffi")
|
|
finally:
|
|
tilelang.enable_cache()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_sync_threads_with_variable_barrier_id():
|
|
@T.prim_func
|
|
def kernel():
|
|
with T.Kernel(1, threads=256) as (bx,):
|
|
barrier_id = T.int32(1)
|
|
T.sync_threads(barrier_id)
|
|
T.sync_threads(barrier_id, 128)
|
|
|
|
kernel = _compile_cuda(kernel)
|
|
src = kernel.get_kernel_source()
|
|
assert "tl::__sync_thread_partial(" in src, src
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_sync_threads_with_computed_barrier_id():
|
|
@T.prim_func
|
|
def kernel():
|
|
with T.Kernel(1, threads=256) as (bx,):
|
|
tx = T.launch_thread("threadIdx.x", 256)
|
|
barrier_id = tx % 4
|
|
T.sync_threads(barrier_id, 256)
|
|
|
|
kernel = _compile_cuda(kernel)
|
|
src = kernel.get_kernel_source()
|
|
assert "tl::__sync_thread_partial(" in src, src
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|