import torch
from task import input_t, output_t
import math

# 检查是否有ROCm支持
if not torch.cuda.is_available() or not torch.version.hip:
    raise RuntimeError("此代码需要ROCm/HIP支持的PyTorch")

def custom_kernel(data: input_t) -> output_t:
    """
    使用PyTorch和HIP优化的block-scale FP8 GEMM实现。
    Args:
        data: 包含以下元素的元组:
            a: torch.Tensor[float8_e4m3fnuz] 形状为 [m, k]
            b: torch.Tensor[float8_e4m3fnuz] 形状为 [n, k]
            a_scale: torch.Tensor[float32] 形状为 [m, k // 128]
            b_scale: torch.Tensor[float32] 形状为 [n // 128, k // 128]
            c: torch.Tensor[bfloat16] 形状为 [m, n]
    Returns:
        包含bf16格式输出的张量
    """
    a, b, a_scale, b_scale, _ = data
    
    # 确保数据连续性
    a = a.contiguous()
    b = b.contiguous()
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()
    
    # 获取维度
    M = a.size(0)
    K = a.size(1)
    N = b.size(0)
    
    # 将FP8转换为更高精度以进行计算
    a_fp32 = a.to(torch.float32)
    b_fp32 = b.to(torch.float32)
    
    # 应用块缩放
    # 重塑张量以便于应用缩放因子
    block_size_k = 128
    k_blocks = K // block_size_k
    
    # 重塑数据以应用缩放
    a_reshaped = a_fp32.view(M, k_blocks, block_size_k)
    b_reshaped = b_fp32.view(N, k_blocks, block_size_k)
    
    # 应用缩放因子
    for m in range(M):
        for kb in range(k_blocks):
            a_reshaped[m, kb] *= a_scale[m, kb]
    
    # 对b应用缩放
    n_block_size = 128
    n_blocks = math.ceil(N / n_block_size)
    
    for nb in range(n_blocks):
        n_start = nb * n_block_size
        n_end = min((nb + 1) * n_block_size, N)
        for kb in range(k_blocks):
            b_reshaped[n_start:n_end, kb] *= b_scale[nb, kb].unsqueeze(0)
    
    # 重塑回原始形状
    a_scaled = a_reshaped.reshape(M, K)
    b_scaled = b_reshaped.reshape(N, K)
    
    # 使用优化的矩阵乘法
    # 转置b以匹配PyTorch的matmul期望
    b_t = b_scaled.t()
    
    # 执行矩阵乘法
    c = torch.matmul(a_scaled, b_t)
    
    # 转换为bfloat16并返回
    return c.to(torch.bfloat16)
python eval.py test /tmp/tmpnzamcuu8