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

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) {
    
    // 使用2D块和线程结构
    int row = blockIdx.y * blockDim.y + threadIdx.y; // 输出行索引
    int col = blockIdx.x * blockDim.x + threadIdx.x; // 输出列索引
    
    if (row >= m || col >= n) return;
    
    int scale_blocks_n = (n + BLOCK - 1) / BLOCK;
    float acc = 0.0f;
    
    // 对k维度分块，每次处理一个BLOCK大小的块
    for (int blk = 0; blk < (k + BLOCK - 1) / BLOCK; ++blk) {
        // 获取此块的缩放因子
        float a_scale = as[row + blk * m];
        float b_scale = bs[(col/BLOCK) + blk * scale_blocks_n];
        float combined_scale = a_scale * b_scale;
        
        // 线程协作加载内存块 - 合并访存模式
        // 每个线程在本块中处理连续的k元素
        float block_result = 0.0f;
        int k_start = blk * BLOCK;
        int k_end = min(k, (blk + 1) * BLOCK);
        
        // 循环展开 - 每次处理4个元素可以提高指令级并行性
        for (int i = k_start; i < k_end; i += 4) {
            // 确保我们不会越界
            if (i + 3 < k_end) {
                // 合并访问a和b中的元素 - 一次加载多个连续元素
                float a_vals[4], b_vals[4];
                
                // 按合并模式加载a矩阵的元素
                #pragma unroll
                for (int j = 0; j < 4; ++j) {
                    a_vals[j] = (float)a[row + (i + j) * m];
                    b_vals[j] = (float)b[col + (i + j) * n];
                }
                
                // 计算部分点积
                #pragma unroll
                for (int j = 0; j < 4; ++j) {
                    block_result += a_vals[j] * b_vals[j];
                }
            } else {
                // 处理剩余元素（不能被4整除的情况）
                for (int j = 0; i + j < k_end; ++j) {
                    float a_val = (float)a[row + (i + j) * m];
                    float b_val = (float)b[col + (i + j) * n];
                    block_result += a_val * b_val;
                }
            }
        }
        
        // 将此块的结果乘以缩放因子，并累加到总结果
        acc += block_result * combined_scale;
    }
    
    // 将结果写入输出矩阵（合并写入）
    c[row * n + col] = (__hip_bfloat16)acc;
}

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);
    
    // 定义块和网格维度以优化内存访问模式
    dim3 threads(16, 16);
    dim3 grid((n + threads.x - 1) / threads.x, (m + threads.y - 1) / threads.y);
    
    custom_kernel<<<grid, threads>>>( 
        (__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
    );
}
"""

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 benchmark /tmp/tmp6a937y0j
[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
