trending/tilelang
1
0
Fork 0
mirror of https://github.com/tile-ai/tilelang.git synced 2026-10-26 07:22:54 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
b27fb9286a0f05f3adffd6911bb71bf54f42c4e8
tilelang/testing/python/language/test_tilelang_language_atomic.py
T
Lei Wang b27fb9286a [Feature] Atomic Reduction Operations and Vectorization Enhancement (#1676) 
* [Feature] Implement atomic reduction operations and enhance atomic add functionality

- Added support for atomic max and min operations with corresponding classes and methods.
- Introduced vectorization capabilities for atomic add operations based on data type and target architecture.
- Refactored atomic add intrinsic calls to improve clarity and consistency.
- Enhanced layout inference for atomic operations and integrated new utility functions for target architecture checks.
- Updated tests to validate the new atomic operations and their vectorization behavior.

* Refactor atomic operation code for improved readability and consistency

- Cleaned up whitespace and formatting in atomic add and reduce implementations for better code clarity.
- Enhanced comments in CUDA code generation for atomic operations to improve understanding.
- Updated function signatures and layout definitions for consistency across atomic operation files.
- Ensured proper alignment of code for better maintainability and readability.

* lint fix

* Enhance loop vectorization and layout handling across multiple operations

- Updated function signatures in atomic add, atomic reduce, copy, and fill operations to include layout map parameters for improved layout handling.
- Refactored vectorization logic to utilize layout maps, ensuring better performance and compatibility with various buffer layouts.
- Enhanced the LowerParallelLoop function to accept layout maps, facilitating more efficient loop transformations.
- Added checks for buffer contiguity in vectorization processes to ensure correctness when using layout maps.
- Updated tests to validate the new layout handling and vectorization behavior.

* Add atomic operation checks in legalize_safe_memory_access

- Introduced a new method to identify atomic operations within the legalize_safe_memory_access transformation.
- Updated the VisitStmt_ function to handle both CallExtern and atomic operations, ensuring recursive condition collection for these cases.
- Enhanced comments for clarity on the handling of atomic operations in the context of memory access legality.

* Add atomic operations for HIP code generation

- Implemented support for various atomic operations including atomic add, atomic load, atomic store, atomic max, and atomic min in the HIP code generation.
- Enhanced the handling of atomic operations to include optional memory order parameters.
- Improved code readability with added comments explaining the purpose of each atomic operation.
- Ensured consistency with existing atomic operation implementations in the codebase.

* Refactor AtomicOpBaseNode and related classes for improved clarity

- Changed the GetElemOpStatic method to a virtual GetElemOp method in AtomicOpBaseNode, enhancing polymorphism.
- Updated AtomicAddNode, AtomicMaxNode, and AtomicMinNode to override the new GetElemOp method.
- Removed unnecessary template parameters from AtomicOpBaseNode, simplifying the class structure.
- Cleaned up includes in atomic_reduce.cc to remove unused dependencies, improving code organization.

* Enhance argument validation for atomic operations

- Added checks in AtomicAdd, AtomicMax, and AtomicMin constructors to ensure at least two arguments are provided, improving error handling and user feedback.
- Removed the unused LowerTMA method declaration from the AtomicAdd class, streamlining the codebase.

* lint fix

* fix

* test fix

* performance fix

* remove useless comments
2026-01-16 00:03:01 +08:00

481 lines
17 KiB
Python
Raw Blame History

import tilelang.testing
import tilelang.layout
import tilelang.language as T
import torch
@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)
print(kernel.get_kernel_source())
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] += 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 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)
print(kernel.get_kernel_source())
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] += 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)
print(B)
print(ref_B)
torch.testing.assert_close(B, ref_B, atol=1e-3, rtol=1e-3)
@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)
@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
def ref_program(A, B):
for k in range(K):
for i in range(M):
for j in range(N):
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_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)
def test_atomic_add():
run_atomic_add(8, 128, 128, 32, 32)
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)
@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.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_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_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)
def test_atomic_return_prev():
run_atomic_return_prev(32, 32, 8, 8)
def test_tile_atomic_add():
run_tile_atomic_add(8, 128, 128, 32, 32)
# ======================= Tile-level atomic max =======================
@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)
print(kernel.get_kernel_source())
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)
def test_tile_atomic_max():
run_tile_atomic_max(8, 128, 128, 32, 32)
# ======================= Tile-level atomic min =======================
@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)
print(kernel.get_kernel_source())
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)
def test_tile_atomic_min():
run_tile_atomic_min(8, 128, 128, 32, 32)
@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()
if __name__ == "__main__":
tilelang.testing.main()
Reference in New Issue View Git Blame Copy Permalink