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': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128}, num_stages=4, num_warps=4),
    Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128}, num_stages=4, num_warps=8),
    Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128}, num_stages=3, num_warps=8),
    
    # 针对不同形状矩阵的优化配置
    Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128}, num_stages=4, num_warps=4),
    Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128}, num_stages=4, num_warps=4),
    Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128}, num_stages=3, num_warps=8),
    Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128}, num_stages=3, num_warps=8),
    
    # 高throughput配置
    Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128}, num_stages=3, num_warps=8),
    Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128}, num_stages=3, num_warps=8),
    
    # 深度流水线配置
    Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 128}, num_stages=5, num_warps=4),
    Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128}, num_stages=5, num_warps=8),
]

@triton.autotune(configs=fp8_gemm_configs, key=['N', 'K'])
@triton.jit
def fp8_gemm_kernel_optimized(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):

    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)
    
    # 计算K维度的块数
    k_blocks = 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, :]
    b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
    a_s_ptrs = a_s_ptr + offs_m * k_blocks
    b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K) * k_blocks

    # 初始化累加器 - 使用更高精度
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    # 主循环 - 优化内存访问模式
    for k in range(k_blocks):
        # 计算有效的K维度大小
        k_remaining = tl.minimum(BLOCK_SIZE_K, K - k * BLOCK_SIZE_K)
        k_mask = offs_k < k_remaining
        
        # 加载数据块
        a_block = tl.load(a_ptrs, mask=k_mask[None, :], other=0.0)
        b_block = tl.load(b_ptrs, mask=k_mask[:, None], other=0.0)
        
        # 加载scale因子
        a_scale = tl.load(a_s_ptrs)
        b_scale = tl.load(b_s_ptrs)
        
        # 执行矩阵乘法并应用scale - 利用MI300的mfma指令
        dot_result = tl.dot(a_block, b_block)
        scaled_result = dot_result * a_scale[:, None] * b_scale[None, :]
        accumulator += scaled_result
        
        # 更新指针
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K
        a_s_ptrs += 1
        b_s_ptrs += 1

    # 转换输出数据类型
    c_result = accumulator.to(c_ptr.dtype.element_ty)
    
    # 存储结果 - 优化边界检查
    offs_m_out = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n_out = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + offs_m_out[:, None] * N + offs_n_out[None, :]
    
    # 边界掩码
    mask_m = offs_m_out < M
    mask_n = offs_n_out < N
    mask = mask_m[:, None] & mask_n[None, :]
    
    tl.store(c_ptrs, c_result, mask=mask)


def apply_strassen_inspired_blocking(M, N, K):
    """
    应用Strassen启发的分块策略，而不是完整的Strassen算法
    返回优化的网格配置
    """
    # 基于Strassen的思想：优先处理较大的正方形块
    # 但避免复杂的递归和大量的临时存储
    
    optimal_block_m = 32
    optimal_block_n = 32
    
    # 根据矩阵大小调整块尺寸
    if min(M, N) >= 128:
        if M >= N:
            optimal_block_m = min(128, M)
            optimal_block_n = min(64, N)
        else:
            optimal_block_m = min(64, M)
            optimal_block_n = min(128, N)
    elif min(M, N) >= 64:
        optimal_block_m = min(64, M)
        optimal_block_n = min(64, N)
    
    # 确保块大小是32的倍数（MI300 mfma指令优化）
    optimal_block_m = ((optimal_block_m + 31) // 32) * 32
    optimal_block_n = ((optimal_block_n + 31) // 32) * 32
    
    return optimal_block_m, optimal_block_n


def custom_kernel(data: input_t) -> output_t:
    """
    Strassen启发优化的block-scale FP8 GEMM实现
    使用简化的分块策略而非完整Strassen算法以避免超时
    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_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()
    a = a.contiguous()
    b = b.contiguous()

    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)
    
    # reshape a to proper 2D matrix
    a_2d = a.view(-1, K)
    
    # 应用Strassen启发的分块策略
    optimal_block_m, optimal_block_n = apply_strassen_inspired_blocking(M, N, K)
    
    # 创建输出张量
    result = a.new_empty(M, N, dtype=torch.get_default_dtype())
    
    # 使用优化的网格配置
    grid = lambda META: (
        triton.cdiv(M, META.get('BLOCK_SIZE_M', optimal_block_m)), 
        triton.cdiv(N, META.get('BLOCK_SIZE_N', optimal_block_n))
    )
    
    # 调用优化的kernel
    fp8_gemm_kernel_optimized[grid](
        a_2d, b, result, 
        a_scale.view(-1, a_scale.size(-1)), b_scale, 
        M, N, K
    )
    
    # 恢复原始形状
    original_shape = list(a.shape[:-1]) + [N]
    result = result.view(original_shape)
    
    return result.to(torch.bfloat16)
python eval.py benchmark /tmp/tmpbyzbbqio