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

# Define the Triton kernel for fp8-blockwise matrix multiplication
@triton.jit
def fp8_grouped_gemm(
    A_ptr, B_ptr, C_ptr,               # Pointers to input matrices A, B, and output matrix C
    lhs_scale_ptr, rhs_scale_ptr,      # Pointers to scaling factors for A and B
    M, N, K,                           # Dimensions of the matrices: A (M x K), B (N x K), C (M x N)
    stride_am, stride_ak,              # Strides for matrix A
    stride_bn, stride_bk,              # Strides for matrix B
    stride_cm, stride_cn,              # Strides for matrix C
    stride_lhs_k,                      # Stride for lhs scaling factor along K dimension
    stride_rhs_nk, stride_rhs_k,       # Strides for rhs scaling factor along N and K dimensions
    BLOCK_M: tl.constexpr,             # Block size along M dimension
    BLOCK_N: tl.constexpr,             # Block size along N dimension
    BLOCK_K: tl.constexpr,             # Block size along K dimension
    GROUP_N: tl.constexpr,             # Grouping size for N dimension in scaling factors
    GROUP_K: tl.constexpr              # Grouping size for K dimension in scaling factors
):
    """
    Triton kernel for an optimized fp8-blockwise matrix multiplication (GEMM) with grouped scaling factors.
    Computes C = A @ B * lhs_scale * rhs_scale, where A and B are fp8 matrices, and scales are applied group-wise.
    """
    # Compute block indices based on program ID
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    # Define offsets for the block
    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)

    # Calculate memory pointers for A and B with strides
    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)

    # Initialize accumulator in fp32 for higher precision
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # Loop over K dimension in blocks
    for k in range(0, K, BLOCK_K):
        # Load blocks of A and B with boundary masking
        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)

        # Compute group index for scaling factors
        group_k = k // GROUP_K

        # Load scaling factors with appropriate masking
        lhs_scale = tl.load(
            lhs_scale_ptr + offs_m[:, None] * stride_am + 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 + group_k * stride_rhs_k,
            mask=(offs_n[None, :] < N),
            other=1.0
        )

        # Perform matrix multiplication and apply scaling factors
        acc += tl.dot(a, b) * lhs_scale * rhs_scale

        # Update pointers for the next block along K
        A += BLOCK_K * stride_ak
        B += BLOCK_K * stride_bk

    # Calculate output pointer and store the result as 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))

# Wrapper function to launch the kernel
def custom_kernel(data: tuple) -> torch.Tensor:
    """
    Launches the fp8_grouped_gemm kernel with optimized parameters for MI300 hardware.
    
    Args:
        data: Tuple containing (a, b, a_scale, b_scale, c), where:
              - a: Input matrix A (M x K, fp8)
              - b: Input matrix B (N x K, fp8)
              - a_scale: Scaling factors for A
              - b_scale: Scaling factors for B
              - c: Output matrix C (M x N, bfloat16)
    
    Returns:
        c: The computed output matrix
    """
    # Unpack input data
    a, b, a_scale, b_scale, c = data

    # Extract matrix dimensions
    m, k = a.shape
    n = b.shape[0]

    # Optimized block and group sizes
    BLOCK_M = 64    # Increased from 32 for better compute intensity
    BLOCK_N = 64    # Adjusted from 128 for balanced resource use
    BLOCK_K = 128   # Kept large to reduce memory accesses
    GROUP_N = 128   # Grouping size for N dimension scaling
    GROUP_K = 128   # Grouping size for K dimension scaling

    # Define 2D grid for kernel launch
    grid = (triton.cdiv(m, BLOCK_M), triton.cdiv(n, BLOCK_N))

    # Launch the kernel
    fp8_grouped_gemm[grid](
        a, b, c,                  # Matrix pointers
        a_scale, b_scale,         # Scaling factor pointers
        m, n, k,                  # Dimensions
        a.stride(0), a.stride(1), # Strides for A
        b.stride(0), b.stride(1), # Strides for B
        c.stride(0), c.stride(1), # Strides for C
        a_scale.stride(1),        # Stride for lhs scaling
        b_scale.stride(0), b_scale.stride(1),  # Strides for rhs scaling
        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 test /tmp/tmpp8j95pt3