# This script provides a template for using load_inline to run a HIP kernel for
import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

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) {
    // 跟随baseline实现但进行性能优化
    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.0f;
    
    // 分块处理k维度
    for(int i = 0; i < k; i += BLOCK) {
        // 在一个块内累积结果
        float block_result = 0.0f;
        
        #pragma unroll 4
        for(int ii = 0; ii < BLOCK; ++ii) {
            // 加载输入矩阵元素并转换为float进行计算
            // 注意输入是列主序的
            float av = (float)a[cx + (i + ii) * m];
            float bv = (float)b[cy + (i + ii) * n];
            block_result += av * bv; 
        }
        
        // 在进入下一个块之前，使用缩放因子缩放当前块的结果
        // 注意as和bs的索引方式
        float a_scale = as[cx + (i/BLOCK) * m];
        float b_scale = bs[(cy/BLOCK) + (i/BLOCK) * sn];
        result += block_result * a_scale * b_scale;
    }
    
    // 最后，将结果写入为bf16格式，注意C是行主序
    c[cx * n + cy] = (__hip_bfloat16)result;
}

// 用于性能优化的内核，特别是针对较大的矩阵
__global__ void optimized_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) {
    // 使用共享内存加速数据访问
    __shared__ float smem_a[32][BLOCK+4]; // 额外的填充以避免bank冲突
    __shared__ float smem_b[32][BLOCK+4];
    
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int bx = blockIdx.x;
    int by = blockIdx.y;
    
    // 当前线程块负责的输出区域
    int row_start = bx * 32;
    int col_start = by * 32;
    
    int sn = (n + BLOCK - 1) / BLOCK;
    
    // 每个线程负责输出块中的一个元素
    int row = row_start + tx;
    int col = col_start + ty;
    
    // 检查边界
    bool valid_output = (row < m) && (col < n);
    
    float result = 0.0f;
    
    // 沿着k维度分块处理
    for (int kb = 0; kb < k; kb += BLOCK) {
        // 协作加载数据块到共享内存
        for (int i = 0; i < BLOCK; i += 32) {
            int k_idx = kb + i + ty;
            if (row < m && k_idx < k) {
                // 从列主序的a加载数据
                smem_a[tx][i+ty] = (float)a[row + k_idx * m];
            }
            
            k_idx = kb + i + tx;
            if (col < n && k_idx < k) {
                // 从列主序的b加载数据
                smem_b[ty][i+tx] = (float)b[col + k_idx * n];
            }
        }
        
        __syncthreads();
        
        // 获取当前块的缩放因子
        float a_scale = (row < m) ? as[row + (kb/BLOCK) * m] : 0.0f;
        float b_scale = (col < n) ? bs[(col/BLOCK) + (kb/BLOCK) * sn] : 0.0f;
        float scale = a_scale * b_scale;
        
        // 计算当前块的内积
        float block_result = 0.0f;
        
        #pragma unroll 8
        for (int k_offset = 0; k_offset < BLOCK; ++k_offset) {
            if (kb + k_offset < k) {
                block_result += smem_a[tx][k_offset] * smem_b[ty][k_offset];
            }
        }
        
        // 应用缩放并累积到结果
        result += block_result * scale;
        
        __syncthreads();
    }
    
    // 将结果写入C（行主序）
    if (valid_output) {
        c[row * n + col] = (__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);
    
    // 根据矩阵大小选择合适的内核
    if (m >= 128 || n >= 128) {
        // 大矩阵使用优化内核
        dim3 block(32, 32);
        dim3 grid((m + 31) / 32, (n + 31) / 32);
        
        optimized_kernel<<<grid, block, 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
        );
    } else {
        // 小矩阵使用基本内核
        dim3 block(16, 16);
        dim3 grid((m + 15) / 16, (n + 15) / 16);
        
        custom_kernel<<<grid, block, 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/tmpd1pbnerm
[92mSuccessfully preprocessed all matching files.[0m

/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]
/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]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 2
[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
