mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-11-01 08:23:24 +08:00
DayuxiaoshuiandLeiWang1999
869f021b2b
* [Feature] Support region as input of T.cumsum - Extend T.cumsum to accept BufferRegion and BufferLoad inputs in addition to Buffer - This enables operations on buffer slices/regions like: T.cumsum(InputG_fragment[i * chunk_size:(i + 1) * chunk_size], dim=0) - Update cumsum_fragment to handle region inputs properly - Add comprehensive tests for 1D and 2D region inputs including normal and reverse modes Fixes #879 * Fix formatting and add docstring for cumsum_fragment - Add comprehensive docstring for cumsum_fragment function - Format code according to ruff style guidelines * Fix CodeRabbit review issues - Fix negative dimension bounds check (dim < -len(shape) instead of dim <= -len(shape)) - Add src/dst shape compatibility validation for out-of-place cumsum - Update copy() type annotation to accept BufferRegion as dst parameter - Fix test in-place mutation issues by using out-of-place cumsum operations - Add non-divisible size test cases for tail region coverage * Fix out-of-bounds access in region tests - Add bounds clamping using T.min() for chunk_end calculations - Prevents accessing beyond tensor bounds for non-divisible sizes - Matches reference implementation behavior - Fixes both 1D and 2D region test cases * Fix region test: use simple slice expressions instead of T.min() - Remove T.min() which cannot be used directly in slice indices - Use chunk_start + chunk_size form instead - Rely on system's automatic bounds checking for non-divisible sizes - Update comments to reflect this approach * Fix cumsum region: use region extents in lowering and update tests for shared memory * Simplify fragment scope check using is_fragment() --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
313 lines
11 KiB
Python
313 lines
11 KiB
Python
from tilelang import tvm as tvm
|
|
import tilelang.testing
|
|
import tilelang as tl
|
|
import torch
|
|
|
|
|
|
def cumsum_smem_test(M, N, block_M, block_N, dim=0, reverse=False, dtype="float32"):
|
|
import tilelang.language as T
|
|
|
|
@T.prim_func
|
|
def cumsum(
|
|
A: T.Tensor((M, N), dtype),
|
|
B: T.Tensor((M, N), dtype),
|
|
):
|
|
# Initialize Kernel Context
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=256) as (bx, by):
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.copy(A[by * block_M, bx * block_N], A_shared)
|
|
T.cumsum(src=A_shared, dim=dim, reverse=reverse)
|
|
T.copy(A_shared, B[by * block_M, bx * block_N])
|
|
|
|
return cumsum
|
|
|
|
|
|
def cumsum_fragment_test(M, N, block_M, block_N, dim=0, reverse=False, dtype="float32"):
|
|
import tilelang.language as T
|
|
|
|
@T.prim_func
|
|
def cumsum(
|
|
A: T.Tensor((M, N), dtype),
|
|
B: T.Tensor((M, N), dtype),
|
|
):
|
|
# Initialize Kernel Context
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=256) as (bx, by):
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
A_fragment = T.alloc_fragment((block_M, block_N), dtype)
|
|
|
|
T.copy(A[by * block_M, bx * block_N], A_shared)
|
|
T.copy(A_shared, A_fragment)
|
|
T.cumsum(src=A_fragment, dim=dim, reverse=reverse)
|
|
T.copy(A_fragment, B[by * block_M, bx * block_N])
|
|
|
|
return cumsum
|
|
|
|
|
|
def run_cumsum(M, N, block_M, block_N, dim=0, reverse=False, dtype="float32", scope="smem"):
|
|
if scope == "smem":
|
|
program = cumsum_smem_test(M, N, block_M, block_N, dim, reverse, dtype)
|
|
elif scope == "fragment":
|
|
program = cumsum_fragment_test(M, N, block_M, block_N, dim, reverse, dtype)
|
|
jit_kernel = tl.compile(program, out_idx=-1)
|
|
|
|
A = torch.randn(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
|
|
def ref_program(A):
|
|
ref_b = torch.empty_like(A)
|
|
for i in range(M // block_M):
|
|
for j in range(N // block_N):
|
|
ref_b[i * block_M : (i + 1) * block_M, j * block_N : (j + 1) * block_N] = A[
|
|
i * block_M : (i + 1) * block_M, j * block_N : (j + 1) * block_N
|
|
].cumsum(dim=dim)
|
|
if reverse:
|
|
ref_b[i * block_M : (i + 1) * block_M, j * block_N : (j + 1) * block_N] = (
|
|
A[i * block_M : (i + 1) * block_M, j * block_N : (j + 1) * block_N]
|
|
.flip(dims=[dim])
|
|
.cumsum(dim=dim)
|
|
.flip(dims=[dim])
|
|
)
|
|
return ref_b
|
|
|
|
tilelang_res = jit_kernel(A)
|
|
ref_res = ref_program(A)
|
|
torch.testing.assert_close(tilelang_res, ref_res, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def cumsum_smem_test_1d(N, block_N, reverse=False, dtype="float32"):
|
|
import tilelang.language as T
|
|
|
|
@T.prim_func
|
|
def cumsum(
|
|
A: T.Tensor((N,), dtype),
|
|
B: T.Tensor((N,), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), threads=block_N) as bx:
|
|
A_shared = T.alloc_shared((block_N,), dtype)
|
|
|
|
T.copy(A[bx * block_N], A_shared)
|
|
T.cumsum(src=A_shared, dim=0, reverse=reverse)
|
|
T.copy(A_shared, B[bx * block_N])
|
|
|
|
return cumsum
|
|
|
|
|
|
def cumsum_fragment_test_1d(N, block_N, reverse=False, dtype="float32"):
|
|
import tilelang.language as T
|
|
|
|
@T.prim_func
|
|
def cumsum(
|
|
A: T.Tensor((N,), dtype),
|
|
B: T.Tensor((N,), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), threads=block_N) as bx:
|
|
A_shared = T.alloc_shared((block_N,), dtype)
|
|
A_fragment = T.alloc_fragment((block_N,), dtype)
|
|
|
|
T.copy(A[bx * block_N], A_shared)
|
|
T.copy(A_shared, A_fragment)
|
|
T.cumsum(src=A_fragment, dim=0, reverse=reverse)
|
|
T.copy(A_fragment, B[bx * block_N])
|
|
|
|
return cumsum
|
|
|
|
|
|
def run_cumsum_1d(N, block_N, reverse=False, dtype="float32", scope="smem"):
|
|
if scope == "smem":
|
|
program = cumsum_smem_test_1d(N, block_N, reverse, dtype)
|
|
elif scope == "fragment":
|
|
program = cumsum_fragment_test_1d(N, block_N, reverse, dtype)
|
|
else:
|
|
raise ValueError(f"Unknown scope {scope}")
|
|
|
|
jit_kernel = tl.compile(program, out_idx=-1)
|
|
A = torch.randn(N, dtype=getattr(torch, dtype)).cuda()
|
|
|
|
def ref_program(A):
|
|
ref_b = torch.empty_like(A)
|
|
num_blocks = (N + block_N - 1) // block_N
|
|
for j in range(num_blocks):
|
|
start = j * block_N
|
|
end = min(start + block_N, N)
|
|
chunk = A[start:end]
|
|
if reverse:
|
|
chunk = torch.flip(chunk, dims=[0])
|
|
chunk = chunk.cumsum(dim=0)
|
|
if reverse:
|
|
chunk = torch.flip(chunk, dims=[0])
|
|
ref_b[start:end] = chunk
|
|
return ref_b
|
|
|
|
tilelang_res = jit_kernel(A)
|
|
ref_res = ref_program(A)
|
|
torch.testing.assert_close(tilelang_res, ref_res, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_cumsum_smem():
|
|
# Test different sizes
|
|
run_cumsum(1024, 1024, 128, 128)
|
|
run_cumsum(1024, 1024, 128, 128, dim=1)
|
|
run_cumsum(1024, 1024, 128, 128, dim=1, reverse=True)
|
|
|
|
# Test different dtypes
|
|
run_cumsum(256, 256, 128, 128, dtype="float32")
|
|
run_cumsum(256, 256, 128, 128, dtype="float32")
|
|
|
|
|
|
def test_cumsum_fragment():
|
|
run_cumsum(1024, 1024, 128, 128, scope="fragment")
|
|
run_cumsum(1024, 1024, 128, 128, dim=1, scope="fragment")
|
|
run_cumsum(1024, 1024, 128, 128, dim=1, reverse=True, scope="fragment")
|
|
|
|
# Test different dtypes
|
|
run_cumsum(256, 256, 128, 128, dtype="float32", scope="fragment")
|
|
run_cumsum(256, 256, 128, 128, dtype="float32", scope="fragment")
|
|
|
|
|
|
def test_cumsum_smem_1d():
|
|
run_cumsum_1d(1024, 128)
|
|
run_cumsum_1d(1024, 128, reverse=True)
|
|
|
|
|
|
def test_cumsum_fragment_1d():
|
|
run_cumsum_1d(1024, 128, scope="fragment")
|
|
run_cumsum_1d(1024, 128, reverse=True, scope="fragment")
|
|
|
|
|
|
def cumsum_region_test_1d(N, chunk_size, reverse=False, dtype="float32"):
|
|
"""Test cumsum with buffer region (slice) as input."""
|
|
import tilelang.language as T
|
|
|
|
@T.prim_func
|
|
def cumsum_region(
|
|
InputG_fragment: T.Tensor((N,), dtype),
|
|
OutputG_fragment: T.Tensor((N,), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, chunk_size), threads=chunk_size) as bx:
|
|
i = bx
|
|
chunk_start = i * chunk_size
|
|
# Copy region to shared memory first (cumsum only supports shared memory)
|
|
A_shared = T.alloc_shared((chunk_size,), dtype)
|
|
T.copy(InputG_fragment[chunk_start : chunk_start + chunk_size], A_shared)
|
|
# Test cumsum with region input - in-place operation on shared memory
|
|
# This demonstrates the feature: T.cumsum(region, dim=0)
|
|
T.cumsum(src=A_shared, dim=0, reverse=reverse)
|
|
# Copy result back to global memory
|
|
T.copy(A_shared, OutputG_fragment[chunk_start : chunk_start + chunk_size])
|
|
|
|
return cumsum_region
|
|
|
|
|
|
def run_cumsum_region_1d(N, chunk_size, reverse=False, dtype="float32"):
|
|
"""Run test for cumsum with region input."""
|
|
program = cumsum_region_test_1d(N, chunk_size, reverse, dtype)
|
|
jit_kernel = tl.compile(program, out_idx=-1)
|
|
A = torch.randn(N, dtype=getattr(torch, dtype)).cuda()
|
|
|
|
def ref_program(A):
|
|
ref_b = torch.empty_like(A)
|
|
num_blocks = (N + chunk_size - 1) // chunk_size
|
|
for j in range(num_blocks):
|
|
start = j * chunk_size
|
|
end = min(start + chunk_size, N)
|
|
chunk = A[start:end].clone()
|
|
if reverse:
|
|
chunk = torch.flip(chunk, dims=[0])
|
|
chunk = chunk.cumsum(dim=0)
|
|
if reverse:
|
|
chunk = torch.flip(chunk, dims=[0])
|
|
ref_b[start:end] = chunk
|
|
return ref_b
|
|
|
|
tilelang_res = jit_kernel(A)
|
|
ref_res = ref_program(A)
|
|
torch.testing.assert_close(tilelang_res, ref_res, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def cumsum_region_test_2d(M, N, block_M, block_N, dim=0, reverse=False, dtype="float32"):
|
|
"""Test cumsum with buffer region (slice) as input in 2D."""
|
|
import tilelang.language as T
|
|
|
|
@T.prim_func
|
|
def cumsum_region(
|
|
InputG_fragment: T.Tensor((M, N), dtype),
|
|
OutputG_fragment: T.Tensor((M, N), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=256) as (bx, by):
|
|
chunk_start_M = by * block_M
|
|
chunk_start_N = bx * block_N
|
|
# Copy region to shared memory first (cumsum only supports shared memory)
|
|
A_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
T.copy(
|
|
InputG_fragment[chunk_start_M : chunk_start_M + block_M, chunk_start_N : chunk_start_N + block_N],
|
|
A_shared,
|
|
)
|
|
# Test cumsum with 2D region input - in-place operation on shared memory
|
|
T.cumsum(src=A_shared, dim=dim, reverse=reverse)
|
|
# Copy result back to global memory
|
|
T.copy(
|
|
A_shared,
|
|
OutputG_fragment[chunk_start_M : chunk_start_M + block_M, chunk_start_N : chunk_start_N + block_N],
|
|
)
|
|
|
|
return cumsum_region
|
|
|
|
|
|
def run_cumsum_region_2d(M, N, block_M, block_N, dim=0, reverse=False, dtype="float32"):
|
|
"""Run test for cumsum with 2D region input."""
|
|
program = cumsum_region_test_2d(M, N, block_M, block_N, dim, reverse, dtype)
|
|
jit_kernel = tl.compile(program, out_idx=-1)
|
|
A = torch.randn(M, N, dtype=getattr(torch, dtype)).cuda()
|
|
|
|
def ref_program(A):
|
|
ref_b = torch.empty_like(A)
|
|
num_blocks_M = (M + block_M - 1) // block_M
|
|
num_blocks_N = (N + block_N - 1) // block_N
|
|
for i in range(num_blocks_M):
|
|
for j in range(num_blocks_N):
|
|
start_M = i * block_M
|
|
end_M = min(start_M + block_M, M)
|
|
start_N = j * block_N
|
|
end_N = min(start_N + block_N, N)
|
|
chunk = A[start_M:end_M, start_N:end_N].clone()
|
|
if reverse:
|
|
chunk = torch.flip(chunk, dims=[dim])
|
|
chunk = chunk.cumsum(dim=dim)
|
|
if reverse:
|
|
chunk = torch.flip(chunk, dims=[dim])
|
|
ref_b[start_M:end_M, start_N:end_N] = chunk
|
|
return ref_b
|
|
|
|
tilelang_res = jit_kernel(A)
|
|
ref_res = ref_program(A)
|
|
torch.testing.assert_close(tilelang_res, ref_res, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_cumsum_region_1d():
|
|
"""Test cumsum with 1D region input."""
|
|
# Test normal cumsum with region input
|
|
run_cumsum_region_1d(1024, 128)
|
|
# Test reverse cumsum with region input
|
|
run_cumsum_region_1d(1024, 128, reverse=True)
|
|
# Test with different chunk sizes
|
|
run_cumsum_region_1d(512, 64)
|
|
run_cumsum_region_1d(2048, 256)
|
|
# Tail coverage (non-divisible size)
|
|
run_cumsum_region_1d(1000, 128)
|
|
|
|
|
|
def test_cumsum_region_2d():
|
|
"""Test cumsum with 2D region input."""
|
|
# Test 2D cumsum along dim 0
|
|
run_cumsum_region_2d(1024, 1024, 128, 128, dim=0)
|
|
# Test 2D cumsum along dim 1
|
|
run_cumsum_region_2d(1024, 1024, 128, 128, dim=1)
|
|
# Test reverse cumsum
|
|
run_cumsum_region_2d(512, 512, 64, 64, dim=1, reverse=True)
|
|
# Tail coverage (non-divisible size)
|
|
run_cumsum_region_2d(1000, 1000, 128, 128, dim=1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|