tilelang/testing/python/kernel/test_tilelang_kernel_gemm_sliced_operand.py

"""A shared operand that is a slice of a wider buffer must match the same GEMM
written with one buffer per slice.

Regression coverage for the WGMMA K-panel stride: a K-major operand whose
``block_K`` spans more than one swizzle atom is stored as (MN, (atom, panels)),
and stepping from one K panel to the next moves by the *buffer's* MN extent. The
offset formulas used to reconstruct that step from the operand's own extent,
which is smaller for a slice, so ``T.gemm(A_s, B_s[0:64, :], ...)`` silently read
the wrong panel for every ``ki`` past the first atom. ``block_K == 64`` (bf16 =
one 128B atom, single panel) never exercised the term and stayed correct.
"""

import pytest
import torch

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

ACCUM = T.float32


def _matmul_sliced(M, N, K, block_M, block_N, block_K, dtype, split, operand, transposed, sliced):
    """Split one GEMM into two half-width ones over a shared operand.

    ``operand`` selects which operand is split ("A" splits M, "B" splits N),
    ``transposed`` gives that operand's ``transpose_`` flag (which decides whether
    the split axis is the buffer's leading or trailing one), and ``sliced`` picks
    the shape under test (one buffer consumed as two slices) versus the reference
    (one buffer per half).
    """
    trans_a = operand == "A" and transposed
    trans_b = operand == "B" and transposed
    a_shape = (K, M) if trans_a else (M, K)
    b_shape = (N, K) if trans_b else (K, N)
    split_full = block_N if operand == "B" else block_M
    rest = split_full - split

    def a_tile(extent):
        return (block_K, extent) if trans_a else (extent, block_K)

    def b_tile(extent):
        return (extent, block_K) if trans_b else (block_K, extent)

    @T.prim_func
    def main(
        A: T.Tensor(a_shape, dtype),
        B: T.Tensor(b_shape, dtype),
        C: T.Tensor((M, N), dtype),
    ):
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
            C0 = T.alloc_fragment((split, block_N) if operand == "A" else (block_M, split), ACCUM)
            C1 = T.alloc_fragment((rest, block_N) if operand == "A" else (block_M, rest), ACCUM)
            T.clear(C0)
            T.clear(C1)

            if operand == "B":
                A_s = T.alloc_shared(a_tile(block_M), dtype)
                if sliced:
                    B_s = T.alloc_shared(b_tile(block_N), dtype)
                else:
                    B_s0 = T.alloc_shared(b_tile(split), dtype)
                    B_s1 = T.alloc_shared(b_tile(rest), dtype)
            else:
                B_s = T.alloc_shared(b_tile(block_N), dtype)
                if sliced:
                    A_s = T.alloc_shared(a_tile(block_M), dtype)
                else:
                    A_s0 = T.alloc_shared(a_tile(split), dtype)
                    A_s1 = T.alloc_shared(a_tile(rest), dtype)

            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=0):
                if operand == "B":
                    if trans_a:
                        T.copy(A[ko * block_K, by * block_M], A_s)
                    else:
                        T.copy(A[by * block_M, ko * block_K], A_s)
                    if sliced:
                        if trans_b:
                            T.copy(B[bx * block_N, ko * block_K], B_s)
                            T.gemm(A_s, B_s[0:split, :], C0, trans_a, trans_b)
                            T.gemm(A_s, B_s[split:block_N, :], C1, trans_a, trans_b)
                        else:
                            T.copy(B[ko * block_K, bx * block_N], B_s)
                            T.gemm(A_s, B_s[:, 0:split], C0, trans_a, trans_b)
                            T.gemm(A_s, B_s[:, split:block_N], C1, trans_a, trans_b)
                    else:
                        if trans_b:
                            T.copy(B[bx * block_N, ko * block_K], B_s0)
                            T.copy(B[bx * block_N + split, ko * block_K], B_s1)
                        else:
                            T.copy(B[ko * block_K, bx * block_N], B_s0)
                            T.copy(B[ko * block_K, bx * block_N + split], B_s1)
                        T.gemm(A_s, B_s0, C0, trans_a, trans_b)
                        T.gemm(A_s, B_s1, C1, trans_a, trans_b)
                else:
                    if trans_b:
                        T.copy(B[bx * block_N, ko * block_K], B_s)
                    else:
                        T.copy(B[ko * block_K, bx * block_N], B_s)
                    if sliced:
                        if trans_a:
                            T.copy(A[ko * block_K, by * block_M], A_s)
                            T.gemm(A_s[:, 0:split], B_s, C0, trans_a, trans_b)
                            T.gemm(A_s[:, split:block_M], B_s, C1, trans_a, trans_b)
                        else:
                            T.copy(A[by * block_M, ko * block_K], A_s)
                            T.gemm(A_s[0:split, :], B_s, C0, trans_a, trans_b)
                            T.gemm(A_s[split:block_M, :], B_s, C1, trans_a, trans_b)
                    else:
                        if trans_a:
                            T.copy(A[ko * block_K, by * block_M], A_s0)
                            T.copy(A[ko * block_K, by * block_M + split], A_s1)
                        else:
                            T.copy(A[by * block_M, ko * block_K], A_s0)
                            T.copy(A[by * block_M + split, ko * block_K], A_s1)
                        T.gemm(A_s0, B_s, C0, trans_a, trans_b)
                        T.gemm(A_s1, B_s, C1, trans_a, trans_b)

            if operand == "B":
                T.copy(C0, C[by * block_M, bx * block_N])
                T.copy(C1, C[by * block_M, bx * block_N + split])
            else:
                T.copy(C0, C[by * block_M, bx * block_N])
                T.copy(C1, C[by * block_M + split, bx * block_N])

    return main


def run_sliced_operand(M, N, K, block_M, block_N, block_K, dtype, split, operand, transposed):
    torch_dtype = getattr(torch, dtype)
    trans_a = operand == "A" and transposed
    trans_b = operand == "B" and transposed

    torch.manual_seed(0)
    a = torch.randn((K, M) if trans_a else (M, K), device="cuda", dtype=torch_dtype)
    b = torch.randn((N, K) if trans_b else (K, N), device="cuda", dtype=torch_dtype)
    ref = (a.T if trans_a else a).float() @ (b.T if trans_b else b).float()

    outputs = {}
    for sliced in (False, True):
        kernel = tilelang.compile(
            _matmul_sliced(M, N, K, block_M, block_N, block_K, dtype, split, operand, transposed, sliced),
            out_idx=[2],
            target="cuda",
            pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
        )
        outputs[sliced] = kernel(a, b).float()

    # The two shapes issue the same instructions over the same values, so they
    # agree bit-for-bit; the fp32 reference only bounds the accumulation error.
    torch.testing.assert_close(outputs[True], outputs[False], rtol=0, atol=0)
    torch.testing.assert_close(outputs[True], ref, rtol=1e-2, atol=1)


@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
@pytest.mark.parametrize("operand", ["A", "B"])
@pytest.mark.parametrize("transposed", [False, True])
@pytest.mark.parametrize("block_K", [64, 128, 256])
def test_gemm_sliced_shared_operand(operand, transposed, block_K):
    """block_K 128/256 put more than one swizzle atom on the K axis."""
    run_sliced_operand(
        M=256,
        N=256,
        K=512,
        block_M=128,
        block_N=128,
        block_K=block_K,
        dtype="bfloat16",
        split=64,
        operand=operand,
        transposed=transposed,
    )


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