mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
418 lines
14 KiB
Python
418 lines
14 KiB
Python
"""CPU atomic op tests (target="c", execution_backend="cython").
|
|
|
|
Two lowering paths are covered:
|
|
- tile-region: T.atomic_add/atomic_max/atomic_min on buffers/regions, lowered
|
|
by src/cpu/op/atomic_add.cc / atomic_reduce.cc to plain serial RMW loops;
|
|
- scalar elem-op: T.atomic_add/max/min/or/load/store/addx2/addx4 on addressed
|
|
elements, rewritten by tl.transform.LowerCPUAtomics to plain RMW.
|
|
|
|
On serial CPU, atomics degenerate to plain read-modify-write: memory_order is
|
|
accepted but ignored, and return_prev yields the pre-update value.
|
|
"""
|
|
|
|
import functools
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
from tilelang import tvm
|
|
|
|
|
|
def _compile_c(func, out_idx):
|
|
with tvm.target.Target("c"):
|
|
return tilelang.compile(
|
|
func,
|
|
out_idx=out_idx,
|
|
target="c",
|
|
target_host="c",
|
|
execution_backend="cython",
|
|
)
|
|
|
|
|
|
def _make_input(shape, dtype, seed=7):
|
|
tdtype = getattr(torch, dtype)
|
|
gen = torch.Generator(device="cpu").manual_seed(seed)
|
|
if tdtype in (torch.int32, torch.int64):
|
|
return torch.randint(-50, 50, shape, dtype=tdtype, generator=gen)
|
|
return torch.randn(shape, dtype=tdtype, generator=gen)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", ["float32", "int32", "float16"])
|
|
def test_cpu_atomic_add_scalar_parallel(dtype):
|
|
"""Scalar atomic_add into a global cell from a T.Parallel loop."""
|
|
N = 128
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((N,), dtype), Init: T.Tensor((1,), dtype), B: T.Tensor((1,), dtype)):
|
|
with T.Kernel(1):
|
|
B[0] = Init[0]
|
|
for i in T.Parallel(N):
|
|
T.atomic_add(B[0], A[i])
|
|
|
|
kernel = _compile_c(main, out_idx=[2])
|
|
A = _make_input((N,), dtype)
|
|
Init = torch.zeros((1,), dtype=getattr(torch, dtype))
|
|
out = kernel(A, Init)
|
|
expected = A.sum() + Init
|
|
if dtype == "float16":
|
|
torch.testing.assert_close(out, expected.reshape((1,)), rtol=1e-2, atol=1e-2)
|
|
else:
|
|
torch.testing.assert_close(out, expected.reshape((1,)))
|
|
|
|
|
|
def test_cpu_atomic_add_scalar_serial_memory_order():
|
|
"""Serial loop + memory_order (accepted, ignored) + codegen assertion."""
|
|
N = 64
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((N,), dtype), Init: T.Tensor((1,), dtype), B: T.Tensor((1,), dtype)):
|
|
with T.Kernel(1):
|
|
B[0] = Init[0]
|
|
for i in T.serial(N):
|
|
T.atomic_add(B[0], A[i], memory_order="seq_cst")
|
|
|
|
kernel = _compile_c(main, out_idx=[2])
|
|
# The intrinsic must be gone from the generated source (plain RMW).
|
|
assert "atomic_add_elem_op" not in kernel.get_kernel_source()
|
|
|
|
A = _make_input((N,), dtype)
|
|
Init = torch.full((1,), 1.0, dtype=torch.float32)
|
|
out = kernel(A, Init)
|
|
torch.testing.assert_close(out, A.sum() + Init, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_cpu_atomic_add_scalar_return_prev():
|
|
"""return_prev yields the pre-update value (exclusive prefix sum)."""
|
|
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_c(main, out_idx=[2, 3])
|
|
A = _make_input((N,), dtype)
|
|
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)
|
|
|
|
|
|
@pytest.mark.parametrize("op", ["max", "min"])
|
|
def test_cpu_atomic_max_min_scalar_return_prev(op):
|
|
"""return_prev with max/min combine: exclusive running prefix."""
|
|
N = 8
|
|
dtype = "float32"
|
|
emit = T.atomic_max if op == "max" else T.atomic_min
|
|
|
|
@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] = emit(B[0], A[i], return_prev=True)
|
|
|
|
kernel = _compile_c(main, out_idx=[2, 3])
|
|
A = _make_input((N,), dtype)
|
|
Init = torch.zeros((1,), dtype=torch.float32)
|
|
out_b, out_p = kernel(A, Init)
|
|
|
|
running = Init.clone()
|
|
expected_p = torch.empty_like(A)
|
|
combine = torch.maximum if op == "max" else torch.minimum
|
|
for i in range(N):
|
|
expected_p[i] = running[0]
|
|
running = combine(running, A[i : i + 1])
|
|
torch.testing.assert_close(out_p, expected_p, rtol=1e-4, atol=1e-4)
|
|
expected_b = combine(Init, A.max() if op == "max" else A.min())
|
|
torch.testing.assert_close(out_b, expected_b, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_cpu_atomic_add_scalar_int_value_cast():
|
|
"""An int value added to a float buffer is cast to the dst dtype."""
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(Init: T.Tensor((1,), dtype), B: T.Tensor((1,), dtype)):
|
|
with T.Kernel(1):
|
|
B[0] = Init[0]
|
|
T.atomic_add(B[0], 3)
|
|
|
|
kernel = _compile_c(main, out_idx=[1])
|
|
Init = torch.full((1,), 1.5, dtype=torch.float32)
|
|
out = kernel(Init)
|
|
torch.testing.assert_close(out, Init + 3.0)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"op,dtype",
|
|
[("max", "float32"), ("min", "float32"), ("max", "float16"), ("min", "int32")],
|
|
)
|
|
def test_cpu_atomic_max_min_scalar(op, dtype):
|
|
N = 64
|
|
emit = T.atomic_max if op == "max" else T.atomic_min
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((N,), dtype), Init: T.Tensor((1,), dtype), B: T.Tensor((1,), dtype)):
|
|
with T.Kernel(1):
|
|
B[0] = Init[0]
|
|
for i in T.Parallel(N):
|
|
emit(B[0], A[i])
|
|
|
|
kernel = _compile_c(main, out_idx=[2])
|
|
A = _make_input((N,), dtype)
|
|
base = A.max() if op == "max" else A.min()
|
|
Init = (base - 1).reshape((1,)) if op == "max" else (base + 1).reshape((1,))
|
|
out = kernel(A, Init)
|
|
expected = torch.maximum(A.max(), Init) if op == "max" else torch.minimum(A.min(), Init)
|
|
if dtype == "float16":
|
|
torch.testing.assert_close(out, expected.reshape((1,)), rtol=1e-2, atol=1e-2)
|
|
else:
|
|
torch.testing.assert_close(out, expected.reshape((1,)))
|
|
|
|
|
|
def test_cpu_atomic_or_scalar():
|
|
N = 32
|
|
dtype = "int32"
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((N,), dtype), Init: T.Tensor((1,), dtype), B: T.Tensor((1,), dtype)):
|
|
with T.Kernel(1):
|
|
B[0] = Init[0]
|
|
for i in T.Parallel(N):
|
|
T.atomic_or(B[0], A[i])
|
|
|
|
kernel = _compile_c(main, out_idx=[2])
|
|
A = _make_input((N,), dtype)
|
|
Init = torch.zeros((1,), dtype=torch.int32)
|
|
out = kernel(A, Init)
|
|
expected = functools.reduce(torch.bitwise_or, list(A.unbind()), Init[0].clone())
|
|
torch.testing.assert_close(out, expected.reshape((1,)))
|
|
|
|
|
|
def test_cpu_atomic_load_store():
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((1,), dtype), B: T.Tensor((1,), dtype)):
|
|
with T.Kernel(1):
|
|
x = T.atomic_load(A[0], memory_order="acquire")
|
|
T.atomic_store(B[0], x + T.float32(1.0), memory_order="release")
|
|
|
|
kernel = _compile_c(main, out_idx=[1])
|
|
A = torch.full((1,), 2.5, dtype=torch.float32)
|
|
out = kernel(A)
|
|
torch.testing.assert_close(out, A + 1.0)
|
|
|
|
|
|
def test_cpu_atomic_addx2_fp16():
|
|
"""x2 vector atomic add, full-buffer form, expanded per element."""
|
|
dtype = "float16"
|
|
|
|
@T.prim_func
|
|
def main(Val: T.Tensor((2,), dtype), Init: T.Tensor((2,), dtype), Dst: T.Tensor((2,), dtype)):
|
|
with T.Kernel(1):
|
|
for i in T.serial(2):
|
|
Dst[i] = Init[i]
|
|
T.atomic_addx2(Dst, Val)
|
|
|
|
kernel = _compile_c(main, out_idx=[2])
|
|
Val = _make_input((2,), dtype)
|
|
Init = _make_input((2,), dtype, seed=11)
|
|
out = kernel(Val, Init)
|
|
torch.testing.assert_close(out, Init + Val, rtol=1e-2, atol=1e-2)
|
|
|
|
|
|
def test_cpu_atomic_addx4_fp32_slice():
|
|
"""x4 vector atomic add, slice form (region min offsets)."""
|
|
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_c(main, out_idx=[2])
|
|
Val = _make_input((4,), dtype)
|
|
Init = _make_input((4,), dtype, seed=11)
|
|
out = kernel(Val, Init)
|
|
torch.testing.assert_close(out, Init + Val, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", ["float32", "int32"])
|
|
def test_cpu_atomic_add_region_global(dtype):
|
|
"""Tile-region atomic_add with a global dst buffer."""
|
|
M, N = 4, 8
|
|
|
|
@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_add(B, A)
|
|
|
|
kernel = _compile_c(main, out_idx=[2])
|
|
A = _make_input((M, N), dtype)
|
|
Init = _make_input((M, N), dtype, seed=11)
|
|
out = kernel(A, Init)
|
|
torch.testing.assert_close(out, Init + A, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
@pytest.mark.parametrize("op", ["max", "min"])
|
|
def test_cpu_atomic_max_min_region_local_dst(op):
|
|
"""Tile-region atomic_max/atomic_min with a local dst buffer."""
|
|
M, N = 4, 8
|
|
dtype = "float32"
|
|
emit = T.atomic_max if op == "max" else T.atomic_min
|
|
|
|
@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):
|
|
dst = T.alloc_local((M, N), dtype)
|
|
for i, j in T.grid(M, N):
|
|
dst[i, j] = Init[i, j]
|
|
emit(dst, A)
|
|
for i, j in T.grid(M, N):
|
|
B[i, j] = dst[i, j]
|
|
|
|
kernel = _compile_c(main, out_idx=[2])
|
|
A = _make_input((M, N), dtype)
|
|
Init = _make_input((M, N), dtype, seed=11)
|
|
out = kernel(A, Init)
|
|
expected = torch.maximum(Init, A) if op == "max" else torch.minimum(Init, A)
|
|
torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_cpu_atomic_add_region_scalar_value():
|
|
"""Tile-region atomic_add with a scalar src value (src_value path)."""
|
|
M, N = 4, 8
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(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_add(B, 2)
|
|
|
|
kernel = _compile_c(main, out_idx=[1])
|
|
Init = _make_input((M, N), dtype)
|
|
out = kernel(Init)
|
|
torch.testing.assert_close(out, Init + 2.0, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_cpu_atomic_add_region_sliced_dst():
|
|
"""Sliced dst region: dst_range carries a non-zero min offset."""
|
|
M, N = 4, 8
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((M - 2, 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_add(B[1 : M - 1, :], A)
|
|
|
|
kernel = _compile_c(main, out_idx=[2])
|
|
A = _make_input((M - 2, N), dtype)
|
|
Init = _make_input((M, N), dtype, seed=11)
|
|
out = kernel(A, Init)
|
|
expected = Init.clone()
|
|
expected[1 : M - 1, :] += A
|
|
torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_cpu_atomic_add_region_memory_order():
|
|
"""memory_order annotation on the tile-region path is accepted/ignored."""
|
|
M, N = 4, 8
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(A: 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] = 0.0
|
|
T.atomic_add(B, A, memory_order="acquire")
|
|
|
|
kernel = _compile_c(main, out_idx=[1])
|
|
A = _make_input((M, N), dtype)
|
|
out = kernel(A)
|
|
torch.testing.assert_close(out, A, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_cpu_atomic_add_use_tma_rejected():
|
|
M, N = 4, 8
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(1):
|
|
T.atomic_add(B, A, use_tma=True)
|
|
|
|
with pytest.raises(Exception, match="use_tma"):
|
|
_compile_c(main, out_idx=[1])
|
|
|
|
|
|
def test_cpu_atomic_addx2_return_prev_rejected():
|
|
"""Vector return_prev is rejected with a readable error on CPU."""
|
|
|
|
@T.prim_func
|
|
def main(
|
|
Dst: T.Tensor((2,), "float16"),
|
|
Val: T.Tensor((2,), "float16"),
|
|
Prev: T.Tensor((2,), "float16"),
|
|
):
|
|
with T.Kernel(1):
|
|
Prev[0:2] = T.atomic_addx2(Dst[0:2], Val[0:2], return_prev=True)
|
|
|
|
with pytest.raises(Exception, match="return_prev"):
|
|
_compile_c(main, out_idx=[2])
|
|
|
|
|
|
def test_cpu_atomic_multiple_return_prev_one_statement():
|
|
"""Two return_prev atomics in one statement: both Bind/Store prefixes are
|
|
spliced, in order, before the statement (pending_prefix_ isolation)."""
|
|
N = 8
|
|
dtype = "float32"
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((N,), dtype),
|
|
B0: T.Tensor((1,), dtype),
|
|
B1: T.Tensor((1,), dtype),
|
|
P: T.Tensor((N,), dtype),
|
|
):
|
|
with T.Kernel(1):
|
|
B0[0] = 0.0
|
|
B1[0] = 100.0
|
|
for i in T.serial(N):
|
|
P[i] = T.atomic_add(B0[0], A[i], return_prev=True) + T.atomic_add(B1[0], A[i], return_prev=True)
|
|
|
|
kernel = _compile_c(main, out_idx=[1, 2, 3])
|
|
A = _make_input((N,), dtype)
|
|
out_b0, out_b1, out_p = kernel(A)
|
|
|
|
prefix = torch.cumsum(A, dim=0) - A # sum(A[:i])
|
|
torch.testing.assert_close(out_p, prefix + (100.0 + prefix), rtol=1e-4, atol=1e-4)
|
|
torch.testing.assert_close(out_b0, A.sum().reshape((1,)), rtol=1e-4, atol=1e-4)
|
|
torch.testing.assert_close(out_b1, 100.0 + A.sum().reshape((1,)), rtol=1e-4, atol=1e-4)
|