tilelang/testing/python/language/test_tilelang_language_transpose.py

"""Tests for T.transpose shared memory transpose primitive."""

import tilelang
import tilelang.testing
import tilelang.language as T
import torch


def tilelang_transpose(M, N, block_M, block_N, dtype=T.float16):
    """Kernel: read tile from A into shared, transpose in shared, write to B.

    A is (M, N), B is (M, N).
    B = A.T.T = A when block_M == M and block_N == N (single tile).
    Actually: we read A tile (block_M, block_N) into shared,
    transpose to (block_N, block_M) in shared, then write to B
    so B[bx*block_N + j, by*block_M + i] = A[by*block_M + i, bx*block_N + j]
    i.e., B = A.T
    """

    @T.prim_func
    def main(
        A: T.Tensor((M, N), dtype),
        B: T.Tensor((N, M), dtype),
    ):
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
            tile = T.alloc_shared((block_M, block_N), dtype)
            tile_T = T.alloc_shared((block_N, block_M), dtype)

            # Load from global to shared
            T.copy(
                A[by * block_M : (by + 1) * block_M, bx * block_N : (bx + 1) * block_N],
                tile,
            )
            # Transpose in shared memory
            T.transpose(tile, tile_T)
            # Store transposed tile back to global
            T.copy(
                tile_T,
                B[bx * block_N : (bx + 1) * block_N, by * block_M : (by + 1) * block_M],
            )

    return main


def run_tilelang_transpose(M=128, N=128, block_M=128, block_N=128, dtype=T.float16):
    program = tilelang_transpose(M, N, block_M, block_N, dtype)
    kernel = tilelang.compile(
        program,
        out_idx=[1],
        pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
    )
    a = torch.randn(M, N, device="cuda", dtype=getattr(torch, dtype))
    b = kernel(a)
    expected = a.T
    torch.testing.assert_close(b, expected, rtol=1e-2, atol=1e-2)
    print(f"PASS: transpose M={M}, N={N}, block_M={block_M}, block_N={block_N}")


def tilelang_transpose_square(M, block_M, dtype=T.float16):
    """Simpler test: square transpose with single tile."""

    @T.prim_func
    def main(
        A: T.Tensor((M, M), dtype),
        B: T.Tensor((M, M), dtype),
    ):
        with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(M, block_M), threads=128) as (bx, by):
            tile = T.alloc_shared((block_M, block_M), dtype)
            tile_T = T.alloc_shared((block_M, block_M), dtype)

            T.copy(
                A[by * block_M : (by + 1) * block_M, bx * block_M : (bx + 1) * block_M],
                tile,
            )
            T.transpose(tile, tile_T)
            T.copy(
                tile_T,
                B[bx * block_M : (bx + 1) * block_M, by * block_M : (by + 1) * block_M],
            )

    return main


def run_tilelang_transpose_square(M=256, block_M=128, dtype=T.float16):
    program = tilelang_transpose_square(M, block_M, dtype)
    kernel = tilelang.compile(
        program,
        out_idx=[1],
        pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
    )
    a = torch.randn(M, M, device="cuda", dtype=getattr(torch, dtype))
    b = kernel(a)
    expected = a.T
    torch.testing.assert_close(b, expected, rtol=1e-2, atol=1e-2)
    print(f"PASS: square transpose M={M}, block_M={block_M}")


def _smem_optin_bytes() -> int:
    """Per-block opt-in dynamic shared memory limit (bytes) of the current CUDA
    device, queried via cudaDevAttrMaxSharedMemoryPerBlockOptin.

    Reference per-block opt-in caps:
        sm_70 (V100)         : 96 KB
        sm_75 (T4 / Turing)  : 64 KB
        sm_80 (A100)         : 164 KB
        sm_86 / sm_89 (Ada)  : 99 KB
        sm_90 (H100)         : 228 KB
        sm_100 (B100)        : 228 KB
        sm_120 (consumer Blackwell, RTX 50-series): 99 KB
    """
    props = torch.cuda.get_device_properties(torch.cuda.current_device())
    return getattr(props, "shared_memory_per_block_optin", 0)


# Transpose test tiers. Each entry is (label, M, N, block_M, block_N).
# `T.transpose` allocates two shared tiles (`tile` + `tile_T`), each
# block_M * block_N * sizeof(dtype) bytes, so the per-block shared-memory
# requirement is 2 * block_M * block_N * elem_bytes plus minor headroom.
# Tiers are ordered smallest -> largest. The top tier matches the original
# (pre-downgrade) test footprint of 128 KB; we never auto-promote past that.
# Smaller tiers exist so consumer cards (e.g. sm_120 with a 99 KB opt-in cap)
# still exercise the codepath instead of skipping outright.
_TRANSPOSE_TIERS = [
    # label              M    N    block_M block_N    # tile pair (fp16) -> total smem
    ("small_64x64", 128, 128, 64, 64),  # 2 * 64*64*2 = 16 KB
    ("default_128x128_square", 128, 128, 128, 128),  # 2 * 128*128*2 = 64 KB
    ("default_128x128_multi_tile", 256, 256, 128, 128),  # 64 KB, 2x2 tiling
    ("wide_128x256", 128, 256, 128, 256),  # 2 * 128*256*2 = 128 KB (original top)
]


@tilelang.testing.requires_cuda
def test_tilelang_transpose():
    import warnings

    optin = _smem_optin_bytes()  # codespell:ignore
    elem_bytes = 2  # fp16
    ran: list[str] = []
    skipped: list[str] = []
    for label, M, N, block_M, block_N in _TRANSPOSE_TIERS:
        required = 2 * block_M * block_N * elem_bytes
        if required > optin:  # codespell:ignore
            skipped.append(f"{label} (needs {required} B)")
            continue
        run_tilelang_transpose(M=M, N=N, block_M=block_M, block_N=block_N)
        ran.append(label)

    # If we did not get to the top tier, emit a runtime warning so it shows
    # up in the pytest warnings summary instead of silently downgrading.
    top_label = _TRANSPOSE_TIERS[-1][0]
    if skipped and top_label not in ran:
        warnings.warn(
            f"test_tilelang_transpose downgraded on this device (shared_memory_per_block_optin={optin} B): ran {ran}; skipped {skipped}",  # codespell:ignore
            stacklevel=1,
        )

    if not ran:
        import pytest

        smallest_required = 2 * _TRANSPOSE_TIERS[0][3] * _TRANSPOSE_TIERS[0][4] * elem_bytes
        pytest.skip(f"No transpose tier fits this device (optin={optin} B, smallest tier needs {smallest_required} B)")  # codespell:ignore


@tilelang.testing.requires_cuda
def test_tilelang_transpose_square():
    run_tilelang_transpose_square(M=128, block_M=128)
    run_tilelang_transpose_square(M=256, block_M=128)
    run_tilelang_transpose_square(M=512, block_M=128)


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