from task import input_t, output_t

from typing import Tuple

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

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

BLOCK_SIZE_M = 64
BLOCK_SIZE_N = 64
BLOCK_SIZE_K = 128
NUM_STAGES = 4
NUM_WARPS = 8

@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):
    """
    Performs a matrix multiplication operation on FP8 matrices with scaling factors.

    Args:
        a_ptr (tl.tensor): Pointer to the first input matrix A. shape:= (M, K) - column-major
        b_ptr (tl.tensor): Pointer to the second input matrix B. shape:= (N, K) - column-major
        c_ptr (tl.tensor): Pointer to the output matrix C. shape:= (M, N) - row-major
        a_s_ptr (tl.tensor): Pointer to the scaling factors for matrix A. shape:= (M, K/BLOCK_SIZE)
        b_s_ptr (tl.tensor): Pointer to the scaling factors for matrix B. shape:= (N/BLOCK_SIZE, K/BLOCK_SIZE)
        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_m = tl.program_id(axis=0) # program row id for C 
    pid_n = tl.program_id(axis=1) # col id for C
    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, :] 
    a_ptrs = a_ptr + offs_m[:, None] + offs_k[None, :] * M # shape:= [BLOCK_SIZE_M, BLOCK_SIZE_K] 
    # b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None] 
    b_ptrs = b_ptr + offs_n[None, :] + offs_k[:, None] * N # (block of B_transposed) shape:= [BLOCK_SIZE_K, BLOCK_SIZE_N]

    a_s_ptrs = a_s_ptr + offs_m
    b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_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 = tl.load(a_ptrs)
        b = tl.load(b_ptrs)

        # a_s_ptrs = a_s_ptr + offs_m + i * M
        # b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K) + i * tl.cdiv(N, BLOCK_SIZE_K)
        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 * M
        b_ptrs += BLOCK_SIZE_K * N
        a_s_ptrs += M
        b_s_ptrs += tl.cdiv(N, BLOCK_SIZE_K)
        
    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, :] # c is row-major matrix
    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):
    """
    Perform a matrix multiplication using FP8 precision.

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

    Returns:
        torch.Tensor: The result of the matrix multiplication.
    """
    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)
    c = a.new_empty(*a.size()[:-1], N, dtype=torch.bfloat16)
    # grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
    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,)
                        #   BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K)
    return c

def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    # print(f'a: {a.shape}, {a.stride()}')
    # print(f'a_scale: {a_scale.shape}, {a_scale.stride()}')
    # print(f'b: {b.shape}, {b.stride()}')
    # print(f'b_scale: {b_scale.shape}, {b_scale.stride()}')
    # print(f'c: {c.shape}, {c.stride()}')
    # print('-'*100)
    c = fp8_gemm(a, a_scale, b, b_scale)
    return c

python eval.py benchmark /tmp/tmpun3aym9h