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

import triton
import triton.language as tl

from task import input_t, output_t


# def get_autotune_config():
#     return [
#         triton.Config(
#             {
#                 "BLOCK_M": 128,
#                 "BLOCK_N": 128,
#                 "BLOCK_K": 128,
#             },
#             num_stages=1,
#             num_warps=2,
#         ),
#         triton.Config(
#             {
#                 "BLOCK_M": 128,
#                 "BLOCK_N": 64,
#                 "BLOCK_K": 128,
#             },
#             num_stages=1,
#             num_warps=2,
#         ),
#         triton.Config(
#             {
#                 "BLOCK_M": 64,
#                 "BLOCK_N": 64,
#                 "BLOCK_K": 128,
#             },
#             num_stages=1,
#             num_warps=2,
#         ),
#         triton.Config(
#             {
#                 "BLOCK_M": 64,
#                 "BLOCK_N": 64,
#                 "BLOCK_K": 128,
#             },
#             num_stages=2,
#             num_warps=4,
#         ),
#         triton.Config(
#             {
#                 "BLOCK_M": 64,
#                 "BLOCK_N": 16,
#                 "BLOCK_K": 128,
#             },
#             num_stages=2,
#             num_warps=4,
#         ),
#     ]


#@triton.autotune(configs=get_autotune_config(), key=["M", "N", "K"])
@triton.jit
def fp8_mm_kernel(
    a_ptr,
    b_ptr,
    a_scale_ptr,
    b_scale_ptr,
    c_ptr,
    M,
    N,
    K,
    SCALE_N,
    # GROUP_N: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    pid = tl.program_id(0)
    grid_m = tl.cdiv(M, BLOCK_M)
    #grid_n = tl.cdiv(N, BLOCK_N)

    # if GROUP_M == 1:
    pid_m = pid % grid_m
    pid_n = pid // grid_m
    # else:
    # grouped_grid_m = grid_m * GROUP_N
    # group_pid = pid // grouped_grid_m
    # first_pid_n = group_pid * GROUP_N
    # group_n = min(grid_n - first_pid_n, GROUP_N)
    # pid_m = (pid % grouped_grid_m) // group_n
    # pid_n = first_pid_n + (pid % group_n)

    m_range = tl.arange(0, BLOCK_M)[:, None] + pid_m * BLOCK_M
    n_range = tl.arange(0, BLOCK_N)[:, None] + pid_n * BLOCK_N
    k_range = tl.arange(0, BLOCK_K)[None, :]

    a_offset = a_ptr + m_range
    b_offset = b_ptr + n_range
    a_scale_offset = a_scale_ptr + m_range
    b_scale_offset = b_scale_ptr + (pid_n * BLOCK_N) // 128
    
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in tl.range(0, tl.cdiv(K, BLOCK_K)):
        a = tl.load(a_offset + k_range * M, (m_range < M) & (k_range < K))
        b = tl.load(b_offset + k_range * N, (n_range < N) & (k_range < K))
        sk_offset = (k * BLOCK_K) // 128
        a_scale = tl.load(a_scale_offset + sk_offset * M)
        b_scale = tl.load(b_scale_offset + sk_offset * SCALE_N)
        acc += tl.dot(a, b.T) * a_scale * b_scale
        k_range += BLOCK_K

    n_range = n_range.T
    tl.store(
        c_ptr + m_range * N + n_range,
        acc.to(tl.bfloat16),
        (m_range < M) & (n_range < N),
    )


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

    M, K = a.shape
    N, K = b.shape
    SCALE_N, SCALE_K = b_scale.shape
    GROUP_N = 8

    BLOCK_M = 32
    BLOCK_N = 16
    BLOCK_K = 128
    num_warps = 4
    num_stages = 2
    #1D grid of blocks
    # def grid(META):
    #     return (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)
    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
    fp8_mm_kernel[grid](
        a,
        b,
        a_scale,
        b_scale,
        c,
        M,
        N,
        K,
        SCALE_N,
        # GROUP_N,
        BLOCK_M,
        BLOCK_N,
        BLOCK_K,
        num_warps=num_warps,
        num_stages=num_stages,
    )
    return c

python eval.py benchmark /tmp/tmp8gnq3bh6