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;
constexpr const int THREAD_TILE_M = 2;  // 每个线程处理2行
constexpr const int THREAD_TILE_N = 2;  // 每个线程处理2列

__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) {
                   
    // 线程的全局基础索引，现在每个线程负责多个元素
    const int thread_row = threadIdx.x + blockDim.x * blockIdx.x;
    const int thread_col = threadIdx.y + blockDim.y * blockIdx.y;
    
    // 线程的行范围
    const int row_start = thread_row * THREAD_TILE_M;
    const int col_start = thread_col * THREAD_TILE_N;
    
    // 预先检查边界条件
    bool valid = row_start < m && col_start < n;
    if (!valid) return;
    
    // 预计算常量，减少重复计算
    const int sn = (n + BLOCK - 1) / BLOCK;
    
    // 为每个tile位置创建结果累加器
    float results[THREAD_TILE_M][THREAD_TILE_N] = {0};
    
    // 对K维度分块处理
    #pragma unroll 1  // 不展开外循环，让编译器决定
    for(int i = 0; i < k; i += BLOCK) {
        // 确定当前块的实际大小（处理边界情况）
        const int current_block_size = min(BLOCK, k - i);
        
        // 为当前块的所有位置创建局部结果
        float block_results[THREAD_TILE_M][THREAD_TILE_N] = {0};
        
        // 处理当前块内的K值
        #pragma unroll 4  // 适度展开内循环
        for(int ii = 0; ii < current_block_size; ++ii) {
            // 预计算基础偏移量，减少重复计算
            const int k_offset = i + ii;
            
            // 对每个tile位置循环
            #pragma unroll
            for(int tile_m = 0; tile_m < THREAD_TILE_M; ++tile_m) {
                const int row = row_start + tile_m;
                if(row >= m) continue;  // 行边界检查
                
                // 预取A矩阵元素，减少重复加载
                const float a_val = (float)a[row + k_offset * m];
                
                #pragma unroll
                for(int tile_n = 0; tile_n < THREAD_TILE_N; ++tile_n) {
                    const int col = col_start + tile_n;
                    if(col >= n) continue;  // 列边界检查
                    
                    // 加载B矩阵元素
                    const float b_val = (float)b[col + k_offset * n];
                    
                    // 累加到当前块的结果中
                    block_results[tile_m][tile_n] += a_val * b_val;
                }
            }
        }
        
        // 应用缩放因子并累加到最终结果
        #pragma unroll
        for(int tile_m = 0; tile_m < THREAD_TILE_M; ++tile_m) {
            const int row = row_start + tile_m;
            if(row >= m) continue;
            
            // 一次性加载A的缩放因子，避免重复加载
            const float a_scale = as[row + i/BLOCK * m];
            
            #pragma unroll
            for(int tile_n = 0; tile_n < THREAD_TILE_N; ++tile_n) {
                const int col = col_start + tile_n;
                if(col >= n) continue;
                
                // 加载B的缩放因子
                const float b_scale = bs[col/BLOCK + i/BLOCK * sn];
                
                // 应用缩放因子并累加到最终结果
                results[tile_m][tile_n] += block_results[tile_m][tile_n] * a_scale * b_scale;
            }
        }
    }
    
    // 将最终结果写回全局内存
    #pragma unroll
    for(int tile_m = 0; tile_m < THREAD_TILE_M; ++tile_m) {
        const int row = row_start + tile_m;
        if(row >= m) continue;
        
        #pragma unroll
        for(int tile_n = 0; tile_n < THREAD_TILE_N; ++tile_n) {
            const int col = col_start + tile_n;
            if(col >= n) continue;
            
            // 写回结果
            c[row * n + col] = (__hip_bfloat16)results[tile_m][tile_n];
        }
    }
}

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架构特点
    dim3 blockDim(8, 8);  // 64个线程 - 接近AMD波前大小
    
    // 每个线程处理多个元素，调整网格大小
    dim3 gridDim((m + blockDim.x * THREAD_TILE_M - 1) / (blockDim.x * THREAD_TILE_M), 
                 (n + blockDim.y * THREAD_TILE_N - 1) / (blockDim.y * THREAD_TILE_N));
    
    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 benchmark /tmp/tmp812itiiu
[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
ninja: no work to do.
