#!POPCORN leaderboard amd-fp8-mm
#!POPCORN gpu MI300

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

torch._dynamo.config.recompile_limit = 10000
NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count

# https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
@triton.jit
def kernel(
    a_ptr,
    b_ptr,
    a_scale_ptr,
    b_scale_ptr,
    c_ptr,
    M: tl.constexpr,
    N: tl.constexpr,
    K: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    SPLIT_K: tl.constexpr = 1,
    BLOCK_Q: tl.constexpr = 128,
    GROUP_M: tl.constexpr = 0,
):
    # tl.static_assert(K % (BLOCK_Q * SPLIT_K) == 0)  # so we don't need to care about remainders

    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)

    pid_k = pid // (num_pid_m * num_pid_n)
    pid = pid % (num_pid_m * num_pid_n)
    if GROUP_M == 0:
        pid_m = pid // num_pid_n
        pid_n = pid % num_pid_n
    else:
        num_pid_in_group = GROUP_M * num_pid_n
        group_id = pid // num_pid_in_group
        first_pid_m = group_id * GROUP_M
        group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
        pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
        pid_n = (pid % num_pid_in_group) // group_size_m

    stride_am, stride_ak = 1, M
    stride_bn, stride_bk = 1, N
    stride_cm, stride_cn = N, 1
    stride_a_scalem, stride_a_scalek = 1, M
    stride_b_scalen, stride_b_scalek = 1, tl.cdiv(N, BLOCK_Q)

    # TODO: separate branch for SPLIT_K=1
    off_m = pid_m * BLOCK_M
    off_n = pid_n * BLOCK_N
    off_k = pid_k * (K // SPLIT_K)

    offs_m = off_m + tl.arange(0, BLOCK_M)
    offs_n = off_n + tl.arange(0, BLOCK_N)
    offs_k = off_k + tl.arange(0, BLOCK_K)
    if M % BLOCK_M == 0:
        offs_m = tl.max_contiguous(tl.multiple_of(offs_m, BLOCK_M), BLOCK_M)
    if N % BLOCK_N == 0:
        offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_N), BLOCK_N)

    tl.static_assert(BLOCK_Q % BLOCK_N == 0)
    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)
    a_scale_ptrs = a_scale_ptr + (off_k // BLOCK_Q) * stride_a_scalek + (offs_m[:, None] * stride_a_scalem)
    b_scale_ptrs = b_scale_ptr + (off_k // BLOCK_Q) * stride_b_scalek + (off_n // BLOCK_Q) * stride_b_scalen

    acc_master = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for _ in range(0, tl.cdiv(K, BLOCK_Q * SPLIT_K)):

        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
        for _ in tl.range(0, BLOCK_Q // BLOCK_K, flatten=True):
            a = tl.load(a_ptrs)  # (BLOCK_M, BLOCK_K)
            b = tl.load(b_ptrs)  # (BLOCK_K, BLOCK_N)
            acc = tl.dot(a, b, acc)  # (BLOCK_M, BLOCK_N)

            a_ptrs += BLOCK_K * stride_ak
            b_ptrs += BLOCK_K * stride_bk

        a_scale = tl.load(a_scale_ptrs)  # (BLOCK_M, 1)
        b_scale = tl.load(b_scale_ptrs)  # ()
        acc_master += acc * (a_scale * b_scale)

        a_scale_ptrs += stride_a_scalek
        b_scale_ptrs += stride_b_scalek

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)[None, :]
    mask = (offs_m < M) & (offs_n < N)
    c_ptrs = c_ptr + (M * N * pid_k) + (offs_m * stride_cm + offs_n * stride_cn)
    tl.store(c_ptrs, acc_master, mask=mask)


# use torch-compile to optimize split-k final reduction
@torch.compile(dynamic=False, mode="max-autotune-no-cudagraphs")
def my_sum(x, y):
    y.copy_(x.sum(0))


def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    # a: (m, k)
    # b: (n, k)
    # a_scale: (m, k/128), maybe padded
    # b_scale: (n/128, k/128), maybe padded
    # c: (m, n)
    # a, b, a_scale, b_scale are column-major
    # c is row-major

    M, K = a.shape
    N, _ = b.shape

    SPLIT_K = 2

    config = dict(
        BLOCK_M=64,
        BLOCK_N=64,
        BLOCK_K=64,
        SPLIT_K=SPLIT_K,
        GROUP_M=0,
        waves_per_eu=2,
        matrix_instr_nonkdim=32,
        # kpack=2,
        num_warps=4,
        num_stages=2,
    )

    if M == 1024 and N <= 576 and K == 7168:  # test case 2 and 7
        configs = [
            triton.Config(
                dict(
                    BLOCK_M=BLOCK_M,
                    BLOCK_N=BLOCK_N,
                    BLOCK_K=BLOCK_K,
                    SPLIT_K=SPLIT_K,
                    GROUP_M=0,
                    waves_per_eu=waves_per_eu,
                    matrix_instr_nonkdim=matrix_instr_nonkdim,
                ),
                num_warps=num_warps,
                num_stages=num_stages,
            )
            for BLOCK_M in [16, 32, 64]
            for BLOCK_N in [16, 32, 64]
            for BLOCK_K in [64]
            for waves_per_eu in [1, 2, 3]
            for matrix_instr_nonkdim in [16, 32]
            for num_warps in [4]
            for num_stages in [2]
        ]
        autotune_kernel = triton.autotune(configs, key=["M", "N", "K"])(kernel)
        c_fp32 = torch.empty(SPLIT_K, M, N, device=c.device, dtype=torch.float32)
        grid = lambda meta: (triton.cdiv(meta["M"], meta["BLOCK_M"]) * triton.cdiv(meta["N"], meta["BLOCK_N"]) * SPLIT_K,)
        autotune_kernel[grid](a, b, a_scale, b_scale, c_fp32, M, N, K)
        my_sum(c_fp32, c)
        for k, v in autotune_kernel.cache.items():
            print(f"{k}: {v}")

    else:
        num_blocks = triton.cdiv(M, config["BLOCK_M"]) * triton.cdiv(N, config["BLOCK_N"]) * SPLIT_K
        if SPLIT_K > 1:
            c_fp32 = torch.empty(SPLIT_K, M, N, device=c.device, dtype=torch.float32)
            kernel[(num_blocks,)](a, b, a_scale, b_scale, c_fp32, M, N, K, **config)
            my_sum(c_fp32, c)
        else:
            kernel[(num_blocks,)](a, b, a_scale, b_scale, c, M, N, K, **config)

    return c

python eval.py benchmark /tmp/tmp26nrnc2c
(1024, 576, 7168, 'torch.float8_e4m3fnuz', 'torch.float8_e4m3fnuz', 'torch.float32', 'torch.float32', 'torch.float32'): BLOCK_M: 64, BLOCK_N: 64, BLOCK_K: 64, SPLIT_K: 2, GROUP_M: 0, waves_per_eu: 2, matrix_instr_nonkdim: 32, num_warps: 4, num_ctas: 1, num_stages: 2, maxnreg: None
(1024, 576, 7168, 'torch.float8_e4m3fnuz', 'torch.float8_e4m3fnuz', 'torch.float32', 'torch.float32', 'torch.float32'): BLOCK_M: 64, BLOCK_N: 64, BLOCK_K: 64, SPLIT_K: 2, GROUP_M: 0, waves_per_eu: 3, matrix_instr_nonkdim: 32, num_warps: 4, num_ctas: 1, num_stages: 2, maxnreg: None
(1024, 576, 7168, 'torch.float8_e4m3fnuz', 'torch.float8_e4m3fnuz', 'torch.float32', 'torch.float32', 'torch.float32'): BLOCK_M: 64, BLOCK_N: 64, BLOCK_K: 64, SPLIT_K: 2, GROUP_M: 0, waves_per_eu: 2, matrix_instr_nonkdim: 32, num_warps: 4, num_ctas: 1, num_stages: 2, maxnreg: None
(1024, 576, 7168, 'torch.float8_e4m3fnuz', 'torch.float8_e4m3fnuz', 'torch.float32', 'torch.float32', 'torch.float32'): BLOCK_M: 64, BLOCK_N: 64, BLOCK_K: 64, SPLIT_K: 2, GROUP_M: 0, waves_per_eu: 3, matrix_instr_nonkdim: 32, num_warps: 4, num_ctas: 1, num_stages: 2, maxnreg: None
(1024, 512, 7168, 'torch.float8_e4m3fnuz', 'torch.float8_e4m3fnuz', 'torch.float32', 'torch.float32', 'torch.float32'): BLOCK_M: 64, BLOCK_N: 64, BLOCK_K: 64, SPLIT_K: 2, GROUP_M: 0, waves_per_eu: 2, matrix_instr_nonkdim: 32, num_warps: 4, num_ctas: 1, num_stages: 2, maxnreg: None
(1024, 512, 7168, 'torch.float8_e4m3fnuz', 'torch.float8_e4m3fnuz', 'torch.float32', 'torch.float32', 'torch.float32'): BLOCK_M: 64, BLOCK_N: 64, BLOCK_K: 64, SPLIT_K: 2, GROUP_M: 0, waves_per_eu: 3, matrix_instr_nonkdim: 32, num_warps: 4, num_ctas: 1, num_stages: 2, maxnreg: None
(1024, 512, 7168, 'torch.float8_e4m3fnuz', 'torch.float8_e4m3fnuz', 'torch.float32', 'torch.float32', 'torch.float32'): BLOCK_M: 64, BLOCK_N: 64, BLOCK_K: 64, SPLIT_K: 2, GROUP_M: 0, waves_per_eu: 3, matrix_instr_nonkdim: 32, num_warps: 4, num_ctas: 1, num_stages: 2, maxnreg: None
(1024, 512, 7168, 'torch.float8_e4m3fnuz', 'torch.float8_e4m3fnuz', 'torch.float32', 'torch.float32', 'torch.float32'): BLOCK_M: 64, BLOCK_N: 64, BLOCK_K: 64, SPLIT_K: 2, GROUP_M: 0, waves_per_eu: 2, matrix_instr_nonkdim: 32, num_warps: 4, num_ctas: 1, num_stages: 2, maxnreg: None
