import torch
import triton
import triton.language as tl
from task import input_t, output_t
from triton import Config

fp8_gemm_configs = [
    Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': block_n, 'BLOCK_SIZE_K': 128, "GROUP_SIZE_M": group, "BLOCK_SCALE": 128}, num_stages=num_stages, num_warps=8)
    for block_m in [16, 32, 64] for block_n in [32, 64, 128] for num_stages in [3, 4, 5, 6] for group in [4]
]

@triton.autotune(configs=fp8_gemm_configs, key=['N', 'K'])
@triton.jit
def deepgemm(a_ptr, b_ptr, c_ptr,
             lhs_ptr, rhs_ptr,
             M, N, K, 
             stride_am, stride_ak, 
             stride_bn, stride_bk,
             stride_cm, stride_cn,
             stride_lhsm, stride_lhsk,
             stride_rhsn, stride_rhsk, 
             BLOCK_SIZE_M: tl.constexpr,
             BLOCK_SIZE_N: tl.constexpr,
             BLOCK_SIZE_K: tl.constexpr,
             GROUP_SIZE_M: tl.constexpr,
             BLOCK_SCALE: tl.constexpr,
             N_SM: tl.constexpr,
             N_XCD: tl.constexpr
             ):
    
    pid = tl.program_id(0)
    # pid = (pid % N_XCD) * (N_SM // N_XCD) + (pid // N_XCD)
    n_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    n_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    tiles = n_pid_m * n_pid_n

    for tile in range(pid, tiles, N_SM):
        n_pid_in_group = GROUP_SIZE_M * n_pid_n
        group_id = tile // n_pid_in_group
        first_pid_m = group_id * GROUP_SIZE_M
        group_size_m = min(n_pid_m - first_pid_m, GROUP_SIZE_M)
        pid_m = first_pid_m + ((tile % n_pid_in_group) % group_size_m)
        pid_n = (tile % n_pid_in_group) // group_size_m

        block_start_m = pid_m * BLOCK_SIZE_M
        block_start_n = pid_n * BLOCK_SIZE_N

        offset_m = block_start_m + tl.arange(0, BLOCK_SIZE_M)
        offset_n = block_start_n + tl.arange(0, BLOCK_SIZE_N)
        offset_m = tl.where(offset_m < M, offset_m, 0)
        offset_n = tl.where(offset_n < N, offset_n, 0)
        offset_m = tl.max_contiguous(tl.multiple_of(offset_m, BLOCK_SIZE_M), BLOCK_SIZE_M)
        offset_n = tl.max_contiguous(tl.multiple_of(offset_n, BLOCK_SIZE_N), BLOCK_SIZE_N)
        offset_k = tl.arange(0, BLOCK_SIZE_K)

        a_base_ptr = a_ptr + (offset_m[:, None] * stride_am + offset_k[None, :] * stride_ak)
        b_base_ptr = b_ptr + (offset_k[:, None] * stride_bk + offset_n[None, :] * stride_bn)

        lhs_base_ptr = lhs_ptr + (offset_m * stride_lhsm)
        rhs_base_ptr = rhs_ptr + ((block_start_n // BLOCK_SCALE) * stride_rhsn)

        acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
            lhs = tl.load(lhs_base_ptr)
            rhs = tl.load(rhs_base_ptr)

            a = tl.load(a_base_ptr)
            b = tl.load(b_base_ptr)
            acc += tl.dot(a,b) * lhs[:, None] * rhs
            
            a_base_ptr += BLOCK_SIZE_K * stride_ak
            b_base_ptr += BLOCK_SIZE_K * stride_bk
            rhs_base_ptr += stride_rhsk 
            lhs_base_ptr += stride_lhsk
        
        c = acc.to(tl.bfloat16)
        offset_m = block_start_m + tl.arange(0, BLOCK_SIZE_M)
        offset_n = block_start_n + tl.arange(0, BLOCK_SIZE_N)
        offset_m = tl.max_contiguous(tl.multiple_of(offset_m, BLOCK_SIZE_M), BLOCK_SIZE_M)
        offset_n = tl.max_contiguous(tl.multiple_of(offset_n, BLOCK_SIZE_N), BLOCK_SIZE_N)
        c_ptrs = c_ptr + (offset_m[:, None] * stride_cm + offset_n[None, :] * stride_cn)
        c_mask = (offset_m[:, None] < M) & (offset_n[None, :] < N)
        tl.store(c_ptrs, c, mask=c_mask)




def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of block-scale fp8 gemm
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k],
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k],
            a_scale: torch.Tensor[float32] of shape [m, k // 128],
            b_scale: torch.Tensor[float32] of shape [n // 128, k // 128],
            c: torch.Tensor[bfloat16] of shape [m, n]
    Returns:
        Tensor containing output in bf16
    """
    # c: [m, n] is pre-allocated memory to avoid timing allocation overhead.
    a, b, lhs, rhs, c = data
    SM = 304
    M, K = a.shape
    N = b.shape[0]

    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
    deepgemm[grid](
        a, b, c,
        lhs, rhs,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        lhs.stride(0), lhs.stride(1),
        rhs.stride(0), rhs.stride(1),
        N_SM=SM, N_XCD=8
    )

    return c


    


python eval.py benchmark /tmp/tmp1fpfzvac