mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Lei Wang
27701c3de2
* Allow dynamic extents in loop partition; warn when layout inversion falls back to NoCheck * add test and introduce predicate * test fix * fix * enhance * inverse with level * test fix * bug fix
73 lines
2.0 KiB
Python
73 lines
2.0 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
import torch
|
|
import tilelang.testing
|
|
import pytest
|
|
|
|
tilelang.testing.set_random_seed()
|
|
|
|
|
|
@tilelang.jit(out_idx=[1])
|
|
def parallel_elementwise_static(length=256, dtype="float32"):
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((length,), dtype),
|
|
B: T.Tensor((length,), dtype),
|
|
):
|
|
with T.Kernel(1, threads=length) as _:
|
|
for i in T.Parallel(length):
|
|
B[i] = A[i] + 1.0
|
|
|
|
return main
|
|
|
|
|
|
@tilelang.jit(out_idx=[1])
|
|
def parallel_elementwise_dynamic(max_len=512, threads=256, dtype="float32"):
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((max_len,), dtype),
|
|
B: T.Tensor((max_len,), dtype),
|
|
valid_len: T.int32,
|
|
):
|
|
with T.Kernel(1, threads=threads) as _:
|
|
for i in T.Parallel(max_len):
|
|
B[i] = 0.0
|
|
span = T.min(valid_len, max_len)
|
|
for i in T.Parallel(span):
|
|
B[i] = A[i] - 1.0
|
|
|
|
return main
|
|
|
|
|
|
def _require_cuda_tensor(shape, dtype=torch.float32):
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA not available")
|
|
try:
|
|
return torch.randn(*shape, device="cuda", dtype=dtype)
|
|
except RuntimeError as err:
|
|
pytest.skip(f"CUDA runtime unavailable: {err}")
|
|
|
|
|
|
def test_parallel_static_extent():
|
|
kernel = parallel_elementwise_static(length=256)
|
|
data = _require_cuda_tensor((256,), torch.float32)
|
|
result = kernel(data)
|
|
torch.testing.assert_close(result, data + 1.0, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
def test_parallel_dynamic_extent():
|
|
kernel = parallel_elementwise_dynamic(max_len=512, threads=256)
|
|
data = _require_cuda_tensor((512,), torch.float32)
|
|
for valid_len in [0, 13, 200, 600]:
|
|
out = kernel(data, valid_len)
|
|
reference = torch.zeros_like(data)
|
|
clip = min(valid_len, data.shape[0])
|
|
reference[:clip] = data[:clip] - 1.0
|
|
torch.testing.assert_close(out, reference, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|