mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
KurisandClaude Opus 4.6
dc60ab8f04
* [Bugfix] Fix alloc_var re-bind warning when assigned with comparison ops EqualOp/NotEqualOp (returned by BufferLoad == / !=) are not PrimExpr subclasses, so `bind()` missed the alloc_var store path and fell through to `bind_immutable`, triggering a spurious "Immutable value is re-bound" warning. Convert them to PrimExpr via `.asobject()` before the store check. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Fix pre-commit lint: double quotes, remove redundant parens Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Fix ruff format Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
65 lines
1.7 KiB
Python
65 lines
1.7 KiB
Python
import warnings
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
import torch
|
|
|
|
|
|
@tilelang.jit
|
|
def kernel_alloc_var_eq(data):
|
|
N = T.dynamic("N")
|
|
data: T.Tensor[[N], T.int32]
|
|
out = T.empty([N], T.int32)
|
|
|
|
with T.Kernel(1) as _:
|
|
for i in T.serial(N):
|
|
cond = T.alloc_var(T.bool)
|
|
cond = data[i] == 2
|
|
if cond:
|
|
out[i] = 1
|
|
else:
|
|
out[i] = 0
|
|
return out
|
|
|
|
|
|
@tilelang.jit
|
|
def kernel_alloc_var_ne(data):
|
|
N = T.dynamic("N")
|
|
data: T.Tensor[[N], T.int32]
|
|
out = T.empty([N], T.int32)
|
|
|
|
with T.Kernel(1) as _:
|
|
for i in T.serial(N):
|
|
cond = T.alloc_var(T.bool)
|
|
cond = data[i] != 0
|
|
if cond:
|
|
out[i] = 1
|
|
else:
|
|
out[i] = 0
|
|
return out
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_alloc_var_eq():
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("error", message="Immutable value.*is re-bound")
|
|
data = torch.tensor([0, 1, 2, 3, 2], dtype=torch.int32, device="cuda")
|
|
result = kernel_alloc_var_eq(data)
|
|
expected = torch.tensor([0, 0, 1, 0, 1], dtype=torch.int32, device="cuda")
|
|
torch.testing.assert_close(result, expected)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_alloc_var_ne():
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("error", message="Immutable value.*is re-bound")
|
|
data = torch.tensor([0, 1, 2, 3, 0], dtype=torch.int32, device="cuda")
|
|
result = kernel_alloc_var_ne(data)
|
|
expected = torch.tensor([0, 1, 1, 1, 0], dtype=torch.int32, device="cuda")
|
|
torch.testing.assert_close(result, expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|