mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
236 lines
7.1 KiB
Python
236 lines
7.1 KiB
Python
"""CPU reduce op tests (target="c", execution_backend="cython").
|
|
|
|
Covers the 8 reduce variants supported by src/cpu/op/reduce.cc on local
|
|
buffers, which is the only path the frontend emits on CPU.
|
|
"""
|
|
|
|
import functools
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
from tilelang import tvm
|
|
|
|
_FLOAT_OPS = ("sum", "max", "min", "abssum", "absmax")
|
|
_BIT_OPS = ("bitand", "bitor", "bitxor")
|
|
|
|
|
|
def _torch_dtype(dtype: str) -> torch.dtype:
|
|
return getattr(torch, dtype)
|
|
|
|
|
|
def _ref_reduce(A: torch.Tensor, op: str, dim: int) -> torch.Tensor:
|
|
legal = dim if dim >= 0 else A.ndim + dim
|
|
if op == "sum":
|
|
return A.sum(dim=legal)
|
|
if op == "max":
|
|
return A.amax(dim=legal)
|
|
if op == "min":
|
|
return A.amin(dim=legal)
|
|
if op == "abssum":
|
|
return A.abs().sum(dim=legal)
|
|
if op == "absmax":
|
|
return A.abs().amax(dim=legal)
|
|
if op in _BIT_OPS:
|
|
binop = {
|
|
"bitand": torch.bitwise_and,
|
|
"bitor": torch.bitwise_or,
|
|
"bitxor": torch.bitwise_xor,
|
|
}[op]
|
|
return functools.reduce(
|
|
lambda acc, i: binop(acc, A.select(legal, i)),
|
|
range(1, A.shape[legal]),
|
|
A.select(legal, 0).clone(),
|
|
)
|
|
raise ValueError(op)
|
|
|
|
|
|
def _emit_reduce(T, op: str, src, dst, dim: int, clear: bool):
|
|
if op == "sum":
|
|
T.reduce_sum(src, dst, dim=dim, clear=clear)
|
|
elif op == "max":
|
|
T.reduce_max(src, dst, dim=dim, clear=clear)
|
|
elif op == "min":
|
|
T.reduce_min(src, dst, dim=dim, clear=clear)
|
|
elif op == "abssum":
|
|
T.reduce_abssum(src, dst, dim=dim)
|
|
elif op == "absmax":
|
|
T.reduce_absmax(src, dst, dim=dim, clear=clear)
|
|
elif op == "bitand":
|
|
T.reduce_bitand(src, dst, dim=dim, clear=clear)
|
|
elif op == "bitor":
|
|
T.reduce_bitor(src, dst, dim=dim, clear=clear)
|
|
elif op == "bitxor":
|
|
T.reduce_bitxor(src, dst, dim=dim, clear=clear)
|
|
else:
|
|
raise ValueError(op)
|
|
|
|
|
|
def _make_kernel_2d(op: str, M: int, N: int, dim: int, dtype: str, clear: bool):
|
|
legal = dim if dim >= 0 else 2 + dim
|
|
dst_shape = (N,) if legal == 0 else (M,)
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((M, N), dtype),
|
|
Init: T.Tensor(dst_shape, dtype),
|
|
B: T.Tensor(dst_shape, dtype),
|
|
):
|
|
with T.Kernel(1):
|
|
src = T.alloc_local((M, N), dtype)
|
|
dst = T.alloc_local(dst_shape, dtype)
|
|
for i, j in T.grid(M, N):
|
|
src[i, j] = A[i, j]
|
|
for i in T.grid(dst_shape[0]):
|
|
dst[i] = Init[i]
|
|
_emit_reduce(T, op, src, dst, dim=legal, clear=clear)
|
|
for i in T.grid(dst_shape[0]):
|
|
B[i] = dst[i]
|
|
|
|
return main
|
|
|
|
|
|
def _make_kernel_3d(op: str, P: int, Q: int, R: int, dim: int, dtype: str):
|
|
legal = dim if dim >= 0 else 3 + dim
|
|
if legal == 0:
|
|
dst_shape = (Q, R)
|
|
elif legal == 1:
|
|
dst_shape = (P, R)
|
|
else:
|
|
dst_shape = (P, Q)
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((P, Q, R), dtype),
|
|
B: T.Tensor(dst_shape, dtype),
|
|
):
|
|
with T.Kernel(1):
|
|
src = T.alloc_local((P, Q, R), dtype)
|
|
dst = T.alloc_local(dst_shape, dtype)
|
|
for p, q, r in T.grid(P, Q, R):
|
|
src[p, q, r] = A[p, q, r]
|
|
_emit_reduce(T, op, src, dst, dim=legal, clear=True)
|
|
if legal == 0:
|
|
for q, r in T.grid(Q, R):
|
|
B[q, r] = dst[q, r]
|
|
elif legal == 1:
|
|
for p, r in T.grid(P, R):
|
|
B[p, r] = dst[p, r]
|
|
else:
|
|
for p, q in T.grid(P, Q):
|
|
B[p, q] = dst[p, q]
|
|
|
|
return main
|
|
|
|
|
|
def _compile_c(func):
|
|
with tvm.target.Target("c"):
|
|
return tilelang.compile(
|
|
func,
|
|
out_idx=[2],
|
|
target="c",
|
|
target_host="c",
|
|
execution_backend="cython",
|
|
)
|
|
|
|
|
|
def _make_input(shape, dtype, op):
|
|
tdtype = _torch_dtype(dtype)
|
|
gen = torch.Generator(device="cpu").manual_seed(7)
|
|
if op in _BIT_OPS:
|
|
return torch.randint(-50, 50, shape, dtype=tdtype, generator=gen)
|
|
return torch.randn(shape, dtype=tdtype, generator=gen)
|
|
|
|
|
|
@pytest.mark.parametrize("op", _FLOAT_OPS)
|
|
@pytest.mark.parametrize("dim", [0, 1, -1])
|
|
def test_cpu_reduce_2d_float(op, dim):
|
|
M, N = 4, 8
|
|
dtype = "float32"
|
|
func = _make_kernel_2d(op, M, N, dim, dtype, clear=True)
|
|
kernel = _compile_c(func)
|
|
|
|
A = _make_input((M, N), dtype, op)
|
|
Init = torch.zeros((N if (dim if dim >= 0 else 2 + dim) == 0 else M,), dtype=_torch_dtype(dtype))
|
|
out = kernel(A, Init)
|
|
expected = _ref_reduce(A, op, dim)
|
|
torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
@pytest.mark.parametrize("op", _BIT_OPS)
|
|
@pytest.mark.parametrize("dim", [0, 1])
|
|
def test_cpu_reduce_2d_bit(op, dim):
|
|
M, N = 4, 8
|
|
dtype = "int32"
|
|
func = _make_kernel_2d(op, M, N, dim, dtype, clear=True)
|
|
kernel = _compile_c(func)
|
|
|
|
A = _make_input((M, N), dtype, op)
|
|
dst_shape = (N,) if (dim if dim >= 0 else 2 + dim) == 0 else (M,)
|
|
Init = torch.zeros(dst_shape, dtype=_torch_dtype(dtype))
|
|
out = kernel(A, Init)
|
|
expected = _ref_reduce(A, op, dim)
|
|
torch.testing.assert_close(out, expected)
|
|
|
|
|
|
@pytest.mark.parametrize("op", ["sum", "max", "bitand"])
|
|
@pytest.mark.parametrize("dim", [0, 1, 2])
|
|
def test_cpu_reduce_3d(op, dim):
|
|
P, Q, R = 2, 3, 4
|
|
dtype = "float32" if op != "bitand" else "int32"
|
|
func = _make_kernel_3d(op, P, Q, R, dim, dtype)
|
|
with tvm.target.Target("c"):
|
|
kernel = tilelang.compile(
|
|
func,
|
|
out_idx=[1],
|
|
target="c",
|
|
target_host="c",
|
|
execution_backend="cython",
|
|
)
|
|
|
|
A = _make_input((P, Q, R), dtype, op)
|
|
out = kernel(A)
|
|
expected = _ref_reduce(A, op, dim)
|
|
if op == "bitand":
|
|
torch.testing.assert_close(out, expected)
|
|
else:
|
|
torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
@pytest.mark.parametrize("op", ["sum", "max", "min"])
|
|
def test_cpu_reduce_clear_false_accumulates(op):
|
|
"""clear=False must accumulate onto the pre-existing dst values."""
|
|
M, N = 4, 8
|
|
dtype = "float32"
|
|
func = _make_kernel_2d(op, M, N, dim=1, dtype=dtype, clear=False)
|
|
kernel = _compile_c(func)
|
|
|
|
A = _make_input((M, N), dtype, op)
|
|
Init = torch.full((M,), 1.0, dtype=_torch_dtype(dtype))
|
|
out = kernel(A, Init)
|
|
|
|
base = _ref_reduce(A, op, dim=1)
|
|
if op == "sum":
|
|
expected = base + Init
|
|
elif op == "max":
|
|
expected = torch.maximum(base, Init)
|
|
else: # min
|
|
expected = torch.minimum(base, Init)
|
|
torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_cpu_reduce_non_divisible_extent():
|
|
"""Reduce extent of 17 (prime) exercises the kSerial loop tail."""
|
|
M, N = 3, 17
|
|
dtype = "float32"
|
|
func = _make_kernel_2d("sum", M, N, dim=1, dtype=dtype, clear=True)
|
|
kernel = _compile_c(func)
|
|
|
|
A = _make_input((M, N), dtype, "sum")
|
|
Init = torch.zeros((M,), dtype=_torch_dtype(dtype))
|
|
out = kernel(A, Init)
|
|
expected = _ref_reduce(A, "sum", dim=1)
|
|
torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-4)
|