tilelang/testing/python/language/test_tilelang_language_scan.py

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 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 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 region tests above).
                run_scan_offset_subregion(H, W, 64, 128, op=op, dim=dim, reverse=reverse)


if __name__ == "__main__":
    tilelang.testing.main()