mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-26 07:22:54 +08:00
74 lines
2.0 KiB
Python
74 lines
2.0 KiB
Python
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from tilelang.utils.language import retrieve_stride
|
|
|
|
|
|
M = 4
|
|
N = 3
|
|
PITCH = 8
|
|
|
|
|
|
def _make_store_kernel():
|
|
@T.prim_func
|
|
def main(dst: T.StridedTensor((M, N), (PITCH, 1), "int32")):
|
|
with T.Kernel(1, threads=1):
|
|
for row in T.serial(M):
|
|
T.stg32(dst[row, 0], T.Cast("uint32", row + 1))
|
|
|
|
return main
|
|
|
|
|
|
def _make_load_kernel():
|
|
@T.prim_func
|
|
def main(
|
|
src: T.StridedTensor((M, N), (PITCH, 1), "int32"),
|
|
out: T.Tensor((M,), "int32"),
|
|
):
|
|
with T.Kernel(1, threads=1):
|
|
for row in T.serial(M):
|
|
out[row] = T.reinterpret(T.ldg32(src[row, 0]), "int32")
|
|
|
|
return main
|
|
|
|
|
|
def test_retrieve_stride_preserves_scalar_rank():
|
|
scalar = T.Tensor((), "float32")
|
|
assert retrieve_stride(scalar) == []
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_strided_stg32_uses_declared_stride():
|
|
kernel = tilelang.compile(_make_store_kernel(), target="cuda")
|
|
physical = torch.zeros(M * PITCH, dtype=torch.int32, device="cuda")
|
|
view = physical.as_strided((M, N), (PITCH, 1))
|
|
|
|
kernel(view)
|
|
|
|
expected = torch.arange(1, M + 1, dtype=torch.int32, device="cuda")
|
|
torch.testing.assert_close(view[:, 0], expected, rtol=0, atol=0)
|
|
source = kernel.get_kernel_source()
|
|
assert "dst[(row * 8)]" in source
|
|
assert "dst[(row * 3)]" not in source
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_strided_ldg32_uses_declared_stride():
|
|
physical = torch.arange(M * PITCH, dtype=torch.int32, device="cuda")
|
|
view = physical.as_strided((M, N), (PITCH, 1))
|
|
out = torch.empty(M, dtype=torch.int32, device="cuda")
|
|
kernel = tilelang.compile(_make_load_kernel(), target="cuda")
|
|
|
|
kernel(view, out)
|
|
|
|
torch.testing.assert_close(out, view[:, 0], rtol=0, atol=0)
|
|
source = kernel.get_kernel_source()
|
|
assert "src[(row * 8)]" in source
|
|
assert "src[(row * 3)]" not in source
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|