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

HARDCODED_CONFIGS = {
    # Config [1024x1536x7168]: 128-128-64-4-8-2
    (1024, 1536, 7168): (128, 128, 64, 4, 8, 2),
    # Config [1024x3072x1536]: 128-128-64-4-8-2
    (1024, 3072, 1536): (128, 128, 64, 4, 8, 2),
    # Config [1024x576x7168]: 128-128-64-4-8-2
    (1024, 576, 7168): (128, 128, 64, 4, 8, 2),
    # Config [1024x7168x256]: 256-128-64-2-8-1
    (1024, 7168, 256): (256, 128, 64, 2, 8, 1),
    # Config [1024x7168x2048]: 256-128-64-2-8-1
    (1024, 7168, 2048): (256, 128, 64, 2, 8, 1),
    # Config [1024x4608x7168]: 128-128-64-4-8-2
    (1024, 4608, 7168): (128, 128, 64, 4, 8, 2),
    # Config [1024x7168x2304]: 256-128-64-2-8-1
    (1024, 7168, 2304): (256, 128, 64, 2, 8, 1),
    # Config [1024x512x7168]: 128-128-64-4-8-2
    (1024, 512, 7168): (128, 128, 64, 4, 8, 2),
    # Config [1024x4096x512]: 128-128-64-4-8-2
    (1024, 4096, 512): (128, 128, 64, 4, 8, 2),
    # Config [6144x1536x7168]: 256-128-64-4-8-1
    (6144, 1536, 7168): (256, 128, 64, 4, 8, 1),
    # Config [6144x3072x1536]: 256-128-64-2-8-1
    (6144, 3072, 1536): (256, 128, 64, 2, 8, 1),
    # Config [6144x576x7168]: 128-128-128-4-4-1
    (6144, 576, 7168): (128, 128, 128, 4, 4, 1),
    # Config [6144x7168x256]: 256-128-64-2-8-1
    (6144, 7168, 256): (256, 128, 64, 2, 8, 1),
    # Config [6144x7168x2048]: 256-128-64-2-8-1
    (6144, 7168, 2048): (256, 128, 64, 2, 8, 1),
    # Config [6144x4608x7168]: 256-128-64-4-8-1
    (6144, 4608, 7168): (256, 128, 64, 4, 8, 1),
    # Config [6144x7168x2304]: 256-128-64-2-8-1
    (6144, 7168, 2304): (256, 128, 64, 2, 8, 1),
    # Config [6144x512x7168]: 128-128-64-4-8-2
    (6144, 512, 7168): (128, 128, 64, 4, 8, 2),
    # Config [6144x4096x512]: 256-128-64-2-8-1
    (6144, 4096, 512): (256, 128, 64, 2, 8, 1),
}
DEFAULT_CONFIG = (256, 128, 64, 4, 8, 1)
# --- End Hardcoded Configurations ---


@triton.jit
def matmul_fp8_kernel(
    a_ptr, b_ptr, a_scale_ptr, b_scale_ptr, c_ptr, # Matrix dimensions
    M, N, K,
    stride_am, stride_ak,   # A (M x K): stride_am=K, stride_ak=1
    stride_bn, stride_bk,   # B (N x K): stride_bn=K, stride_bk=1
    stride_a_scale_m, stride_a_scale_k, # A scales (M x K/BLK_K_S): stride_a_scale_m=K/BLK_K_S, stride_a_scale_k=1
    stride_b_scale_n, stride_b_scale_k, # B scales (N/BLK_N_S x K/BLK_K_S): stride_b_scale_n=K/BLK_K_S, stride_b_scale_k=1
    stride_cm, stride_cn,   # C (M x N): stride_cm=N, stride_cn=1
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
    BLOCK_SIZE_K_SCALE: tl.constexpr, BLOCK_SIZE_N_SCALE: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    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 = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_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) # Corrected: Used later for B and C pointers

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

    # Calculate total number of K_SCALE blocks
    num_k_scale_blocks = tl.cdiv(K, BLOCK_SIZE_K_SCALE)

    # Outer loop iterates over K_SCALE blocks
    for k_scale_idx in range(0, num_k_scale_blocks):
        k_scale_offset = k_scale_idx * BLOCK_SIZE_K_SCALE

        # --- Load scales ONCE per k_scale_idx ---
        # Load a_scale_vec
        a_scale_k_idx = k_scale_idx # Scale K index matches scale block index
        a_scale_ptrs = a_scale_ptr + offs_m[:, None] * stride_a_scale_m + a_scale_k_idx * stride_a_scale_k
        # Masking for M dimension and K_scale dimension
        a_scale_mask = (offs_m[:, None] < M) # K check implicit via loop bound? Let's keep it for safety: & (a_scale_k_idx < tl.cdiv(K, BLOCK_SIZE_K_SCALE))
        a_scale_vec = tl.load(a_scale_ptrs, mask=a_scale_mask, other=1.0)

        # Load b_scale
        # Calculate the N block index within the scaled N dimension
        b_scale_n_block_idx = (pid_n * BLOCK_SIZE_N) // BLOCK_SIZE_N_SCALE
        b_scale_k_idx = k_scale_idx # Scale K index matches scale block index
        b_scale_ptr_scalar = b_scale_ptr + (b_scale_n_block_idx * stride_b_scale_n + b_scale_k_idx * stride_b_scale_k)
        # Masking for N_scale dimension and K_scale dimension
        b_scale_mask = (b_scale_n_block_idx < tl.cdiv(N, BLOCK_SIZE_N_SCALE)) # K check implicit via loop bound? Let's keep it: & (b_scale_k_idx < tl.cdiv(K, BLOCK_SIZE_K_SCALE))
        b_scale = tl.load(b_scale_ptr_scalar, mask=b_scale_mask, other=1.0)
        # --- End Scale Loading ---

        # Inner loop iterates over K blocks *within* the current K_SCALE block
        k_inner_start = k_scale_offset
        k_inner_end = k_scale_offset + BLOCK_SIZE_K_SCALE

        for k_inner_offset in range(k_inner_start, k_inner_end, BLOCK_SIZE_K):
            # Calculate current K offsets for A and B loading
            offs_k = k_inner_offset + tl.arange(0, BLOCK_SIZE_K)

            # --- Load A ---
            a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
            a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K)
            # Use cache modifier .ca if num_stages >= 2 (likely true based on configs)
            a = tl.load(a_ptrs, mask=a_mask, other=0.0, cache_modifier=".ca")

            # --- Load B ---
            # B needs K x N block, loaded from N x K storage (stride_bn=K, stride_bk=1)
            b_base_ptrs = b_ptr + offs_n[None, :] * stride_bn # Base pointers for N elements in this block
            b_ptrs = b_base_ptrs + offs_k[:, None] * stride_bk # Offset by K
            b_mask = (offs_k[:, None] < K) & (offs_n[None, :] < N)
            # Use cache modifier .ca if num_stages >= 2
            b = tl.load(b_ptrs, mask=b_mask, other=0.0, cache_modifier=".ca")

            # --- Compute and Accumulate ---
            # Apply the single loaded scale factors for this entire k_scale_idx block
            accumulator += tl.dot(a, b, out_dtype=tl.float32) * a_scale_vec * b_scale

    # --- Store C ---
    accumulator = accumulator.to(c_ptr.dtype.element_ty)
    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, accumulator, mask=c_mask)


def custom_kernel(data: input_t) -> output_t:
    """
    Wrapper function to launch the FP8 matrix multiplication kernel
    with hardcoded configurations based on input dimensions.
    """
    a, b, a_scale, b_scale, c = data

    M, K = a.shape
    N = b.shape[0] # B is row-major N x K
    BLOCK_SIZE_K_SCALE = 128
    BLOCK_SIZE_N_SCALE = 128

    config_key = (M, N, K)
    selected_config = HARDCODED_CONFIGS.get(config_key, DEFAULT_CONFIG)

    block_m, block_n, block_k, group_m, num_warps, num_stages = selected_config
    grid = (triton.cdiv(M, block_m) * triton.cdiv(N, block_n),)

    # Launch the kernel
    matmul_fp8_kernel[grid](
        a, b, a_scale, b_scale, c,   # Data pointers
        M, N, K,                     # Dimensions
        # Strides for row-major tensors
        a.stride(0), a.stride(1),    # A: K, 1
        b.stride(0), b.stride(1),    # B: K, 1
        a_scale.stride(0), a_scale.stride(1), # a_scale: K/BLK_K_S, 1
        b_scale.stride(0), b_scale.stride(1), # b_scale: K/BLK_K_S, 1
        c.stride(0), c.stride(1),    # C: N, 1
        # Pass kernel meta-parameters explicitly
        BLOCK_SIZE_M=block_m,
        BLOCK_SIZE_N=block_n,
        BLOCK_SIZE_K=block_k,
        GROUP_SIZE_M=group_m,
        BLOCK_SIZE_K_SCALE=BLOCK_SIZE_K_SCALE, # Pass the fixed scale block size
        BLOCK_SIZE_N_SCALE=BLOCK_SIZE_N_SCALE, # Pass the fixed scale block size
        # Pass launch parameters explicitly
        num_warps=num_warps,
        num_stages=num_stages,
    )

    return c

python eval.py benchmark /tmp/tmpawz5z710