mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
Yongqi ZhuoandClaude Fable 5
2169b925ae
* [BugFix] Always vectorize `T.Parallel` loops
* Fix CI failures
* Restore always-vectorize accidentally reverted by the CI-fix commit
f447919b captured a transient working-tree state from a concurrent
baseline-vs-patched benchmark flip, reverting the gate removal in src/
while keeping the new-policy tests: CI then ran the old gated compiler
against them (scalar var-loop -> no float2; positional if/for check
broken by the copy loops the larger tile introduces).
Bring back the parameter-less LowerParallelLoop and rewrite the issue
1734 assertion as containment: the guarded loop must sit inside the
hoisted if block, which holds on both sm90 and sm100 regardless of how
many copy loops precede it.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
95 lines
2.7 KiB
Python
95 lines
2.7 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=T.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=T.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}")
|
|
|
|
|
|
PARALLEL_DYNAMIC_VALID_LENGTHS = [0, 13, 200, 600]
|
|
|
|
|
|
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)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"valid_len",
|
|
PARALLEL_DYNAMIC_VALID_LENGTHS,
|
|
ids=[f"valid_len={value}" for value in PARALLEL_DYNAMIC_VALID_LENGTHS],
|
|
)
|
|
def test_parallel_dynamic_extent(valid_len):
|
|
kernel = parallel_elementwise_dynamic(max_len=512, threads=256)
|
|
data = _require_cuda_tensor((512,), torch.float32)
|
|
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)
|
|
|
|
|
|
@tilelang.jit
|
|
def _parallel_vectorize_local_and_var():
|
|
with T.Kernel(1) as _:
|
|
x = T.alloc_fragment([256], T.float32)
|
|
y = T.alloc_fragment([256], T.float32)
|
|
z = T.alloc_var(T.float32)
|
|
for i in T.parallel(256):
|
|
y[i] = x[i] * z
|
|
|
|
|
|
def test_parallel_vectorize_var():
|
|
source = _parallel_vectorize_local_and_var.get_kernel_source()
|
|
# register-register loops vectorize too; the var operand becomes a
|
|
# scalar broadcast inside the packed multiply
|
|
assert "float2" in source
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|