mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
102 lines
2.9 KiB
Python
102 lines
2.9 KiB
Python
import torch
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit
|
|
def get_select_kernel_1():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor[(128, 8), T.float32],
|
|
B: T.Tensor[(128, 8), T.float32],
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
tx = T.get_thread_binding(0)
|
|
tmp = T.alloc_var(T.bfloat16)
|
|
tmp = A[tx, 3]
|
|
|
|
B[tx, 0] = T.Select(True, A[tx, 0], 0.0)
|
|
B[tx, 1] = T.Select(False, 1.0, A[tx, 1])
|
|
B[tx, 2] = T.Select(T.cast(A[tx, 3], T.bfloat16) == tmp, A[tx, 2], T.cast(tmp, T.float32))
|
|
B[tx, 3] = T.Select(B[tx, 0] != 0.0, T.if_then_else(B[tx, 1] != 0.0, A[tx, 3], 0.0), 0.0)
|
|
|
|
for i in T.serial(4):
|
|
B[tx, i + 4] = T.Select(
|
|
A[tx, 0] == 0, T.if_then_else(T.Select(True, False, True), 1.0, 2.0), T.Select(True, A[tx, i + 4], 3.0)
|
|
)
|
|
|
|
return main
|
|
|
|
|
|
def test_select_correctness():
|
|
A = torch.randn((128, 8), dtype=torch.float32, device="cuda")
|
|
B = torch.empty((128, 8), dtype=torch.float32, device="cuda")
|
|
kernel = get_select_kernel_1()
|
|
|
|
A = torch.clamp(A, min=1e-4)
|
|
kernel(A, B)
|
|
assert torch.allclose(A, B)
|
|
|
|
|
|
@tilelang.jit
|
|
def get_select_kernel_2():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor[(128, 8), T.float32],
|
|
B: T.Tensor[(128, 8), T.float32],
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
tx = T.get_thread_binding(0)
|
|
tmp = T.alloc_var(T.bfloat16)
|
|
tmp = A[tx, 3]
|
|
|
|
B[tx, 0] = T.Select(True, A[tx, 0], 0.0)
|
|
B[tx, 1] = T.Select(False, 1.0, A[tx, 1])
|
|
B[tx, 2] = T.Select(T.cast(A[tx, 3], T.bfloat16) == tmp, A[tx, 2], T.cast(tmp, T.float32))
|
|
B[tx, 3] = T.Select(B[tx, 0] != 0.0, T.Select(B[tx, 1] != 0.0, A[tx, 3], 0.0), 0.0)
|
|
|
|
for i in T.serial(4):
|
|
B[tx, i + 4] = T.Select(
|
|
A[tx, 0] == 0,
|
|
T.Select(T.Select(True, False, True), 1.0, 2.0),
|
|
T.Select(T.sin(A[tx, 2]) == 1.0, T.sin(A[tx, 0]), T.cos(A[tx, 1])),
|
|
)
|
|
|
|
return main
|
|
|
|
|
|
def test_select_codegen_no_if():
|
|
kernel = get_select_kernel_2()
|
|
source = kernel.get_kernel_source()
|
|
assert "if (" not in source
|
|
|
|
|
|
@tilelang.jit
|
|
def get_parallel_select_kernel():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor[(1024,), T.float32],
|
|
B: T.Tensor[(1024,), T.float32],
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
for i in T.Parallel(1024):
|
|
B[i] = T.Select(i < 512, A[i], 0.0)
|
|
|
|
return main
|
|
|
|
|
|
def test_parallel_select_vectorized_condition():
|
|
kernel = get_parallel_select_kernel()
|
|
A = torch.randn((1024,), dtype=torch.float32, device="cuda")
|
|
B = torch.empty_like(A)
|
|
|
|
kernel(A, B)
|
|
|
|
expected = torch.zeros_like(A)
|
|
expected[:512] = A[:512]
|
|
torch.testing.assert_close(B, expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|