from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

CPP_WRAPPER = """
void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c);
"""

CUDA_SRC = """
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

// 保持原有的块大小，但优化内部实现
constexpr const int BLOCK = 128;

__global__ void custom_kernel(const __hip_fp8_e4m3_fnuz* a, const __hip_fp8_e4m3_fnuz* b, const float* as, const float* bs, 
                   __hip_bfloat16* c, int m, int n, int k) {
    // 保持原有的线程索引计算
    int cx = threadIdx.x + blockDim.x * blockIdx.x;
    int cy = threadIdx.y + blockDim.y * blockIdx.y;
    if(cx >= m || cy >= n) return;
    
    int sn = (n + BLOCK - 1) / BLOCK;
    
    float result = 0;
    
    // 主循环保持不变，仍按块处理
    for(int i = 0; i < k; i += BLOCK) {
        float block_result = 0;

        // 循环展开，每次处理4个元素，减少循环开销
        // 假设BLOCK是4的倍数（128符合条件）
        #pragma unroll 4
        for(int ii = 0; ii < BLOCK; ii += 4) {
            // 优化：预取数据并进行更高效的计算
            
            // 首先计算内存地址，避免重复计算
            int idx_a0 = cx + (i + ii) * m;
            int idx_a1 = cx + (i + ii + 1) * m;
            int idx_a2 = cx + (i + ii + 2) * m;
            int idx_a3 = cx + (i + ii + 3) * m;
            
            int idx_b0 = cy + (i + ii) * n;
            int idx_b1 = cy + (i + ii + 1) * n;
            int idx_b2 = cy + (i + ii + 2) * n;
            int idx_b3 = cy + (i + ii + 3) * n;
            
            // 预取并转换数据
            float a0 = (float)a[idx_a0];
            float a1 = (float)a[idx_a1];
            float a2 = (float)a[idx_a2];
            float a3 = (float)a[idx_a3];
            
            float b0 = (float)b[idx_b0];
            float b1 = (float)b[idx_b1];
            float b2 = (float)b[idx_b2];
            float b3 = (float)b[idx_b3];
            
            // 计算点积，避免多次乘法累加的开销
            block_result += a0 * b0 + a1 * b1 + a2 * b2 + a3 * b3;
        }
        
        // 缩放计算保持不变，确保正确性
        float scale_a = as[cx + i/BLOCK * m];
        float scale_b = bs[cy/BLOCK + i/BLOCK * sn];
        result += block_result * scale_a * scale_b;
    }
    
    // 写回结果
    c[cx * n + cy] = (__hip_bfloat16)result;
}

// 优化线程块配置但保持基本算法不变
void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c) {
    int m = a.size(0);
    int n = b.size(0);
    int k = a.size(1);
    
    // 根据AMD GPU架构特性优化线程块大小
    // 通常AMD GPU的波前大小是64
    dim3 blockDim(16, 4);  // 64线程，接近波前大小，有利于计算效率
    dim3 gridDim((m + blockDim.x - 1) / blockDim.x, (n + blockDim.y - 1) / blockDim.y);
    
    custom_kernel<<<gridDim, blockDim, 0, 0>>> (
        (__hip_fp8_e4m3_fnuz*)a.data_ptr(), 
        (__hip_fp8_e4m3_fnuz*)b.data_ptr(), 
        as.data_ptr<float>(), 
        bs.data_ptr<float>(), 
        (__hip_bfloat16*)c.data_ptr(), 
        m, n, k
    );
}
"""

import os
os.environ["CXX"] = "clang++"

module = load_inline(
    name='fp8_mm',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['fp8_mm'],
    verbose=True,
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20"],
)

def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    module.fp8_mm(a, b, a_scale, b_scale, c)
    return c
python eval.py leaderboard /tmp/tmpkcr8ie9h
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [skipped, already hipified]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
[1/1] clang++ main.o hip.cuda.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o fp8_mm.so
