mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
Lei Wang
b3d5981b00
Reorder the `mask` parameter to the end with a default of `_FULL_WARP_MASK` for shfl_xor/down/up/sync, any_sync, all_sync, ballot_sync, and match_any/all_sync. Full-warp is the overwhelmingly common case, so making mask an optional trailing kwarg removes boilerplate at call sites and lets old `(value, delta)`-style callers keep working. Update the test suite and docs/instructions.md to match the new signatures.
63 lines
1.4 KiB
Python
63 lines
1.4 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
import torch
|
|
from tvm import tir
|
|
import tilelang.testing
|
|
|
|
|
|
@tilelang.jit
|
|
def kernel_with_warp_sync():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((1,), "int32"),
|
|
B: T.Tensor((1,), "int32"),
|
|
):
|
|
with T.Kernel(1, threads=32):
|
|
tx = T.get_thread_binding()
|
|
if tx == 0:
|
|
tir.call_extern("void", "__nanosleep", 100)
|
|
A[0] = -1
|
|
T.sync_warp()
|
|
if tx == 1:
|
|
B[0] = A[0]
|
|
|
|
return main
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_warp_sync():
|
|
a = torch.empty((1), device="cuda", dtype=torch.int32)
|
|
b = torch.empty((1), device="cuda", dtype=torch.int32)
|
|
kernel = kernel_with_warp_sync()
|
|
assert "__syncwarp" in kernel.get_kernel_source()
|
|
kernel(a, b)
|
|
assert b[0] == -1
|
|
|
|
|
|
@tilelang.jit
|
|
def kernel_with_shfl_sync():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((32,), "int32"),
|
|
):
|
|
with T.Kernel(1, threads=32):
|
|
tx = T.get_thread_binding()
|
|
val = tx * 10
|
|
broadcast = T.shfl_sync(val, 31)
|
|
A[tx] = broadcast
|
|
|
|
return main
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_shfl_sync():
|
|
a = torch.empty((32), device="cuda", dtype=torch.int32)
|
|
kernel = kernel_with_shfl_sync()
|
|
assert "__shfl_sync" in kernel.get_kernel_source()
|
|
kernel(a)
|
|
assert torch.all(a == 310)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|