mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Yongqi Zhuo
fb1b11f1b8
This PR adds a new set of annotations for describing arbitrary warp specialization scheme, with the abstraction of pipelines. A pipeline supports 4 kinds of operations: producer_acquire, producer_commit, consumer_wait, and consumer_release. Behind the scenes, mbarriers are created to implement the synchronization between warps. This PR introduces a new pass, MaterializeWSSchedule, that transforms an annotated kernel to a warp-specialized kernel. The set of infrastructure enables further development of fully automatic warp specialization in TileLang. I am working on this.
716 lines
25 KiB
Python
716 lines
25 KiB
Python
from tilelang import tvm as tvm
|
|
import tilelang.testing
|
|
import tilelang as tl
|
|
import torch
|
|
import tilelang.language as T
|
|
|
|
|
|
def _torch_cummax(chunk, dim, reverse):
|
|
if reverse:
|
|
return torch.flip(torch.flip(chunk, dims=[dim]).cummax(dim=dim).values, dims=[dim])
|
|
return chunk.cummax(dim=dim).values
|
|
|
|
|
|
def cumsum_smem_test(M, N, block_M, block_N, dim=0, reverse=False, dtype=T.float32):
|
|
@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=T.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 test_fragment_scan_annotations_preserved():
|
|
@T.prim_func
|
|
def scan(A: T.Tensor((32,), T.float32)):
|
|
with T.Kernel(1, threads=32):
|
|
sum_frag = T.alloc_fragment((32,), T.float32)
|
|
max_frag = T.alloc_fragment((32,), T.float32)
|
|
T.copy(A, sum_frag)
|
|
T.copy(A, max_frag)
|
|
T.cumsum(sum_frag, annotations={"test.scan": "cumsum"})
|
|
T.cummax(max_frag, annotations={"test.scan": "cummax"})
|
|
|
|
lines = scan.script().splitlines()
|
|
cumsum_line = next(line for line in lines if "T.cumsum(" in line)
|
|
cummax_line = next(line for line in lines if "T.cummax(" in line)
|
|
assert 'test.scan="cumsum"' in cumsum_line
|
|
assert 'test.scan="cummax"' in cummax_line
|
|
|
|
|
|
def run_cumsum(M, N, block_M, block_N, dim=0, reverse=False, dtype=T.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=T.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=T.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=T.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(256, 256, 64, 64)
|
|
run_cumsum(256, 256, 64, 64, dim=1)
|
|
run_cumsum(256, 256, 64, 64, dim=1, reverse=True)
|
|
run_cumsum(192, 160, 64, 32, dim=0)
|
|
run_cumsum(192, 160, 64, 32, dim=0, reverse=True)
|
|
|
|
# Test different dtypes
|
|
run_cumsum(128, 128, 64, 64, dtype=T.float32)
|
|
|
|
|
|
def test_cumsum_fragment():
|
|
run_cumsum(256, 256, 64, 64, scope="fragment")
|
|
run_cumsum(256, 256, 64, 64, dim=1, scope="fragment")
|
|
run_cumsum(256, 256, 64, 64, dim=1, reverse=True, scope="fragment")
|
|
|
|
# Test different dtypes
|
|
run_cumsum(128, 128, 64, 64, dtype=T.float32, scope="fragment")
|
|
|
|
|
|
def test_cumsum_smem_1d():
|
|
run_cumsum_1d(512, 64)
|
|
run_cumsum_1d(512, 64, reverse=True)
|
|
|
|
|
|
def test_cumsum_fragment_1d():
|
|
run_cumsum_1d(512, 64, scope="fragment")
|
|
run_cumsum_1d(512, 64, reverse=True, scope="fragment")
|
|
|
|
|
|
def cumsum_region_test_1d(N, chunk_size, reverse=False, dtype=T.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=T.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=T.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=T.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(512, 64)
|
|
# Test reverse cumsum with region input
|
|
run_cumsum_region_1d(512, 64, reverse=True)
|
|
# Test with different chunk sizes
|
|
run_cumsum_region_1d(384, 128)
|
|
# Tail coverage (non-divisible size)
|
|
run_cumsum_region_1d(250, 64)
|
|
|
|
|
|
def test_cumsum_region_2d():
|
|
"""Test cumsum with 2D region input."""
|
|
# Test 2D cumsum along dim 0
|
|
run_cumsum_region_2d(256, 256, 64, 64, dim=0)
|
|
# Test 2D cumsum along dim 1
|
|
run_cumsum_region_2d(256, 256, 64, 64, dim=1)
|
|
# Test reverse cumsum
|
|
run_cumsum_region_2d(192, 192, 64, 64, dim=1, reverse=True)
|
|
# Tail coverage (non-divisible size)
|
|
run_cumsum_region_2d(250, 250, 64, 64, dim=1)
|
|
|
|
|
|
def cummax_smem_test(M, N, block_M, block_N, dim=0, reverse=False, dtype=T.float32):
|
|
@T.prim_func
|
|
def cummax(
|
|
A: T.Tensor((M, N), dtype),
|
|
B: T.Tensor((M, N), dtype),
|
|
):
|
|
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.cummax(src=A_shared, dim=dim, reverse=reverse)
|
|
T.copy(A_shared, B[by * block_M, bx * block_N])
|
|
|
|
return cummax
|
|
|
|
|
|
def cummax_fragment_test(M, N, block_M, block_N, dim=0, reverse=False, dtype=T.float32):
|
|
@T.prim_func
|
|
def cummax(
|
|
A: T.Tensor((M, N), dtype),
|
|
B: T.Tensor((M, N), dtype),
|
|
):
|
|
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.cummax(src=A_fragment, dim=dim, reverse=reverse)
|
|
T.copy(A_fragment, B[by * block_M, bx * block_N])
|
|
|
|
return cummax
|
|
|
|
|
|
def cummax_smem_out_test(M, N, block_M, block_N, dim=0, reverse=False, dtype=T.float32):
|
|
@T.prim_func
|
|
def cummax(
|
|
A: T.Tensor((M, N), dtype),
|
|
B: T.Tensor((M, N), dtype),
|
|
):
|
|
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)
|
|
B_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.copy(A[by * block_M, bx * block_N], A_shared)
|
|
T.cummax(src=A_shared, dst=B_shared, dim=dim, reverse=reverse)
|
|
T.copy(B_shared, B[by * block_M, bx * block_N])
|
|
|
|
return cummax
|
|
|
|
|
|
def run_cummax(M, N, block_M, block_N, dim=0, reverse=False, dtype=T.float32, scope="smem"):
|
|
if scope == "smem":
|
|
program = cummax_smem_test(M, N, block_M, block_N, dim, reverse, dtype)
|
|
elif scope == "fragment":
|
|
program = cummax_fragment_test(M, N, block_M, block_N, dim, reverse, dtype)
|
|
elif scope == "smem_out":
|
|
program = cummax_smem_out_test(M, N, block_M, block_N, dim, reverse, dtype)
|
|
else:
|
|
raise ValueError(f"Unknown scope {scope}")
|
|
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 - 1) // block_M):
|
|
for j in range((N + block_N - 1) // block_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)
|
|
ref_b[start_m:end_m, start_n:end_n] = _torch_cummax(A[start_m:end_m, start_n:end_n], dim, reverse)
|
|
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 cummax_smem_test_1d(N, block_N, reverse=False, dtype=T.float32, threads=None):
|
|
if threads is None:
|
|
threads = block_N
|
|
|
|
@T.prim_func
|
|
def cummax(
|
|
A: T.Tensor((N,), dtype),
|
|
B: T.Tensor((N,), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), threads=threads) as bx:
|
|
A_shared = T.alloc_shared((block_N,), dtype)
|
|
|
|
T.copy(A[bx * block_N], A_shared)
|
|
T.cummax(src=A_shared, dim=0, reverse=reverse)
|
|
T.copy(A_shared, B[bx * block_N])
|
|
|
|
return cummax
|
|
|
|
|
|
def cummax_fragment_test_1d(N, block_N, reverse=False, dtype=T.float32):
|
|
@T.prim_func
|
|
def cummax(
|
|
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.cummax(src=A_fragment, dim=0, reverse=reverse)
|
|
T.copy(A_fragment, B[bx * block_N])
|
|
|
|
return cummax
|
|
|
|
|
|
def run_cummax_1d(N, block_N, reverse=False, dtype=T.float32, scope="smem", negative_input=False, threads=None):
|
|
if scope == "smem":
|
|
program = cummax_smem_test_1d(N, block_N, reverse, dtype, threads)
|
|
elif scope == "fragment":
|
|
program = cummax_fragment_test_1d(N, block_N, reverse, dtype)
|
|
else:
|
|
raise ValueError(f"Unknown scope {scope}")
|
|
|
|
jit_kernel = tl.compile(program, out_idx=-1)
|
|
torch_dtype = getattr(torch, dtype)
|
|
if negative_input:
|
|
A = -torch.arange(1, N + 1, dtype=torch.float32, device="cuda").to(torch_dtype)
|
|
else:
|
|
A = torch.randn(N, dtype=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)
|
|
ref_b[start:end] = _torch_cummax(A[start:end], 0, reverse)
|
|
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_cummax_smem():
|
|
run_cummax(256, 256, 64, 64)
|
|
run_cummax(256, 256, 64, 64, dim=1)
|
|
run_cummax(256, 256, 64, 64, dim=1, reverse=True)
|
|
run_cummax(192, 160, 64, 32, dim=0)
|
|
run_cummax(192, 160, 64, 32, dim=0, reverse=True)
|
|
|
|
|
|
def test_cummax_fragment():
|
|
run_cummax(256, 256, 64, 64, scope="fragment")
|
|
run_cummax(256, 256, 64, 64, dim=1, scope="fragment")
|
|
run_cummax(256, 256, 64, 64, dim=1, reverse=True, scope="fragment")
|
|
|
|
|
|
def test_cummax_out_of_place():
|
|
run_cummax(128, 128, 64, 64, dim=1, scope="smem_out")
|
|
|
|
|
|
def test_cummax_smem_1d():
|
|
run_cummax_1d(512, 64)
|
|
run_cummax_1d(512, 64, reverse=True)
|
|
run_cummax_1d(80, 40, reverse=True, negative_input=True, threads=64)
|
|
|
|
|
|
def test_cummax_fragment_1d():
|
|
run_cummax_1d(512, 64, scope="fragment")
|
|
run_cummax_1d(512, 64, reverse=True, scope="fragment")
|
|
|
|
|
|
def cumsum_strided_region_test(M, N, NBIG, dim=0, reverse=False, dtype=T.float32):
|
|
|
|
@T.prim_func
|
|
def cumsum_strided(
|
|
A: T.Tensor((M, N), dtype),
|
|
B: T.Tensor((M, N), dtype),
|
|
):
|
|
with T.Kernel(1, threads=256) as _:
|
|
big = T.alloc_shared((M, NBIG), dtype)
|
|
for i, j in T.Parallel(M, NBIG):
|
|
big[i, j] = T.cast(0, dtype)
|
|
for i, j in T.Parallel(M, N):
|
|
big[i, j] = A[i, j]
|
|
T.cumsum(src=big[0:M, 0:N], dst=big[0:M, 0:N], dim=dim, reverse=reverse)
|
|
for i, j in T.Parallel(M, N):
|
|
B[i, j] = big[i, j]
|
|
|
|
return cumsum_strided
|
|
|
|
|
|
def run_cumsum_strided(M, N, NBIG, dim=0, reverse=False, dtype=T.float32):
|
|
program = cumsum_strided_region_test(M, N, NBIG, dim, reverse, dtype)
|
|
jit_kernel = tl.compile(program, out_idx=-1)
|
|
|
|
A = torch.randint(-2, 3, (M, N), dtype=torch.float32).cuda()
|
|
|
|
if reverse:
|
|
ref = A.flip(dims=[dim]).cumsum(dim=dim).flip(dims=[dim])
|
|
else:
|
|
ref = A.cumsum(dim=dim)
|
|
|
|
tilelang_res = jit_kernel(A)
|
|
torch.testing.assert_close(tilelang_res, ref)
|
|
|
|
|
|
def cummax_strided_region_test(M, N, NBIG, dim=0, reverse=False, dtype=T.float32):
|
|
|
|
@T.prim_func
|
|
def cummax_strided(
|
|
A: T.Tensor((M, N), dtype),
|
|
B: T.Tensor((M, N), dtype),
|
|
):
|
|
with T.Kernel(1, threads=256) as _:
|
|
big = T.alloc_shared((M, NBIG), dtype)
|
|
for i, j in T.Parallel(M, NBIG):
|
|
big[i, j] = T.cast(0, dtype)
|
|
for i, j in T.Parallel(M, N):
|
|
big[i, j] = A[i, j]
|
|
T.cummax(src=big[0:M, 0:N], dst=big[0:M, 0:N], dim=dim, reverse=reverse)
|
|
for i, j in T.Parallel(M, N):
|
|
B[i, j] = big[i, j]
|
|
|
|
return cummax_strided
|
|
|
|
|
|
def run_cummax_strided(M, N, NBIG, dim=0, reverse=False, dtype=T.float32):
|
|
program = cummax_strided_region_test(M, N, NBIG, dim, reverse, dtype)
|
|
jit_kernel = tl.compile(program, out_idx=-1)
|
|
|
|
A = torch.randint(-2, 3, (M, N), dtype=torch.float32).cuda()
|
|
|
|
if reverse:
|
|
ref = A.flip(dims=[dim]).cummax(dim=dim).values.flip(dims=[dim])
|
|
else:
|
|
ref = A.cummax(dim=dim).values
|
|
|
|
tilelang_res = jit_kernel(A)
|
|
torch.testing.assert_close(tilelang_res, ref)
|
|
|
|
|
|
def scan_strided_out_of_place_test(M, N, src_pitch, dst_pitch, op="cumsum", dim=0, reverse=False, dtype=T.float32):
|
|
|
|
@T.prim_func
|
|
def scan_strided_out_of_place(
|
|
A: T.Tensor((M, N), dtype),
|
|
B: T.Tensor((M, N), dtype),
|
|
):
|
|
with T.Kernel(1, threads=256) as _:
|
|
src_big = T.alloc_shared((M, src_pitch), dtype)
|
|
dst_big = T.alloc_shared((M, dst_pitch), dtype)
|
|
for i, j in T.Parallel(M, src_pitch):
|
|
src_big[i, j] = T.cast(0, dtype)
|
|
for i, j in T.Parallel(M, dst_pitch):
|
|
dst_big[i, j] = T.cast(0, dtype)
|
|
for i, j in T.Parallel(M, N):
|
|
src_big[i, j] = A[i, j]
|
|
scan = T.cumsum if op == "cumsum" else T.cummax
|
|
scan(
|
|
src=src_big[0:M, 0:N],
|
|
dst=dst_big[0:M, 0:N],
|
|
dim=dim,
|
|
reverse=reverse,
|
|
)
|
|
for i, j in T.Parallel(M, N):
|
|
B[i, j] = dst_big[i, j]
|
|
|
|
return scan_strided_out_of_place
|
|
|
|
|
|
def run_scan_strided_out_of_place(M, N, src_pitch, dst_pitch, op="cumsum", dim=0, reverse=False, dtype=T.float32):
|
|
program = scan_strided_out_of_place_test(M, N, src_pitch, dst_pitch, op, dim, reverse, dtype)
|
|
jit_kernel = tl.compile(program, out_idx=-1)
|
|
|
|
A = torch.arange(M * N, dtype=getattr(torch, dtype), device="cuda").reshape(M, N) - N
|
|
if op == "cumsum":
|
|
ref = A.cumsum(dim=dim)
|
|
else:
|
|
ref = A.cummax(dim=dim).values
|
|
if reverse:
|
|
flipped = A.flip(dims=[dim])
|
|
if op == "cumsum":
|
|
ref = flipped.cumsum(dim=dim).flip(dims=[dim])
|
|
else:
|
|
ref = flipped.cummax(dim=dim).values.flip(dims=[dim])
|
|
|
|
tilelang_res = jit_kernel(A)
|
|
torch.testing.assert_close(tilelang_res, ref)
|
|
|
|
|
|
def test_cumsum_strided_region():
|
|
"""cumsum over a non-contiguous 2-D shared sub-region."""
|
|
for M, N, NBIG, dim, reverse in [
|
|
(8, 40, 64, 0, False),
|
|
(8, 40, 64, 1, False),
|
|
(8, 40, 64, 1, True),
|
|
(8, 40, 64, 0, True),
|
|
]:
|
|
run_cumsum_strided(M, N, NBIG, dim, reverse)
|
|
|
|
|
|
def test_cummax_strided_region():
|
|
"""cummax over a non-contiguous 2-D shared sub-region."""
|
|
for M, N, NBIG, dim, reverse in [
|
|
(8, 40, 64, 0, False),
|
|
(8, 40, 64, 1, False),
|
|
(8, 40, 64, 0, True),
|
|
(8, 40, 64, 1, True),
|
|
]:
|
|
run_cummax_strided(M, N, NBIG, dim, reverse)
|
|
|
|
|
|
def test_scan_strided_out_of_place():
|
|
"""Out-of-place scan with distinct source and destination row pitches."""
|
|
for op in ("cumsum", "cummax"):
|
|
for dim in (0, 1):
|
|
run_scan_strided_out_of_place(8, 40, 64, 80, op=op, dim=dim, reverse=dim == 1)
|
|
|
|
|
|
def scan_offset_subregion_test(H, W, r0, r1, op="cumsum", dim=0, reverse=False, dtype=T.float32):
|
|
"""Feed a row-offset 2D sub-region of shared memory directly to the scan.
|
|
|
|
Regression for #2536: MakeAccessPtrFromRegion dropped the innermost dims'
|
|
``min`` from the access-pointer offset, so a sub-region like
|
|
``A_shared[r0:r1, :]`` with ``r0 != 0`` silently scanned rows ``[0:r1-r0]``.
|
|
"""
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((H, W), dtype),
|
|
B: T.Tensor((H, W), dtype),
|
|
):
|
|
with T.Kernel(1, threads=128):
|
|
A_shared = T.alloc_shared((H, W), dtype)
|
|
T.copy(A, A_shared)
|
|
scan = T.cumsum if op == "cumsum" else T.cummax
|
|
# Offset sub-region fed straight into the scan.
|
|
scan(src=A_shared[r0:r1, :], dim=dim, reverse=reverse)
|
|
T.copy(A_shared, B)
|
|
|
|
return main
|
|
|
|
|
|
def run_scan_offset_subregion(H, W, r0, r1, op="cumsum", dim=0, reverse=False, dtype=T.float32):
|
|
program = scan_offset_subregion_test(H, W, r0, r1, op, dim, reverse, dtype)
|
|
jit_kernel = tl.compile(program, out_idx=-1)
|
|
A = torch.randn(H, W, dtype=getattr(torch, dtype)).cuda()
|
|
|
|
def ref_program(A):
|
|
ref_b = A.clone() # rows outside [r0:r1] must be passed through untouched
|
|
chunk = A[r0:r1, :]
|
|
if op == "cumsum":
|
|
if reverse:
|
|
chunk = chunk.flip(dims=[dim]).cumsum(dim=dim).flip(dims=[dim])
|
|
else:
|
|
chunk = chunk.cumsum(dim=dim)
|
|
else:
|
|
chunk = _torch_cummax(chunk, dim, reverse)
|
|
ref_b[r0:r1, :] = 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_scan_offset_subregion():
|
|
"""Regression for #2536: row-offset 2D shared sub-regions fed to the scan."""
|
|
H, W = 128, 8
|
|
for op in ("cumsum", "cummax"):
|
|
for dim in (0, 1):
|
|
for reverse in (False, True):
|
|
# r0 == 64 is the regressing case (r0 == 0 is already covered by
|
|
# the full-region tests above).
|
|
run_scan_offset_subregion(H, W, 64, 128, op=op, dim=dim, reverse=reverse)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|