from task import input_t, output_t
from typing import Tuple

import torch
import triton
import triton.language as tl
from triton import Config

import os
os.environ["OPTIMIZE_EPILOGUE"] = "1"


fp8_gemm_configs = [
    Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128,
            # 'waves_per_eu': wpeu,
            'matrix_instr_nonkdim': 32},
           num_stages=ns, num_warps=num_warps)
    for ns in [2]#, 4, 5]
    # for wpeu in [0, 1, 2]
    for block_m in [16, 64, 128]#, 64]
    for num_warps in [4, 8, 16]
]

@triton.autotune(configs=fp8_gemm_configs, key=['M', 'N', 'K'], warmup=2)
@triton.jit
def fp8_gemm_kernel(a_ptr, b_ptr, c_ptr,
                    a_s_ptr, b_s_ptr,
                    M, N: tl.constexpr, K: tl.constexpr,
                    BLOCK_SIZE_M: tl.constexpr,
                    BLOCK_SIZE_N: tl.constexpr,
                    BLOCK_SIZE_K: tl.constexpr):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)
    
    k = tl.cdiv(K, BLOCK_SIZE_K)
    
    offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    
    a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
    b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
    
    a_s_ptrs = a_s_ptr + offs_m * k
    b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K) * k

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for i in range(k):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - i * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - i * BLOCK_SIZE_K, other=0.0)
        a_s = tl.load(a_s_ptrs)
        b_s = tl.load(b_s_ptrs)
        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K
        a_s_ptrs += 1
        b_s_ptrs += 1
    c = accumulator.to(c_ptr.dtype.element_ty)
    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)
    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=mask)


def fp8_gemm(a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor, c: torch.Tensor):
    assert a.is_contiguous() and b.is_contiguous(), 'Input tensors must be contiguous'
    assert a_s.is_contiguous() and b_s.is_contiguous(), 'Scaling factor tensors must be contiguous'
    M, K = a.shape
    N = b.size(0)
    if c is None:
        c = torch.empty((M, N), dtype=torch.float32, device=a.device)
    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']), triton.cdiv(N, META['BLOCK_SIZE_N']))
    fp8_gemm_kernel[grid](a, b, c, a_s, b_s, M, N, K)
    return c.to(torch.bfloat16)


def custom_kernel(data: input_t) -> output_t:
    """
    Optimized implementation of block-scale FP8 GEMM using Triton kernel for MI300 GPU.

    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k], COLUMN-MAJOR ORDER
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k], COLUMN-MAJOR ORDER
            a_scale: torch.Tensor[float32] of shape [m, k // 128], COLUMN-MAJOR ORDER
            b_scale: torch.Tensor[float32] of shape [n // 128, k // 128], COLUMN-MAJOR ORDER
            c: torch.Tensor[bfloat16] of shape [m, n], ROW-MAJOR ORDER
    Returns:
        Tensor containing output in bf16
    """
    a, b, a_scale, b_scale, c = data

    # Ensure all tensors are contiguous (do NOT transpose)
    a = a.contiguous()
    a_scale = a_scale.contiguous()
    b = b.contiguous()
    b_scale = b_scale.contiguous()

    return fp8_gemm(a, a_scale, b, b_scale, c)
python eval.py benchmark /tmp/tmpj9pwcc45