submission 757537

# EVOLVE-BLOCK-START

import torch
import triton
import triton.language as tl
from typing import Tuple

# Enable reduced‑precision reduction for FP16 matmul on Hopper GPUs.
# This can give a noticeable speed‑up while staying within the 0.01 tolerance.
if hasattr(torch.backends.cuda, "matmul"):
    try:
        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True
    except Exception:
        pass

# ------------------------------------------------------------------
# Simple tiled FP16 → FP32 → FP16 Triton kernel (kept for completeness).
# It is not used in the default path, but provides a ready‑to‑use
# Triton implementation should the user wish to experiment with it.
# ------------------------------------------------------------------
@triton.autotune(
    configs=[
        triton.Config(
            {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
            num_warps=4,
            num_stages=3,
        ),
        triton.Config(
            {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64},
            num_warps=8,
            num_stages=4,
        ),
        triton.Config(
            {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64},
            num_warps=8,
            num_stages=5,
        ),
        triton.Config(
            {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 128},
            num_warps=16,
            num_stages=5,
        ),
    ],
    key=["M", "N", "K"],
)
@triton.jit
def triton_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,
):
    pid_m = tl.program_id(0)  # tile row
    pid_n = tl.program_id(1)  # tile column

    # Tile start indices
    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)

    # Bounds masks
    mask_m = offs_m < M
    mask_n = offs_n < N

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

    # Loop over K dimension in tiles
    for k in range(0, K, BLOCK_SIZE_K):
        offs_k = k + tl.arange(0, BLOCK_SIZE_K)
        mask_k = offs_k < K

        # Load A tile (M × K)
        a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
        a = tl.load(a_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0)

        # Load B tile (K × N)
        b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)
        b = tl.load(b_ptrs, mask=mask_k[:, None] & mask_n[None, :], other=0.0)

        # Matrix‑multiply‑accumulate
        acc = tl.dot(a, b, acc)

    # Write result back to C
    c_ptrs = c_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
    c_mask = mask_m[:, None] & mask_n[None, :]
    tl.store(c_ptrs, acc.to(tl.float16), mask=c_mask)


def custom_kernel(data: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]) -> torch.Tensor:
    """
    Compute C = A @ B.
    For all problem sizes we rely on cuBLAS via torch.mm, which on Hopper
    GPUs benefits from the reduced‑precision reduction flag enabled above.
    The Triton kernel is retained for future experimentation.
    """
    a, b, c = data
    # The inputs are assumed to be CUDA, contiguous, float16 as per the spec.
    torch.mm(a, b, out=c)
    return c
# EVOLVE-BLOCK-END