mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
jpwang
e41fadbe3d
* [BugFix] Fix warp_reduce truncating int64/uint64 to 32 bits on sm_80+ * Add tests to resolve review comments
130 lines
3.7 KiB
Python
130 lines
3.7 KiB
Python
import torch
|
|
import pytest
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit
|
|
def get_kernel(reduce_op: str, dtype: str):
|
|
assert reduce_op in ["sum", "max", "min", "bitand", "bitor"]
|
|
|
|
@T.prim_func
|
|
def main(x: T.Tensor((32), dtype)):
|
|
with T.Kernel(1, threads=32):
|
|
tx = T.get_thread_binding(0)
|
|
local_val = T.alloc_local([1], dtype)
|
|
local_val[0] = x[tx]
|
|
reduced_val = T.alloc_local([1], dtype)
|
|
if reduce_op == "sum":
|
|
reduced_val[0] = T.warp_reduce_sum(local_val[0])
|
|
elif reduce_op == "max":
|
|
reduced_val[0] = T.warp_reduce_max(local_val[0])
|
|
elif reduce_op == "min":
|
|
reduced_val[0] = T.warp_reduce_min(local_val[0])
|
|
elif reduce_op == "bitand":
|
|
reduced_val[0] = T.warp_reduce_bitand(local_val[0])
|
|
elif reduce_op == "bitor":
|
|
reduced_val[0] = T.warp_reduce_bitor(local_val[0])
|
|
x[tx] = reduced_val[0]
|
|
|
|
return main
|
|
|
|
|
|
def test_warp_reduce_sum():
|
|
a = torch.randn((32,), dtype=torch.float32, device="cuda")
|
|
kernel = get_kernel("sum", T.float32)
|
|
ref = torch.full_like(a, a.sum())
|
|
kernel(a)
|
|
torch.testing.assert_close(a, ref)
|
|
|
|
|
|
def test_warp_reduce_max():
|
|
a = torch.randn((32,), dtype=torch.float32, device="cuda")
|
|
kernel = get_kernel("max", T.float32)
|
|
print(kernel.get_kernel_source())
|
|
ref = torch.full_like(a, a.max())
|
|
kernel(a)
|
|
torch.testing.assert_close(a, ref)
|
|
|
|
|
|
def test_warp_reduce_min():
|
|
a = torch.randn((32,), dtype=torch.float32, device="cuda")
|
|
kernel = get_kernel("min", T.float32)
|
|
ref = torch.full_like(a, a.min())
|
|
kernel(a)
|
|
torch.testing.assert_close(a, ref)
|
|
|
|
|
|
def test_warp_reduce_bitand():
|
|
a = torch.randint(0, 100, size=(32,), dtype=torch.int32, device="cuda")
|
|
kernel = get_kernel("bitand", T.int32)
|
|
ref_val = a[0]
|
|
for i in range(1, a.shape[0]):
|
|
ref_val = ref_val & a[i]
|
|
ref = torch.full_like(a, ref_val)
|
|
kernel(a)
|
|
torch.testing.assert_close(a, ref)
|
|
|
|
|
|
def test_warp_reduce_bitor():
|
|
a = torch.randint(0, 100, size=(32,), dtype=torch.int32, device="cuda")
|
|
kernel = get_kernel("bitor", T.int32)
|
|
ref_val = a[0]
|
|
for i in range(1, a.shape[0]):
|
|
ref_val = ref_val | a[i]
|
|
ref = torch.full_like(a, ref_val)
|
|
kernel(a)
|
|
torch.testing.assert_close(a, ref)
|
|
|
|
|
|
WARP_REDUCE_CASES_64 = [
|
|
# (op, dtype, N)
|
|
("sum", "int64", 32),
|
|
("max", "int64", 32),
|
|
("min", "int64", 32),
|
|
("bitand", "int64", 32),
|
|
("bitor", "int64", 32),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("op", "dtype", "N"),
|
|
WARP_REDUCE_CASES_64,
|
|
)
|
|
def test_warp_reduce_64(op, dtype, N):
|
|
def warp_reduce_ref(a):
|
|
if op == "sum":
|
|
return torch.full_like(a, a.sum())
|
|
elif op == "max":
|
|
return torch.full_like(a, a.max())
|
|
elif op == "min":
|
|
return torch.full_like(a, a.min())
|
|
elif op == "bitand":
|
|
ref_val = a[0]
|
|
for i in range(1, a.shape[0]):
|
|
ref_val = ref_val & a[i]
|
|
return torch.full_like(a, ref_val)
|
|
elif op == "bitor":
|
|
ref_val = a[0]
|
|
for i in range(1, a.shape[0]):
|
|
ref_val = ref_val | a[i]
|
|
return torch.full_like(a, ref_val)
|
|
raise AssertionError(f"Unknown op: {op}")
|
|
|
|
torch_dtype = getattr(torch, dtype)
|
|
tl_dtype = getattr(T, dtype)
|
|
|
|
a = torch.randint(1 << 32, (1 << 63) - 1, (N,), dtype=torch_dtype, device="cuda")
|
|
ref = warp_reduce_ref(a)
|
|
|
|
kernel = get_kernel(op, tl_dtype)
|
|
kernel(a)
|
|
|
|
torch.testing.assert_close(a, ref)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|