# import torch
# from task import input_t, output_t

# 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, a_scale, b_scale, c = data

#     # a is M x K in column-major order, we convert here for simplicity.
#     a = a.contiguous()
#     a_scale = a_scale.contiguous()
#     b_scale = b_scale.contiguous()

#     # constants
#     m = a.shape[0]
#     n = b.shape[0]
#     k = a.shape[1]
#     block_shape_n = 128
#     block_shape_k = 128
#     scale_n = b_scale.shape[0]
#     scale_k = b_scale.shape[1]

#     # Apply scaling to input 'a'
#     a_scale = a_scale.unsqueeze(-1).repeat(1, 1, block_shape_k)  # Shape: [m, scale_k, block_shape_k]
#     a_scale = a_scale.reshape(m, scale_k * block_shape_k) 
#     a_scale = a_scale[:, :k]

#     # Dequantize 'a', in your implementation you should do this at the end.
#     a = a.to(a_scale.dtype) * a_scale 

#     # Apply scaling to input 'b'
#     b_scale = (
#         b_scale.view(-1, 1)
#         .repeat(1, block_shape_n * block_shape_k)
#         .view(scale_n, scale_k, block_shape_n, block_shape_k)
#         .permute(0, 2, 1, 3)  # Reorder dimensions: [scale_n, blk_n, scale_k, blk_k]
#         .reshape(scale_n * block_shape_n, scale_k * block_shape_k)
#     )
#     b_scale = b_scale[:n, :k]

#     # Dequantize 'b', in your implementation you should do this at the end.
#     b = b.to(b_scale.dtype) * b_scale 

#     c[...] = (a @ b.T).to(torch.bfloat16)
#     return c

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


fp8_gemm_configs = [
 Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': block_n, 'BLOCK_SIZE_K': 128,"GROUP_SIZE_M": 2}, num_stages=num_stages, num_warps=num_warps)
    for block_m in [32] for block_n in [128] for num_stages in [3] for num_warps in [8] #, 4, 5, 6
]
@triton.autotune(configs=fp8_gemm_configs, key=['N', 'K'])
@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,
                    GROUP_SIZE_M: tl.constexpr, 
                    ):
    """
    Performs a matrix multiplication operation on FP8 matrices with scaling factors.

    Args:
        a_ptr (tl.tensor): Pointer to the first input matrix A.
        b_ptr (tl.tensor): Pointer to the second input matrix B.
        c_ptr (tl.tensor): Pointer to the output matrix C.
        a_s_ptr (tl.tensor): Pointer to the scaling factors for matrix A.
        b_s_ptr (tl.tensor): Pointer to the scaling factors for matrix B.
        M (int): Number of rows in matrix A and C.
        N (tl.constexpr): Number of columns in matrix B and C.
        K (tl.constexpr): Number of columns in matrix A and rows in matrix B.
        BLOCK_SIZE_M (tl.constexpr): Block size for the M dimension.
        BLOCK_SIZE_N (tl.constexpr): Block size for the N dimension.
        BLOCK_SIZE_K (tl.constexpr): Block size for the K dimension.

    Returns:
        None
    """
    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

    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):
    """
    Perform a matrix multiplication using FP8 precision.

    Args:
        a (torch.Tensor): The first input matrix, must be contiguous.
        a_s (torch.Tensor): The scaling factor for the first input matrix, must be contiguous.
        b (torch.Tensor): The second input matrix, must be contiguous.
        b_s (torch.Tensor): The scaling factor for the second input matrix, must be contiguous.

    Returns:
        torch.Tensor: The result of the matrix multiplication.
    """
    assert a.is_contiguous() and b.is_contiguous() and c.is_contiguous() 
    # assert a_s.is_contiguous() and b_s.is_contiguous()
    a_s.contiguous() 
    b_s.contiguous()
    assert a.shape[1] == b.shape[1], "Incompatible dimensions"  
    assert a.shape[0] == c.shape[0], "Incompatible dimensions" 
    assert b.shape[0] == c.shape[1], "Incompatible dimensions" 
    assert a.dtype == b.dtype, "Incompatible dtypes"

    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)
    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


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, a_scale, b_scale, c = data

    # # a is M x K in column-major order, we convert here for simplicity.
    a = a.contiguous()
    b = b.contiguous()
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()

    # # constants
    # m = a.shape[0]
    # n = b.shape[0]
    # k = a.shape[1]
    # block_shape_n = 128
    # block_shape_k = 128
    # scale_n = b_scale.shape[0]
    # scale_k = b_scale.shape[1]

    # Apply scaling to input 'a'
    # a_scale = a_scale.unsqueeze(-1).repeat(1, 1, block_shape_k)  # Shape: [m, scale_k, block_shape_k]
    # a_scale = a_scale.reshape(m, scale_k * block_shape_k) 
    # a_scale = a_scale[:, :k]

    # # Dequantize 'a', in your implementation you should do this at the end.
    # a = a.to(a_scale.dtype) * a_scale 

    # # Apply scaling to input 'b'
    # b_scale = (
    #     b_scale.view(-1, 1)
    #     .repeat(1, block_shape_n * block_shape_k)
    #     .view(scale_n, scale_k, block_shape_n, block_shape_k)
    #     .permute(0, 2, 1, 3)  # Reorder dimensions: [scale_n, blk_n, scale_k, blk_k]
    #     .reshape(scale_n * block_shape_n, scale_k * block_shape_k)
    # )
    # b_scale = b_scale[:n, :k]

    # # Dequantize 'b', in your implementation you should do this at the end.
    # b = b.to(b_scale.dtype) * b_scale 

    # c[...] = (a @ b.T).to(torch.bfloat16)
    # c = fp8_gemm(a, a_scale, b, b_scale, c)

    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)
    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_scale, b_scale, M, N, K)
    return c
python eval.py benchmark /tmp/tmplf1sywpm