mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
xiaoyouPREG
ffeda9f394
* [BugFix] Pack 32-lane 8-bit CUDA vectors correctly * [Test] Avoid SM100-only store in int8 broadcast regression * [BugFix] Move int8x32 packing into CUDA templates
62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
import pytest
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
LANES = 32
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dtype,torch_dtype,value",
|
|
[("int8", torch.int8, -37), ("uint8", torch.uint8, 213)],
|
|
)
|
|
@tilelang.testing.requires_cuda
|
|
def test_runtime_int8_broadcast_to_32_lanes(dtype, torch_dtype, value):
|
|
@T.prim_func
|
|
def main(src: T.Tensor((1,), dtype), out: T.Tensor((LANES,), dtype)):
|
|
with T.Kernel(1, threads=1) as _:
|
|
# Scalar global stores keep this test runnable on pre-SM100 GPUs.
|
|
packed = T.alloc_local((LANES,), dtype)
|
|
packed[0:LANES] = T.Broadcast(src[0], LANES)
|
|
for i in T.serial(LANES):
|
|
out[i] = packed[i]
|
|
|
|
kernel = tilelang.compile(main, target="cuda")
|
|
constructor = "make_ulonglong4" if dtype == "uint8" else "make_longlong4"
|
|
assert constructor in kernel.get_kernel_source()
|
|
src = torch.tensor([value], dtype=torch_dtype, device="cuda")
|
|
out = torch.empty(LANES, dtype=torch_dtype, device="cuda")
|
|
kernel(src, out)
|
|
|
|
torch.testing.assert_close(out, torch.full_like(out, value), rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dtype,torch_dtype,value",
|
|
[("int8", torch.int8, -37), ("uint8", torch.uint8, 213)],
|
|
)
|
|
@tilelang.testing.requires_cuda
|
|
def test_constant_int8_broadcast_to_32_lanes(dtype, torch_dtype, value):
|
|
@T.prim_func
|
|
def main(out: T.Tensor((LANES,), dtype)):
|
|
with T.Kernel(1, threads=1) as _:
|
|
packed = T.alloc_local((LANES,), dtype)
|
|
packed[0:LANES] = T.Broadcast(T.cast(value, dtype), LANES)
|
|
for i in T.serial(LANES):
|
|
out[i] = packed[i]
|
|
|
|
kernel = tilelang.compile(main, target="cuda")
|
|
constructor = "make_ulonglong4" if dtype == "uint8" else "make_longlong4"
|
|
assert constructor in kernel.get_kernel_source()
|
|
out = torch.empty(LANES, dtype=torch_dtype, device="cuda")
|
|
kernel(out)
|
|
|
|
torch.testing.assert_close(out, torch.full_like(out, value), rtol=0, atol=0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|