#!POPCORN leaderboard amd-fp8-mm
#!POPCORN gpu MI300
import torch
from task import input_t, output_t
import triton
import triton.language as tl

@triton.jit
def fp8_grouped_gemm(
    A_ptr, B_ptr, C_ptr,
    lhs_scale_ptr, rhs_scale_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bn, stride_bk,
    stride_cm, stride_cn,
    stride_lhs_k,
    stride_rhs_nk, stride_rhs_k,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    GROUP_N: tl.constexpr,
    GROUP_K: tl.constexpr
):
    # Grid indices
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    # Compute offsets
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    # Pointers to initial blocks
    A = A_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    B = B_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)

    # Accumulator in fp32
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # K-dimension loop
    for k in range(0, K, BLOCK_K):
        # Load fp8 blocks with boundary masks
        a = tl.load(A, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
        b = tl.load(B, mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)

        # Load scaling factors
        lhs_scale = tl.load(
            lhs_scale_ptr + offs_m[:, None] * stride_am + (k // GROUP_K) * stride_lhs_k,
            mask=(offs_m[:, None] < M),
            other=1.0
        )
        rhs_scale = tl.load(
            rhs_scale_ptr + (offs_n[None, :] // GROUP_N) * stride_rhs_nk + (k // GROUP_K) * stride_rhs_k,
            mask=(offs_n[None, :] < N),
            other=1.0
        )

        # Compute and accumulate
        acc += tl.dot(a, b) * lhs_scale * rhs_scale

        # Advance pointers
        A += BLOCK_K * stride_ak
        B += BLOCK_K * stride_bk

    # Store to C in bfloat16
    c = C_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
    tl.store(c, acc.to(tl.bfloat16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

def custom_kernel(data: input_t) -> output_t:
    """
    Optimized fp8 block-wise GEMM for MI300.
    Args:
        data: Tuple (a, b, a_scale, b_scale, c)
            - a: [M, K] fp8 e4m3fnuz, column-major
            - b: [N, K] fp8 e4m3fnuz, column-major
            - a_scale: [M, K//128] fp32, column-major
            - b_scale: [N//128, K//128] fp32, column-major
            - c: [M, N] bf16, row-major
    Returns:
        c: Updated output tensor
    """
    a, b, a_scale, b_scale, c = data

    m, k = a.shape
    n = b.shape[0]

    # Optimized block sizes
    BLOCK_M = 128
    BLOCK_N = 128
    BLOCK_K = 128
    GROUP_N = 128
    GROUP_K = 128

    # Launch grid
    grid = (triton.cdiv(m, BLOCK_M), triton.cdiv(n, BLOCK_N))

    # Kernel invocation
    fp8_grouped_gemm[grid](
        a, b, c,
        a_scale, b_scale,
        m, n, k,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        a_scale.stride(1),
        b_scale.stride(0), b_scale.stride(1),
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        GROUP_N=GROUP_N,
        GROUP_K=GROUP_K
    )
    return c
python eval.py benchmark /tmp/tmp5ktyrhwv