mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-26 07:22:54 +08:00
Lei Wang
aa0b1090c9
* Add atomic operations for CUDA templates in new atomic.h file - Introduced atomic functions including AtomicMax, AtomicMin, AtomicAdd, and their return variants for various data types. - Implemented support for half, bfloat16, and float types with appropriate memory ordering. - Moved atomic-related utilities from common.h to the new atomic.h file for better organization. - Added Python bindings for atomic operations in tilelang, including atomic_max, atomic_min, atomic_add, and their vectorized counterparts. - Updated customize.py to utilize the new atomic functions, enhancing modularity and maintainability. * Refactor atomic operations in CUDA templates for improved readability - Reformatted atomic operation implementations in atomic.h for better code clarity. - Adjusted function signatures in tilelang's atomic.py to enhance readability by aligning parameters. - Cleaned up unnecessary whitespace and comments in customize.py to streamline the codebase. * Add thread storage synchronization configuration option - Introduced a new configuration option `tl.disable_thread_storage_sync` to control the automatic insertion of thread synchronization barriers in shared memory access. - Updated the `ThreadSync` pass to check this configuration and bypass synchronization if disabled. - Enhanced documentation in `builtin.h` and `pass_config.py` to clarify the purpose and usage of the new option. * Refactor thread storage sync configuration retrieval - Simplified the retrieval of the thread storage sync configuration in the `ThreadSync` pass by removing unnecessary intermediate variables. - Ensured that the inclusion of `builtin.h` is consistent by moving it to the appropriate location in the file. * test fix * Update atomic operations and tests for improved functionality - Updated atomic operations in CUDA templates to remove unnecessary address_of calls, enhancing performance and readability. - Refactored atomic operation signatures in tilelang's atomic.py to accept references instead of pointers. - Added new atomic operations and corresponding test cases for atomic add, max, min, and load/store functionalities in the testing suite. - Updated the TVM subproject to the latest commit for better compatibility. * Update attention sink examples to use 32 heads - Modified the `heads` parameter in both `example_gqa_sink_fwd_bhsd_wgmma_pipelined.py` and `example_mha_sink_fwd_bhsd_wgmma_pipelined.py` from 1 to 32 to enhance performance in attention mechanisms. - Ensured consistency across example scripts for improved usability and testing. * Refactor atomic add handling in vectorization - Simplified the extraction of buffer loads for atomic add operations by removing unnecessary address_of calls, improving code clarity and performance. - Updated the data type retrieval for vectorization size calculation to directly access the buffer load node, enhancing efficiency. * Add loop break functionality and enhance thread synchronization - Introduced a new `loop_break` function in `customize.py` to allow breaking out of loops, returning a call to the `tl.loop_break` intrinsic. - Updated the `sync_threads` function in `builtin.py` to accept optional parameters for `barrier_id` and `arrive_count`, improving its flexibility for thread synchronization. - Added necessary imports in `__init__.py` to include the new `loop_break` function for broader accessibility. * test fix
382 lines
13 KiB
Python
382 lines
13 KiB
Python
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_add_program(K, M, N, block_M, block_N, dtype="float"):
|
|
|
|
@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="float32"):
|
|
kernel = atomic_add_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 tile_atomic_add_program(K, M, N, block_M, block_N, dtype="float"):
|
|
|
|
@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="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="float"):
|
|
|
|
@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="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="float"):
|
|
|
|
@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="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="float"):
|
|
|
|
@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="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="float"):
|
|
|
|
@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="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):
|
|
|
|
@T.prim_func
|
|
def atomic_addx2(A: T.Tensor((M, N), "float16"), B: T.Tensor((M, N), "float16")):
|
|
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):
|
|
kernel = atomic_addx2_program(M, N, block_M, block_N)
|
|
import torch
|
|
|
|
A = torch.randn(M, N, dtype=torch.float16).cuda()
|
|
B = torch.zeros(M, N, dtype=torch.float16).cuda()
|
|
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_different_memory_orders_program(M, N, block_M, block_N, dtype="float"):
|
|
|
|
@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="relaxed")
|
|
T.atomic_max(C[idx_i, idx_j], val, memory_order="acquire")
|
|
T.atomic_min(D[idx_i, idx_j], val, memory_order="release")
|
|
|
|
return atomic_different_orders
|
|
|
|
|
|
def run_atomic_different_memory_orders(M, N, block_M, block_N, dtype="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))
|
|
|
|
|
|
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)
|
|
|
|
|
|
def test_atomic_load_store():
|
|
run_atomic_load_store(64, 64, 16, 16)
|
|
|
|
|
|
def test_atomic_memory_order():
|
|
run_atomic_memory_order(4, 64, 64, 16, 16)
|
|
|
|
|
|
def test_atomic_addx2():
|
|
run_atomic_addx2(32, 64, 8, 16)
|
|
|
|
|
|
@tilelang.jit
|
|
def atomic_addx4_program(M, N, block_M, block_N):
|
|
|
|
@T.prim_func
|
|
def atomic_addx4(A: T.Tensor((M, N), "float32"), B: T.Tensor((M, N), "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="float"):
|
|
|
|
@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="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)
|
|
|
|
|
|
def test_atomic_different_memory_orders():
|
|
run_atomic_different_memory_orders(32, 32, 8, 8)
|
|
|
|
|
|
def test_atomic_addx4():
|
|
run_atomic_addx4(16, 64, 4, 4)
|
|
|
|
|
|
def test_atomic_return_prev():
|
|
run_atomic_return_prev(32, 32, 8, 8)
|
|
|
|
|
|
# TODO(lei): test failed and this is experimental
|
|
# CC @dyq
|
|
# def test_tile_atomic_add():
|
|
# run_tile_atomic_add(8, 128, 128, 32, 32)
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|