submission 583248

import argparse
import os
import inspect

import triton
import triton.language as tl
import torch


@triton.jit
def _matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    """Fused matrix multiplication kernel: C = A @ B
    
    Single fused kernel: tiled dot-product accumulation over K dimension,
    with fp32 accumulator cast to fp16 on store.
    """
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_in_group = pid % num_pid_in_group
    pid_m = first_pid_m + (pid_in_group % group_size_m)
    pid_n = pid_in_group // group_size_m

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)

    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    m_mask = offs_m < M
    n_mask = offs_n < N
    n_full_k_blocks = K // BLOCK_SIZE_K

    for _ in range(0, n_full_k_blocks):
        a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
        b = tl.load(b_ptrs, mask=n_mask[None, :], other=0.0)
        accumulator = tl.dot(a, b, accumulator)
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

    k_remaining = K - n_full_k_blocks * BLOCK_SIZE_K
    if k_remaining > 0:
        k_mask = offs_k < k_remaining
        a = tl.load(a_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0)
        b = tl.load(b_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0)
        accumulator = tl.dot(a, b, accumulator)

    c = accumulator.to(tl.float16)

    c_ptrs = c_ptr + stride_cm * offs_m[:, None] + stride_cn * offs_n[None, :]
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)


def kernel_function(inputs):
    """Wrapper for matmul kernel: C = A @ B
    
    All computation is fused into a single Triton kernel launch.
    Wrapper only validates, allocates, and launches.
    """
    a, b, c = inputs

    M, K = a.shape
    _, N = b.shape

    a = a.contiguous()
    b = b.contiguous()

    BLOCK_SIZE_M = 128
    BLOCK_SIZE_N = 128
    BLOCK_SIZE_K = 64
    GROUP_SIZE_M = 8

    grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),)

    _matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K,
        GROUP_SIZE_M=GROUP_SIZE_M,
        num_warps=8,
        num_stages=3,
    )

    return c

def custom_kernel(input):
    sig = inspect.signature(kernel_function)
    num_params = len(sig.parameters)

    if len(input) == num_params:
        return kernel_function(*input)
    return kernel_function(input)


# Ensure deterministic cuBLAS.
if os.environ.get("CUBLAS_WORKSPACE_CONFIG", "") not in (":4096:8", ":16:8"):
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"


def _build_inputs(m: int, n: int, k: int, device: str, dtype: torch.dtype):
    a = torch.randn((m, k), device=device, dtype=dtype)
    b = torch.randn((k, n), device=device, dtype=dtype)
    c = torch.empty((m, n), device=device, dtype=dtype)
    return a, b, c


def _check_correctness(m: int, n: int, k: int, device: str, dtype: torch.dtype):
    a, b, c = _build_inputs(m, n, k, device, dtype)
    out = custom_kernel((a, b, c))
    ref = torch.matmul(a, b)
    ok = torch.allclose(out, ref, rtol=2e-2, atol=2e-2)
    max_diff = (out - ref).abs().max().item()
    print(f"check={'PASS' if ok else 'FAIL'} max_diff={max_diff:.6f}")
    return ok


def main():
    parser = argparse.ArgumentParser(description="Launch Triton matmul kernel.")
    parser.add_argument("--m", type=int, default=2048)
    parser.add_argument("--n", type=int, default=2048)
    parser.add_argument("--k", type=int, default=2048)
    parser.add_argument("--warmup", type=int, default=10)
    parser.add_argument("--iters", type=int, default=50)
    parser.add_argument("--dtype", choices=["fp16", "fp32"], default="fp16")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--check", action="store_true")
    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required to run this kernel launcher.")

    torch.manual_seed(args.seed)
    dtype = torch.float16 if args.dtype == "fp16" else torch.float32

    if args.check:
        _check_correctness(args.m, args.n, args.k, "cuda", dtype)
        return

    inputs = _build_inputs(args.m, args.n, args.k, "cuda", dtype)

    for _ in range(args.warmup):
        custom_kernel(inputs)
    torch.cuda.synchronize()

    for _ in range(args.iters):
        custom_kernel(inputs)
    torch.cuda.synchronize()

    print(
        f"Completed launches: warmup={args.warmup}, iters={args.iters}, "
        f"shape=({args.m},{args.k})x({args.k},{args.n}), dtype={args.dtype}"
    )


if __name__ == "__main__":
    main()