import torch
from task import input_t, output_t

# 检查是否有ROCm/HIP支持
def check_hip_available():
    if torch.cuda.is_available():
        device_name = torch.cuda.get_device_name()
        if "AMD" in device_name or "MI" in device_name:
            return True
    return False

# 尝试导入特定于AMD GPU的优化库
try:
    import torch_xla
    has_xla = True
except ImportError:
    has_xla = False

def custom_kernel(data: input_t) -> output_t:
    """
    Optimized implementation of block-scale FP8 GEMM.
    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
    
    # 检查并使用HIP设备
    if torch.cuda.is_available():
        device = torch.device("cuda")
        # 对于MI300，尝试设置更优的内存分配策略
        if check_hip_available():
            # 对于ROCm/HIP，设置流和缓存选项
            current_stream = torch.cuda.current_stream()
            torch.cuda.set_stream(current_stream)

    with torch.no_grad():
        # 将张量移动到GPU (如果可用)
        if torch.cuda.is_available():
            a = a.to(device)
            b = b.to(device)
            a_scale = a_scale.to(device)
            b_scale = b_scale.to(device)
            c = c.to(device)
            
        # Extract dimensions
        m, k = a.shape
        n = b.shape[0]
        block_shape_n = 128
        block_shape_k = 128

        # Process a_scale
        scale_k = a_scale.shape[1]
        a_scale = a_scale.view(m, scale_k, 1)  # Reshape to [m, scale_k, 1]
        a_scale = a_scale.expand(-1, -1, block_shape_k)  # Expand to [m, scale_k, block_shape_k]
        a_scale = a_scale.reshape(m, -1)[:, :k]  # Flatten and slice to match 'k'

        # Process b_scale
        scale_n, scale_k = b_scale.shape
        b_scale = b_scale.view(scale_n, scale_k, 1, 1)  # Reshape to [scale_n, scale_k, 1, 1]
        b_scale = b_scale.expand(-1, -1, block_shape_n, block_shape_k)  # Expand dimensions
        b_scale = b_scale.permute(0, 2, 1, 3)  # Permute to [scale_n, block_shape_n, scale_k, block_shape_k]
        b_scale = b_scale.reshape(scale_n * block_shape_n, scale_k * block_shape_k)  # Flatten
        b_scale = b_scale[:n, :k]  # Slice to match [n, k]

        # Explicit type conversion - 使用HIP优化的内存访问模式
        a_fp32 = a.to(torch.float32) * a_scale
        b_fp32 = b.to(torch.float32) * b_scale

        # 确保数据连续性以提高HIP性能
        a_fp32 = a_fp32.contiguous()
        b_fp32 = b_fp32.contiguous()

        # 对于大矩阵，分块计算以优化内存使用
        if m * n * k > 1e9:  # 如果矩阵非常大
            # 分块大小
            chunk_size = 1024
            result = torch.zeros((m, n), dtype=torch.float32, device=device)
            
            for i in range(0, m, chunk_size):
                i_end = min(i + chunk_size, m)
                for j in range(0, n, chunk_size):
                    j_end = min(j + chunk_size, n)
                    # 分块矩阵乘法
                    result[i:i_end, j:j_end] = torch.matmul(
                        a_fp32[i:i_end], 
                        b_fp32[j:j_end].T
                    )
            c.copy_(result.to(torch.bfloat16))
        else:
            # 标准矩阵乘法
            result = torch.matmul(a_fp32, b_fp32.T)  # FP32 matrix multiplication
            c.copy_(result.to(torch.bfloat16))  # Convert result to BF16 and copy to pre-allocated memory

        # 确保所有HIP操作完成
        if torch.cuda.is_available():
            torch.cuda.synchronize()

    return c
python eval.py benchmark /tmp/tmpg605b3df