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()

# 针对MI300优化的配置
fp8_gemm_configs = [
    Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': block_n, 'BLOCK_SIZE_K': block_k}, 
           num_stages=num_stages, num_warps=num_warps)
    for block_m in [32, 64, 128] 
    for block_n in [32, 64, 128]
    for block_k in [32, 64, 128] 
    for num_stages in [2, 3]
    for num_warps in [4, 8]
]

# 标准矩阵乘法 - 安全的实现
@triton.autotune(configs=fp8_gemm_configs, key=['M', 'N', 'K'])
@triton.jit
def fp8_gemm_kernel(
    a_ptr, b_ptr, c_ptr,
    a_s_ptr, b_s_ptr,
    M, N, K,
    stride_am=1, stride_ak=None,
    stride_bn=None, stride_bk=1,
    stride_cm=1, stride_cn=None,
    BLOCK_SIZE_M: tl.constexpr = 64,
    BLOCK_SIZE_N: tl.constexpr = 64,
    BLOCK_SIZE_K: tl.constexpr = 32
):
    # 处理默认步长
    stride_ak = K if stride_ak is None else stride_ak
    stride_bn = N if stride_bn is None else stride_bn
    stride_cn = N if stride_cn is None else stride_cn
    
    # 获取程序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)
    
    # 边界检查掩码
    mask_m = offs_m < M
    mask_n = offs_n < N
    
    # 计算指针
    a_ptrs = a_ptr + offs_m[:, None] * stride_ak + offs_k[None, :] * stride_am
    b_ptrs = b_ptr + offs_n[None, :] * stride_bk + offs_k[:, None] * stride_bn
    
    # 计算缩放因子指针 - 确保正确处理K=128的情况
    scale_blocks = tl.maximum(K // 128, 1)  # 至少有1个缩放块
    a_s_ptrs = a_s_ptr + offs_m * scale_blocks
    b_s_ptrs = b_s_ptr + (offs_n // 128) * scale_blocks
    
    # 初始化累加器
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    # 计算循环
    k_blocks = tl.cdiv(K, BLOCK_SIZE_K)
    for i in range(k_blocks):
        k_idx = i * BLOCK_SIZE_K
        
        # 加载数据时确保边界安全
        a_mask = mask_m[:, None] & (offs_k[None, :] < K - k_idx)
        b_mask = mask_n[None, :] & (offs_k[:, None] < K - k_idx)
        
        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
        
        # 加载缩放因子 (K=128时特殊处理)
        if K == 128:
            a_s = tl.load(a_s_ptrs, mask=mask_m, other=1.0)
            b_s = tl.load(b_s_ptrs, mask=mask_n, other=1.0)
        else:
            a_s_idx = k_idx // 128
            b_s_idx = k_idx // 128
            a_s = tl.load(a_s_ptrs + a_s_idx, mask=mask_m, other=1.0)
            b_s = tl.load(b_s_ptrs + b_s_idx, mask=mask_n, other=1.0)
        
        # 应用缩放因子
        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        
        # 更新指针
        a_ptrs += BLOCK_SIZE_K * stride_am
        b_ptrs += BLOCK_SIZE_K * stride_bn
    
    # 存储结果
    c_ptrs = c_ptr + offs_m[:, None] * stride_cn + offs_n[None, :] * stride_cm
    
    # 确保边界安全
    output_mask = mask_m[:, None] & mask_n[None, :]
    c = accumulator.to(c_ptr.dtype.element_ty)
    tl.store(c_ptrs, c, mask=output_mask)

# 针对小矩阵的特殊优化实现 (m=64, n=64, k=128)
@triton.jit
def small_gemm_kernel(
    a_ptr, b_ptr, c_ptr,
    a_s_ptr, b_s_ptr,
    M, N, K
):
    # 获取程序ID
    pid = tl.program_id(axis=0)
    
    # 固定块大小
    BLOCK_SIZE_M: tl.constexpr = 32
    BLOCK_SIZE_N: tl.constexpr = 32
    BLOCK_SIZE_K: tl.constexpr = 32
    
    # 计算网格尺寸
    grid_n = tl.cdiv(N, BLOCK_SIZE_N)
    
    # 计算块位置
    pid_m = pid // grid_n
    pid_n = pid % grid_n
    
    # 计算偏移
    m_offset = pid_m * BLOCK_SIZE_M
    n_offset = pid_n * BLOCK_SIZE_N
    
    # 边界检查
    if m_offset < M and n_offset < N:
        # 计算块内偏移
        offs_m = m_offset + tl.arange(0, BLOCK_SIZE_M)
        offs_n = n_offset + tl.arange(0, BLOCK_SIZE_N)
        offs_k = tl.arange(0, BLOCK_SIZE_K)
        
        # 边界检查掩码
        mask_m = offs_m < M
        mask_n = offs_n < N
        
        # 计算指针
        a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
        b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
        
        # 缩放因子指针 - 对于K=128的特殊情况
        a_s_ptrs = a_s_ptr + offs_m
        b_s_ptrs = b_s_ptr + (offs_n // 128)
        
        # 初始化累加器
        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        
        # 计算循环 - 针对K=128的特殊处理
        for i in range(4):  # 128/32 = 4
            k_idx = i * BLOCK_SIZE_K
            
            # 加载数据时确保边界安全
            a_mask = mask_m[:, None] & (offs_k[None, :] < K - k_idx)
            b_mask = mask_n[None, :] & (offs_k[:, None] < K - k_idx)
            
            a = tl.load(a_ptrs, mask=a_mask, other=0.0)
            b = tl.load(b_ptrs, mask=b_mask, other=0.0)
            
            # 加载缩放因子 (K=128特殊处理)
            a_s = tl.load(a_s_ptrs, mask=mask_m, other=1.0)
            b_s = tl.load(b_s_ptrs, mask=mask_n, other=1.0)
            
            # 应用缩放因子
            accumulator += tl.dot(a, b) * 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, :]
        
        # 确保边界安全
        output_mask = mask_m[:, None] & mask_n[None, :]
        c = accumulator.to(c_ptr.dtype.element_ty)
        tl.store(c_ptrs, c, mask=output_mask)

# 针对中等大小矩阵的优化实现
@triton.jit
def medium_gemm_kernel(
    a_ptr, b_ptr, c_ptr,
    a_s_ptr, b_s_ptr,
    M, N, K
):
    # 获取程序ID
    pid = tl.program_id(axis=0)
    
    # 固定块大小
    BLOCK_SIZE_M: tl.constexpr = 64
    BLOCK_SIZE_N: tl.constexpr = 64
    BLOCK_SIZE_K: tl.constexpr = 32
    
    # 计算网格尺寸
    grid_n = tl.cdiv(N, BLOCK_SIZE_N)
    
    # 计算块位置
    pid_m = pid // grid_n
    pid_n = pid % grid_n
    
    # 计算偏移
    m_offset = pid_m * BLOCK_SIZE_M
    n_offset = pid_n * BLOCK_SIZE_N
    
    # 边界检查
    if m_offset < M and n_offset < N:
        # 计算块内偏移
        offs_m = m_offset + tl.arange(0, BLOCK_SIZE_M)
        offs_n = n_offset + tl.arange(0, BLOCK_SIZE_N)
        offs_k = tl.arange(0, BLOCK_SIZE_K)
        
        # 边界检查掩码
        mask_m = offs_m < M
        mask_n = offs_n < N
        
        # 计算指针
        a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
        b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
        
        # 缩放因子指针 - 确保正确处理
        scale_blocks = tl.maximum(K // 128, 1)  # 至少有1个缩放块
        a_s_ptrs = a_s_ptr + offs_m * scale_blocks
        b_s_ptrs = b_s_ptr + (offs_n // 128) * scale_blocks
        
        # 初始化累加器
        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        
        # 计算循环
        k_blocks = tl.cdiv(K, BLOCK_SIZE_K)
        for i in range(k_blocks):
            k_idx = i * BLOCK_SIZE_K
            
            # 加载数据时确保边界安全
            a_mask = mask_m[:, None] & (offs_k[None, :] < K - k_idx)
            b_mask = mask_n[None, :] & (offs_k[:, None] < K - k_idx)
            
            a = tl.load(a_ptrs, mask=a_mask, other=0.0)
            b = tl.load(b_ptrs, mask=b_mask, other=0.0)
            
            # 加载缩放因子
            if K == 128:
                a_s = tl.load(a_s_ptrs, mask=mask_m, other=1.0)
                b_s = tl.load(b_s_ptrs, mask=mask_n, other=1.0)
            else:
                a_s_idx = k_idx // 128
                b_s_idx = k_idx // 128
                a_s = tl.load(a_s_ptrs + a_s_idx, mask=mask_m, other=1.0)
                b_s = tl.load(b_s_ptrs + b_s_idx, mask=mask_n, other=1.0)
            
            # 应用缩放因子
            accumulator += tl.dot(a, b) * 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, :]
        
        # 确保边界安全
        output_mask = mask_m[:, None] & mask_n[None, :]
        c = accumulator.to(c_ptr.dtype.element_ty)
        tl.store(c_ptrs, c, mask=output_mask)

# 针对大矩阵的优化实现
@triton.jit
def large_gemm_kernel(
    a_ptr, b_ptr, c_ptr,
    a_s_ptr, b_s_ptr,
    M, N, K
):
    # 获取程序ID
    pid = tl.program_id(axis=0)
    
    # 使用更大的块大小处理大矩阵
    BLOCK_SIZE_M: tl.constexpr = 128
    BLOCK_SIZE_N: tl.constexpr = 128
    BLOCK_SIZE_K: tl.constexpr = 64
    
    # 计算网格尺寸
    grid_n = tl.cdiv(N, BLOCK_SIZE_N)
    
    # 计算块位置
    pid_m = pid // grid_n
    pid_n = pid % grid_n
    
    # 计算偏移
    m_offset = pid_m * BLOCK_SIZE_M
    n_offset = pid_n * BLOCK_SIZE_N
    
    # 边界检查
    if m_offset < M and n_offset < N:
        # 计算块内偏移
        offs_m = m_offset + tl.arange(0, BLOCK_SIZE_M)
        offs_n = n_offset + tl.arange(0, BLOCK_SIZE_N)
        offs_k = tl.arange(0, BLOCK_SIZE_K)
        
        # 边界检查掩码
        mask_m = offs_m < M
        mask_n = offs_n < N
        
        # 计算指针
        a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
        b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
        
        # 缩放因子指针 - 确保正确处理
        scale_blocks = tl.maximum(K // 128, 1)  # 至少有1个缩放块
        a_s_ptrs = a_s_ptr + offs_m * scale_blocks
        b_s_ptrs = b_s_ptr + (offs_n // 128) * scale_blocks
        
        # 初始化累加器
        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        
        # 计算循环
        k_blocks = tl.cdiv(K, BLOCK_SIZE_K)
        for i in range(k_blocks):
            k_idx = i * BLOCK_SIZE_K
            
            # 加载数据时确保边界安全
            a_mask = mask_m[:, None] & (offs_k[None, :] < K - k_idx)
            b_mask = mask_n[None, :] & (offs_k[:, None] < K - k_idx)
            
            a = tl.load(a_ptrs, mask=a_mask, other=0.0)
            b = tl.load(b_ptrs, mask=b_mask, other=0.0)
            
            # 加载缩放因子
            if K == 128:
                a_s = tl.load(a_s_ptrs, mask=mask_m, other=1.0)
                b_s = tl.load(b_s_ptrs, mask=mask_n, other=1.0)
            else:
                a_s_idx = k_idx // 128
                b_s_idx = k_idx // 128
                a_s = tl.load(a_s_ptrs + a_s_idx, mask=mask_m, other=1.0)
                b_s = tl.load(b_s_ptrs + b_s_idx, mask=mask_n, other=1.0)
            
            # 应用缩放因子
            accumulator += tl.dot(a, b) * 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, :]
        
        # 确保边界安全
        output_mask = mask_m[:, None] & mask_n[None, :]
        c = accumulator.to(c_ptr.dtype.element_ty)
        tl.store(c_ptrs, c, mask=output_mask)

def custom_kernel(data: input_t) -> output_t:
    """
    Optimized implementation of block-scale FP8 GEMM using Triton.
    基于矩阵尺寸自适应选择优化策略。
    
    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()
    
    # 获取矩阵维度
    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)
    
    # 创建输出张量
    c_temp = torch.empty((M, N), dtype=torch.float32, device=a.device)
    
    # 特殊处理 m=64, n=64, k=128 的情况
    if M == 64 and N == 64 and K == 128:
        # 使用专门针对这种情况优化的内核
        grid = (triton.cdiv(M, 32) * triton.cdiv(N, 32),)
        small_gemm_kernel[grid](a, b, c_temp, a_scale, b_scale, M, N, K)
    elif M <= 128 and N <= 128:
        # 小矩阵使用标准的GEMM
        grid = (triton.cdiv(M, 64), triton.cdiv(N, 64))
        fp8_gemm_kernel[grid](a, b, c_temp, a_scale, b_scale, M, N, K)
    elif M <= 512 and N <= 512:
        # 中等大小矩阵使用专门的优化
        grid = (triton.cdiv(M, 64) * triton.cdiv(N, 64),)
        medium_gemm_kernel[grid](a, b, c_temp, a_scale, b_scale, M, N, K)
    else:
        # 大矩阵使用更大块的分块处理
        grid = (triton.cdiv(M, 128) * triton.cdiv(N, 128),)
        large_gemm_kernel[grid](a, b, c_temp, a_scale, b_scale, M, N, K)
    
    # 转换为bfloat16并返回
    return c_temp.to(torch.bfloat16)
python eval.py benchmark /tmp/tmp9h5zprkj