mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Jayce Suanddingsg
39ade83ff2
Use PTX v4 atomics for 16-bit CUDA atomic_addx4 The initial fp16/bf16 atomic_addx4 fix avoided the generic float4 path by splitting each operation into two AtomicAddx2 calls. That fixed correctness, but it still used the simulated implementation on Hopper where PTX provides native vector atomics. Add sm90+ half_t and bfloat16_t AtomicAddx4 overloads backed by the PTX vector forms atom.global.v4.f16.add.noftz and atom.global.v4.bf16.add.noftz, including the existing memory-order variants. Keep the x2 fallback for pre-sm90 targets because PTX vector f16/bf16 atomics require sm90 or newer. Keep the CUDA language regression focused on fp16/bf16 atomic_addx4 codegen and offset coverage. Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
723 lines
25 KiB
Python
723 lines
25 KiB
Python
import tilelang.testing
|
|
import tilelang.layout
|
|
import tilelang.language as T
|
|
import torch
|
|
|
|
|
|
# ======================= Thread-level atomic add =======================
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_add_program(K, M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_add(A: T.Tensor((K, M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), K, threads=32) as (bx, by, bz):
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.copy(A[bz, bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], A_shared)
|
|
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
T.atomic_add(B[bx * block_M + i, by * block_N + j], A_shared[i, j])
|
|
|
|
return atomic_add
|
|
|
|
|
|
def run_atomic_add(K, M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = atomic_add_program(K, M, N, block_M, block_N, dtype=dtype)
|
|
import torch
|
|
|
|
A = torch.randn(K, M, N, dtype=getattr(torch, dtype)).cuda()
|
|
B = torch.zeros(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
ref_B = B + A.sum(dim=0)
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_memory_order_program(K, M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_with_memory_order(A: T.Tensor((K, M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), K, threads=32) as (bx, by, bz):
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.copy(A[bz, bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], A_shared)
|
|
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
T.atomic_add(B[bx * block_M + i, by * block_N + j], A_shared[i, j], memory_order="relaxed")
|
|
|
|
return atomic_with_memory_order
|
|
|
|
|
|
def run_atomic_memory_order(K, M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = atomic_memory_order_program(K, M, N, block_M, block_N, dtype=dtype)
|
|
import torch
|
|
|
|
A = torch.randn(K, M, N, dtype=getattr(torch, dtype)).cuda()
|
|
B = torch.zeros(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
ref_B = B + A.sum(dim=0)
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_addx2_program(M, N, block_M, block_N, dtype=T.float16):
|
|
@T.prim_func
|
|
def atomic_addx2(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=32) as (bx, by):
|
|
for i, j in T.Parallel(block_M, block_N // 2):
|
|
idx_i = bx * block_M + i
|
|
idx_j = by * block_N + j * 2
|
|
T.atomic_addx2(B[idx_i, idx_j], A[idx_i, idx_j])
|
|
|
|
return atomic_addx2
|
|
|
|
|
|
def run_atomic_addx2(M, N, block_M, block_N, dtype=T.float16):
|
|
kernel = atomic_addx2_program(M, N, block_M, block_N, dtype=dtype)
|
|
|
|
import torch
|
|
|
|
A = torch.randn(M, N, dtype=torch.float32).cuda().to(getattr(torch, dtype))
|
|
B = torch.zeros(M, N, dtype=torch.float32).cuda().to(getattr(torch, dtype))
|
|
ref_B = B.clone()
|
|
|
|
for i in range(M):
|
|
for j in range(0, N - 1, 2):
|
|
ref_B[i, j] += A[i, j]
|
|
ref_B[i, j + 1] += A[i, j + 1]
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_add_mixed_dtype_program(N, src_dtype, dst_dtype):
|
|
@T.prim_func
|
|
def atomic_add(Src: T.Tensor((N,), src_dtype), Out: T.Tensor((N,), dst_dtype)):
|
|
with T.Kernel(threads=1):
|
|
frag = T.alloc_fragment((N,), src_dtype)
|
|
for i in T.Parallel(N):
|
|
frag[i] = Src[i]
|
|
for i in T.Parallel(N):
|
|
T.atomic_add(Out[i], frag[i])
|
|
|
|
return atomic_add
|
|
|
|
|
|
def run_atomic_add_mixed_dtype(N, src_dtype, dst_dtype):
|
|
kernel = atomic_add_mixed_dtype_program(N, src_dtype, dst_dtype)
|
|
assert "AtomicAddx2" in kernel.get_kernel_source()
|
|
|
|
src = torch.arange(1, N + 1, dtype=getattr(torch, src_dtype)).cuda()
|
|
out = torch.zeros(N, dtype=getattr(torch, dst_dtype)).cuda()
|
|
kernel(src, out)
|
|
torch.testing.assert_close(out, src.to(getattr(torch, dst_dtype)), atol=1e-2, rtol=1e-2)
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_addx2_mixed_dtype_program(M, N, block_M, block_N, src_dtype, dst_dtype):
|
|
@T.prim_func
|
|
def atomic_addx2(A: T.Tensor((M, N), src_dtype), B: T.Tensor((M, N), dst_dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=32) as (bx, by):
|
|
for i, j in T.Parallel(block_M, block_N // 2):
|
|
idx_i = bx * block_M + i
|
|
idx_j = by * block_N + j * 2
|
|
T.atomic_addx2(B[idx_i, idx_j], A[idx_i, idx_j])
|
|
|
|
return atomic_addx2
|
|
|
|
|
|
def run_atomic_addx2_mixed_dtype(M, N, block_M, block_N, src_dtype, dst_dtype):
|
|
kernel = atomic_addx2_mixed_dtype_program(M, N, block_M, block_N, src_dtype, dst_dtype)
|
|
assert "AtomicAddx2" in kernel.get_kernel_source()
|
|
|
|
A = torch.randn(M, N, dtype=getattr(torch, src_dtype)).cuda()
|
|
B = torch.zeros(M, N, dtype=getattr(torch, dst_dtype)).cuda()
|
|
ref_B = A.to(getattr(torch, dst_dtype))
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-2, rtol=1e-2)
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_different_memory_orders_program(M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_different_orders(
|
|
A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype), C: T.Tensor((M, N), dtype), D: T.Tensor((M, N), dtype)
|
|
):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=32) as (bx, by):
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
idx_i = bx * block_M + i
|
|
idx_j = by * block_N + j
|
|
if idx_i < M and idx_j < N:
|
|
val = A[idx_i, idx_j]
|
|
T.atomic_add(B[idx_i, idx_j], val, memory_order="release")
|
|
T.atomic_max(C[idx_i, idx_j], val, memory_order="relaxed")
|
|
T.atomic_min(D[idx_i, idx_j], val, memory_order="relaxed")
|
|
|
|
return atomic_different_orders
|
|
|
|
|
|
def run_atomic_different_memory_orders(M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = atomic_different_memory_orders_program(M, N, block_M, block_N, dtype=dtype)
|
|
import torch
|
|
|
|
A = torch.randn(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
B = torch.zeros(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
C = torch.zeros(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
D = torch.full((M, N), float("inf"), dtype=getattr(torch, dtype)).cuda()
|
|
|
|
kernel(A, B, C, D)
|
|
|
|
torch.testing.assert_close(B, A, atol=1e-3, rtol=1e-3)
|
|
torch.testing.assert_close(C, torch.maximum(torch.zeros_like(A), A))
|
|
torch.testing.assert_close(D, torch.minimum(torch.full_like(A, float("inf")), A))
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_addx4_program(M, N, block_M, block_N):
|
|
@T.prim_func
|
|
def atomic_addx4(A: T.Tensor((M, N), T.float32), B: T.Tensor((M, N), T.float32)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=32) as (bx, by):
|
|
for i, j in T.Parallel(block_M, block_N // 4):
|
|
idx_i = bx * block_M + i
|
|
idx_j = by * block_N + j * 4
|
|
T.atomic_addx4(B[idx_i, idx_j], A[idx_i, idx_j])
|
|
|
|
return atomic_addx4
|
|
|
|
|
|
def run_atomic_addx4(M, N, block_M, block_N):
|
|
kernel = atomic_addx4_program(M, N, block_M, block_N)
|
|
import torch
|
|
|
|
A = torch.randn(M, N, dtype=torch.float32).cuda()
|
|
B = torch.zeros(M, N, dtype=torch.float32).cuda()
|
|
ref_B = B.clone()
|
|
|
|
for i in range(M):
|
|
for j in range(0, N - 3, 4):
|
|
ref_B[i, j] += A[i, j]
|
|
ref_B[i, j + 1] += A[i, j + 1]
|
|
ref_B[i, j + 2] += A[i, j + 2]
|
|
ref_B[i, j + 3] += A[i, j + 3]
|
|
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_addx4_16bit_program(dtype, offset, nthreads):
|
|
@T.prim_func
|
|
def atomic_addx4(A: T.Tensor((16,), dtype), B: T.Tensor((16,), dtype)):
|
|
with T.Kernel(1, threads=nthreads):
|
|
T.atomic_addx4(B[offset], A[offset])
|
|
|
|
return atomic_addx4
|
|
|
|
|
|
def run_atomic_addx4_16bit(dtype, offset, nthreads):
|
|
kernel = atomic_addx4_16bit_program(dtype, offset, nthreads)
|
|
source = kernel.get_kernel_source()
|
|
assert "AtomicAddx4" in source
|
|
|
|
torch_dtype = getattr(torch, str(dtype))
|
|
A = torch.zeros(16, dtype=torch_dtype, device="cuda")
|
|
B_init = torch.zeros(16, dtype=torch_dtype, device="cuda")
|
|
A[offset : offset + 4] = torch.tensor([1, 2, 3, 4], dtype=torch_dtype, device="cuda")
|
|
B_init[offset : offset + 4] = torch.tensor([10, 20, 30, 40], dtype=torch_dtype, device="cuda")
|
|
B = B_init.clone()
|
|
|
|
ref_B = B_init.float()
|
|
ref_B[offset : offset + 4] += nthreads * A[offset : offset + 4].float()
|
|
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B.float(), ref_B, atol=0, rtol=0)
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_return_prev_program(M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_with_return_prev(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype), old_vals: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=32) as (bx, by):
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
idx_i = bx * block_M + i
|
|
idx_j = by * block_N + j
|
|
if idx_i < M and idx_j < N:
|
|
old_vals[idx_i, idx_j] = T.atomic_add(B[idx_i, idx_j], A[idx_i, idx_j], return_prev=True)
|
|
|
|
return atomic_with_return_prev
|
|
|
|
|
|
def run_atomic_return_prev(M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = atomic_return_prev_program(M, N, block_M, block_N, dtype=dtype)
|
|
import torch
|
|
|
|
A = torch.ones(M, N, dtype=getattr(torch, dtype)).cuda() * 5.0
|
|
B = torch.ones(M, N, dtype=getattr(torch, dtype)).cuda() * 2.0
|
|
old_vals = torch.zeros(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
|
|
initial_B = B.clone()
|
|
kernel(A, B, old_vals)
|
|
|
|
torch.testing.assert_close(old_vals, initial_B, atol=1e-3, rtol=1e-3)
|
|
torch.testing.assert_close(B, initial_B + A, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def tma_atomic_add_program(out, explicit_swizzle=False):
|
|
out: T.Tensor[(16, 16), T.float32]
|
|
|
|
with T.Kernel(
|
|
1,
|
|
):
|
|
out_shared = T.alloc_shared((16, 16), dtype=T.float32)
|
|
if explicit_swizzle:
|
|
T.annotate_layout({out_shared: tilelang.layout.make_swizzled_layout(out_shared)})
|
|
T.fill(out_shared, 1)
|
|
for _ in range(16):
|
|
T.atomic_add(out, out_shared, use_tma=True)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_tma_atomic_add():
|
|
out = torch.zeros((16, 16), dtype=torch.float32, device="cuda")
|
|
tma_atomic_add_program(out)
|
|
torch.testing.assert_close(out, torch.ones((16, 16), dtype=torch.float32, device="cuda") * 16)
|
|
|
|
kernel = tma_atomic_add_program.compile(out=T.Tensor[(16, 16), T.float32])
|
|
assert "tma_store_add" in kernel.get_kernel_source()
|
|
assert "desc" in kernel.get_kernel_source() # Ensure using cp.reduce.async.bulk.tensor
|
|
|
|
kernel_with_explicit_swizzle = tma_atomic_add_program.compile(out=T.Tensor[(16, 16), T.float32], explicit_swizzle=True)
|
|
# Ensure auto swizzled layout is applied
|
|
assert kernel.get_kernel_source() == kernel_with_explicit_swizzle.get_kernel_source()
|
|
|
|
|
|
def run_atomic_add_auto_vectorized(K, M, N, block_M, block_N, dtype=T.float32):
|
|
tilelang.disable_cache()
|
|
kernel = atomic_add_program(K, M, N, block_M, block_N, dtype=dtype)
|
|
assert "AtomicAddx4" in kernel.get_kernel_source()
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_add_auto_vectorized_unit_test(vec_size: int, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_addx2(A: T.Tensor((vec_size,), dtype)):
|
|
with T.Kernel(threads=1):
|
|
A_local = T.alloc_fragment((vec_size,), dtype)
|
|
for i in T.Parallel(vec_size):
|
|
T.atomic_add(A[i], A_local[i])
|
|
|
|
return atomic_addx2
|
|
|
|
|
|
def run_atomic_add_auto_vectorized_unit_test(vec_size: int, dtype=T.float32):
|
|
kernel = atomic_add_auto_vectorized_unit_test(vec_size, dtype)
|
|
assert f"AtomicAddx{vec_size}" in kernel.get_kernel_source()
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_add_complicated_parallel_program(K, M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_add(A: T.Tensor((K, M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), K, threads=32) as (bx, by, bz):
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.copy(A[bz, bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], A_shared)
|
|
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
value = A_shared[i, j]
|
|
T.atomic_add(B[bx * block_M + i, by * block_N + j], value)
|
|
|
|
return atomic_add
|
|
|
|
|
|
def run_atomic_add_complicated_parallel(K, M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = atomic_add_complicated_parallel_program(K, M, N, block_M, block_N, dtype=dtype)
|
|
assert "float4 value" in kernel.get_kernel_source()
|
|
assert "AtomicAddx4" in kernel.get_kernel_source()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_atomic_memory_order():
|
|
run_atomic_memory_order(4, 64, 64, 16, 16)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_atomic_addx2_half():
|
|
run_atomic_addx2(32, 64, 8, 16, dtype=T.float16)
|
|
|
|
|
|
def test_atomic_addx2_float():
|
|
run_atomic_addx2(32, 64, 8, 16, dtype=T.float32)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_atomic_add_mixed_dtype_fp16():
|
|
run_atomic_add_mixed_dtype(8, T.float32, T.float16)
|
|
run_atomic_addx2_mixed_dtype(32, 64, 8, 16, T.float32, T.float16)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_ge(8, 0)
|
|
def test_atomic_add_mixed_dtype_bf16():
|
|
run_atomic_add_mixed_dtype(8, T.float32, T.bfloat16)
|
|
run_atomic_addx2_mixed_dtype(32, 64, 8, 16, T.float32, T.bfloat16)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_atomic_different_memory_orders():
|
|
run_atomic_different_memory_orders(32, 32, 8, 8, dtype=T.float32)
|
|
run_atomic_different_memory_orders(32, 32, 8, 8, dtype=T.float16)
|
|
run_atomic_different_memory_orders(32, 32, 8, 8, dtype=T.bfloat16)
|
|
|
|
|
|
def test_atomic_addx4():
|
|
run_atomic_addx4(16, 64, 4, 4)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_ge(8, 0)
|
|
def test_atomic_addx4_16bit():
|
|
for dtype in (T.float16, T.bfloat16):
|
|
for offset in (0, 4):
|
|
run_atomic_addx4_16bit(dtype, offset=offset, nthreads=2)
|
|
|
|
|
|
def test_atomic_return_prev():
|
|
run_atomic_return_prev(32, 32, 8, 8)
|
|
|
|
|
|
def test_atomic_add():
|
|
run_atomic_add(8, 128, 128, 32, 32)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
|
|
def test_atomic_add_auto_vectorized():
|
|
run_atomic_add_auto_vectorized(8, 128, 128, 32, 32, dtype=T.float32)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
|
|
def test_atomic_add_auto_vectorized_unit_test():
|
|
run_atomic_add_auto_vectorized_unit_test(2, dtype=T.float32)
|
|
run_atomic_add_auto_vectorized_unit_test(4, dtype=T.float32)
|
|
run_atomic_add_auto_vectorized_unit_test(2, dtype=T.float16)
|
|
run_atomic_add_auto_vectorized_unit_test(2, dtype=T.bfloat16)
|
|
|
|
|
|
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
|
|
def test_atomic_add_complicated_parallel():
|
|
run_atomic_add_complicated_parallel(8, 128, 128, 32, 32, dtype=T.float32)
|
|
|
|
|
|
# ======================= Tile-level atomic add =======================
|
|
|
|
|
|
@tilelang.jit
|
|
def tile_atomic_add_program(K, M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_add(A: T.Tensor((K, M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), K, threads=32) as (bx, by, bz):
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.copy(A[bz, bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], A_shared)
|
|
|
|
T.atomic_add(B[bx * block_M, by * block_N], A_shared)
|
|
|
|
return atomic_add
|
|
|
|
|
|
def run_tile_atomic_add(K, M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = tile_atomic_add_program(K, M, N, block_M, block_N, dtype=dtype)
|
|
import torch
|
|
|
|
A = torch.randn(K, M, N, dtype=getattr(torch, dtype)).cuda()
|
|
B = torch.zeros(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
ref_B = B + A.sum(dim=0)
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def tile_atomic_add_expr_program(M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_add(A: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=32) as (bx, by):
|
|
T.atomic_add(A[bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], 1.0)
|
|
|
|
return atomic_add
|
|
|
|
|
|
def run_tile_atomic_add_expr(M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = tile_atomic_add_expr_program(M, N, block_M, block_N, dtype=dtype)
|
|
import torch
|
|
|
|
def ref_program(A):
|
|
for i in range(M):
|
|
for j in range(N):
|
|
A[i, j] += 1
|
|
|
|
A = torch.zeros(M, N, dtype=torch.float32).cuda()
|
|
ref_A = A.clone()
|
|
ref_program(ref_A)
|
|
kernel(A)
|
|
torch.testing.assert_close(A, ref_A, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def tile_atomic_add_scalar_program(dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_add(A: T.Tensor((1), dtype), B: T.Tensor((1), dtype)):
|
|
with T.Kernel(
|
|
1,
|
|
) as _:
|
|
A_local = T.alloc_local([1], dtype)
|
|
T.copy(A, A_local)
|
|
T.clear(B)
|
|
T.atomic_add(B, A_local)
|
|
T.atomic_add(B, 1)
|
|
|
|
return atomic_add
|
|
|
|
|
|
def run_tile_atomic_add_scalar(dtype=T.float32):
|
|
kernel = tile_atomic_add_scalar_program(dtype=dtype)
|
|
import torch
|
|
|
|
def ref_program(A, B):
|
|
B[0] = A[0] + 1
|
|
|
|
A = torch.randn(1, dtype=getattr(torch, dtype)).cuda()
|
|
B = torch.zeros(1, dtype=getattr(torch, dtype)).cuda()
|
|
ref_B = B.clone()
|
|
ref_program(A, ref_B)
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_tile_atomic_add():
|
|
run_tile_atomic_add(8, 128, 128, 32, 32)
|
|
|
|
|
|
def test_tile_atomic_add_expr():
|
|
run_tile_atomic_add_expr(128, 128, 32, 32)
|
|
|
|
|
|
def test_tile_atomic_add_scalar():
|
|
run_tile_atomic_add_scalar()
|
|
|
|
|
|
# ======================= Thread-level atomic max/min/load store =======================
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_max_program(K, M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_max(A: T.Tensor((K, M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), K, threads=32) as (bx, by, bz):
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.copy(A[bz, bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], A_shared)
|
|
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
T.atomic_max(B[bx * block_M + i, by * block_N + j], A_shared[i, j])
|
|
|
|
return atomic_max
|
|
|
|
|
|
def run_atomic_max(K, M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = atomic_max_program(K, M, N, block_M, block_N, dtype=dtype)
|
|
import torch
|
|
|
|
def ref_program(A, B):
|
|
for k in range(K):
|
|
for i in range(M):
|
|
for j in range(N):
|
|
B[i, j] = max(B[i, j], A[k, i, j])
|
|
|
|
A = torch.randn(K, M, N, dtype=getattr(torch, dtype)).cuda()
|
|
B = torch.zeros(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
ref_B = B.clone()
|
|
ref_program(A, ref_B)
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_min_program(K, M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_min(A: T.Tensor((K, M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), K, threads=32) as (bx, by, bz):
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.copy(A[bz, bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], A_shared)
|
|
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
T.atomic_min(B[bx * block_M + i, by * block_N + j], A_shared[i, j])
|
|
|
|
return atomic_min
|
|
|
|
|
|
def run_atomic_min(K, M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = atomic_min_program(K, M, N, block_M, block_N, dtype=dtype)
|
|
import torch
|
|
|
|
def ref_program(A, B):
|
|
for k in range(K):
|
|
for i in range(M):
|
|
for j in range(N):
|
|
B[i, j] = min(B[i, j], A[k, i, j])
|
|
|
|
A = torch.randn(K, M, N, dtype=getattr(torch, dtype)).cuda()
|
|
B = torch.full((M, N), float("inf"), dtype=getattr(torch, dtype)).cuda()
|
|
ref_B = B.clone()
|
|
ref_program(A, ref_B)
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_load_store_program(M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_load_store(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=32) as (bx, by):
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
idx_i = bx * block_M + i
|
|
idx_j = by * block_N + j
|
|
if idx_i < M and idx_j < N:
|
|
val = T.atomic_load(A[idx_i, idx_j])
|
|
T.atomic_store(B[idx_i, idx_j], val)
|
|
|
|
return atomic_load_store
|
|
|
|
|
|
def run_atomic_load_store(M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = atomic_load_store_program(M, N, block_M, block_N, dtype=dtype)
|
|
import torch
|
|
|
|
A = torch.randn(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
B = torch.zeros(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, A, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_atomic_max():
|
|
run_atomic_max(4, 64, 64, 16, 16)
|
|
|
|
|
|
def test_atomic_min():
|
|
run_atomic_min(4, 64, 64, 16, 16)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_atomic_load_store():
|
|
run_atomic_load_store(64, 64, 16, 16)
|
|
|
|
|
|
# ======================= Tile-level atomic max/min =======================
|
|
|
|
|
|
@tilelang.jit
|
|
def tile_atomic_max_program(K, M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def tile_atomic_max(A: T.Tensor((K, M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), K, threads=32) as (bx, by, bz):
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.copy(A[bz, bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], A_shared)
|
|
|
|
T.atomic_max(B[bx * block_M, by * block_N], A_shared)
|
|
|
|
return tile_atomic_max
|
|
|
|
|
|
def run_tile_atomic_max(K, M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = tile_atomic_max_program(K, M, N, block_M, block_N, dtype=dtype)
|
|
|
|
def ref_program(A, B):
|
|
for k in range(K):
|
|
for i in range(M):
|
|
for j in range(N):
|
|
B[i, j] = max(B[i, j], A[k, i, j])
|
|
|
|
A = torch.randn(K, M, N, dtype=getattr(torch, dtype)).cuda()
|
|
B = torch.full((M, N), float("-inf"), dtype=getattr(torch, dtype)).cuda()
|
|
ref_B = B.clone()
|
|
ref_program(A, ref_B)
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def tile_atomic_min_program(K, M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def tile_atomic_min(A: T.Tensor((K, M, N), dtype), B: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), K, threads=32) as (bx, by, bz):
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.copy(A[bz, bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], A_shared)
|
|
|
|
T.atomic_min(B[bx * block_M, by * block_N], A_shared)
|
|
|
|
return tile_atomic_min
|
|
|
|
|
|
def run_tile_atomic_min(K, M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = tile_atomic_min_program(K, M, N, block_M, block_N, dtype=dtype)
|
|
|
|
def ref_program(A, B):
|
|
for k in range(K):
|
|
for i in range(M):
|
|
for j in range(N):
|
|
B[i, j] = min(B[i, j], A[k, i, j])
|
|
|
|
A = torch.randn(K, M, N, dtype=getattr(torch, dtype)).cuda()
|
|
B = torch.full((M, N), float("inf"), dtype=getattr(torch, dtype)).cuda()
|
|
ref_B = B.clone()
|
|
ref_program(A, ref_B)
|
|
kernel(A, B)
|
|
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@tilelang.jit
|
|
def tile_atomic_max_expr_program(M, N, block_M, block_N, dtype=T.float32):
|
|
@T.prim_func
|
|
def atomic_max(A: T.Tensor((M, N), dtype)):
|
|
with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=32) as (bx, by):
|
|
T.atomic_max(A[bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N], 0.5)
|
|
|
|
return atomic_max
|
|
|
|
|
|
def run_tile_atomic_max_expr(M, N, block_M, block_N, dtype=T.float32):
|
|
kernel = tile_atomic_max_expr_program(M, N, block_M, block_N, dtype=dtype)
|
|
import torch
|
|
|
|
def ref_program(A):
|
|
for i in range(M):
|
|
for j in range(N):
|
|
A[i, j] = max(A[i, j], 0.5)
|
|
|
|
A = torch.randn(M, N, dtype=torch.float32).cuda()
|
|
ref_A = A.clone()
|
|
ref_program(ref_A)
|
|
kernel(A)
|
|
torch.testing.assert_close(A, ref_A, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_tile_atomic_max():
|
|
run_tile_atomic_max(8, 128, 128, 32, 32)
|
|
|
|
|
|
def test_tile_atomic_min():
|
|
run_tile_atomic_min(8, 128, 128, 32, 32)
|
|
|
|
|
|
def test_tile_atomic_max_expr():
|
|
run_tile_atomic_max_expr(128, 128, 32, 32)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|