mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-28 07:43:01 +08:00
92 lines
3.0 KiB
Python
92 lines
3.0 KiB
Python
"""LLVM atomic op tests (shares the CPU atomic impls with the `c` target).
|
|
|
|
Lightweight mirror of testing/python/cpu/test_tilelang_cpu_atomic.py covering
|
|
one representative case per lowering path, since `llvm` and `c` share the same
|
|
registries (TargetIsCPU matches kDLCPU) and the same LowerCPUAtomics pass:
|
|
- scalar add + return_prev (tl.transform.LowerCPUAtomics),
|
|
- tile-region max (src/cpu/op/atomic_reduce.cc),
|
|
- addx4 vector expansion (tl.transform.LowerCPUAtomics).
|
|
"""
|
|
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from tilelang import tvm
|
|
|
|
|
|
def _compile_llvm(func, out_idx):
|
|
with tvm.target.Target("llvm"):
|
|
return tilelang.compile(func, out_idx=out_idx, target="llvm", execution_backend="tvm_ffi")
|
|
|
|
|
|
@tilelang.testing.requires_llvm
|
|
def test_llvm_atomic_add_scalar_return_prev():
|
|
N = 8
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((N,), dtype),
|
|
Init: T.Tensor((1,), dtype),
|
|
B: T.Tensor((1,), dtype),
|
|
P: T.Tensor((N,), dtype),
|
|
):
|
|
with T.Kernel(1):
|
|
B[0] = Init[0]
|
|
for i in T.serial(N):
|
|
P[i] = T.atomic_add(B[0], A[i], return_prev=True)
|
|
|
|
kernel = _compile_llvm(main, out_idx=[2, 3])
|
|
gen = torch.Generator(device="cpu").manual_seed(7)
|
|
A = torch.randn((N,), dtype=torch.float32, generator=gen)
|
|
Init = torch.full((1,), 1.0, dtype=torch.float32)
|
|
out_b, out_p = kernel(A, Init)
|
|
expected_p = Init + torch.cumsum(A, dim=0) - A
|
|
torch.testing.assert_close(out_p, expected_p, rtol=1e-4, atol=1e-4)
|
|
torch.testing.assert_close(out_b, A.sum() + Init, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
@tilelang.testing.requires_llvm
|
|
def test_llvm_atomic_max_region():
|
|
M, N = 4, 8
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((M, N), dtype), Init: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(1):
|
|
for i, j in T.grid(M, N):
|
|
B[i, j] = Init[i, j]
|
|
T.atomic_max(B, A)
|
|
|
|
kernel = _compile_llvm(main, out_idx=[2])
|
|
gen = torch.Generator(device="cpu").manual_seed(7)
|
|
A = torch.randn((M, N), dtype=torch.float32, generator=gen)
|
|
Init = torch.randn((M, N), dtype=torch.float32, generator=gen)
|
|
out = kernel(A, Init)
|
|
torch.testing.assert_close(out, torch.maximum(Init, A), rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
@tilelang.testing.requires_llvm
|
|
def test_llvm_atomic_addx4():
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(Val: T.Tensor((4,), dtype), Init: T.Tensor((4,), dtype), Dst: T.Tensor((4,), dtype)):
|
|
with T.Kernel(1):
|
|
for i in T.serial(4):
|
|
Dst[i] = Init[i]
|
|
T.atomic_addx4(Dst[0:4], Val[0:4])
|
|
|
|
kernel = _compile_llvm(main, out_idx=[2])
|
|
gen = torch.Generator(device="cpu").manual_seed(7)
|
|
Val = torch.randn((4,), dtype=torch.float32, generator=gen)
|
|
Init = torch.randn((4,), dtype=torch.float32, generator=gen)
|
|
out = kernel(Val, Init)
|
|
torch.testing.assert_close(out, Init + Val, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|