#!POPCORN leaderboard amd-fp8-mm
import torch
from task import input_t, output_t

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

DEVICE = triton.runtime.driver.active.get_active_torch_device()

# 大幅简化配置空间，减少自动调优时间
fp8_gemm_configs = [
    Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': block_n, 'BLOCK_SIZE_K': 128}, 
           num_stages=4, num_warps=8)
    for block_m in [32, 64] 
    for block_n in [64, 128]
] + [
    # 添加几个专为大矩阵优化的配置
    Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128}, num_stages=5, num_warps=16),
    Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128}, num_stages=4, num_warps=16)
]

@triton.autotune(configs=fp8_gemm_configs, key=['M', 'N'])
@triton.jit
def fp8_gemm_kernel(
    a_ptr, b_ptr, c_ptr,
    a_s_ptr, b_s_ptr,
    M, N, K,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
):
    # 计算程序 ID 和偏移
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)
    
    # 使用优化的块分配提高缓存命中率
    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)
    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]
    
    # 初始化累加器 - 使用BF16进行中间计算可能提高MI300上的性能
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    # 计算K维度上的块数
    k_blocks = tl.cdiv(K, BLOCK_SIZE_K)
    
    # 预先计算比例因子指针
    a_scale_base = a_s_ptr + offs_m * k_blocks
    b_scale_base = b_s_ptr + (offs_n // BLOCK_SIZE_K) * k_blocks
    
    # 使用固定大小的K块进行循环，避免动态大小的内存访问
    for k in range(0, k_blocks):
        # 计算有效K值和掩码
        k_remain = min(BLOCK_SIZE_K, K - k * BLOCK_SIZE_K)
        
        # 优化的加载掩码，减少条件分支
        a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < k_remain)
        b_mask = (offs_n[None, :] < N) & (offs_k[:, None] < k_remain)
        
        # 使用向量化加载，MI300支持高效的向量指令
        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
        
        # 加载比例因子，优化边界检查
        a_s = tl.load(a_scale_base + k, mask=offs_m < M, other=0.0)
        b_s = tl.load(b_scale_base + k, mask=offs_n < N, other=0.0)
        
        # 使用矩阵乘法指令，启用快速数学运算
        accumulator += tl.dot(a, b, allow_tf32=True) * a_s[:, None] * b_s[None, :]
        
        # 更新指针，使用固定步长提高循环效率
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K
    
    # 计算输出指针和掩码
    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
    out_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    
    # 转换并存储结果
    c_out = accumulator.to(c_ptr.dtype.element_ty)
    tl.store(c_ptrs, c_out, mask=out_mask)


def custom_kernel(data: input_t) -> output_t:
    """
    Fast implementation of block-scale fp8 gemm for AMD MI300
    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
    
    # 确保内存连续性
    a = a.contiguous()
    b = b.contiguous()
    a_scale = a_scale.contiguous() 
    b_scale = b_scale.contiguous()
    
    # 获取维度
    M, K = a.shape
    N, K_ = b.shape
    assert K == K_, "Incompatible dimensions for matrix multiplication"
    
    # 创建输出张量
    c = torch.empty((M, N), dtype=torch.bfloat16, device=a.device)
    
    # 优化网格划分以减少调度开销
    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/tmp0uw5q_sh