# This script provides a template for using load_inline to run a HIP kernel for
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_K = 128;  // 块大小与缩放因子的块大小一致
constexpr const int BLOCK_M = 64;   // 针对MI300优化的块大小
constexpr const int BLOCK_N = 64;   // 针对MI300优化的块大小
constexpr const int THREAD_M = 8;   // 每个线程处理的M维度元素数
constexpr const int THREAD_N = 8;   // 每个线程处理的N维度元素数
constexpr const int THREAD_K = 8;   // 每次迭代处理的K维度元素数

// 使用共享内存缓存数据
__global__ void custom_kernel(
    const __hip_fp8_e4m3_fnuz* __restrict__ a, 
    const __hip_fp8_e4m3_fnuz* __restrict__ b, 
    const float* __restrict__ a_scale, 
    const float* __restrict__ b_scale, 
    __hip_bfloat16* __restrict__ c, 
    int m, int n, int k) {
    
    // 计算块索引
    int block_m = blockIdx.x;
    int block_n = blockIdx.y;
    
    // 计算线程索引
    int thread_x = threadIdx.x % (BLOCK_M / THREAD_M);
    int thread_y = threadIdx.x / (BLOCK_M / THREAD_M);
    
    // 计算全局行列索引
    int row = block_m * BLOCK_M + thread_x * THREAD_M;
    int col = block_n * BLOCK_N + thread_y * THREAD_N;
    
    // 检查边界
    if (row >= m || col >= n) return;
    
    // 计算缩放因子索引
    int sn = (n + BLOCK_K - 1) / BLOCK_K;
    
    // 为每个线程分配累加器
    float acc[THREAD_M][THREAD_N] = {0.0f};
    
    // 遍历K维度的块
    for (int bk = 0; bk < k; bk += BLOCK_K) {
        // 获取当前块的缩放因子
        float a_scales[THREAD_M];
        float b_scales[THREAD_N];
        
        // 预取A矩阵的缩放因子
        #pragma unroll
        for (int tm = 0; tm < THREAD_M; tm++) {
            if (row + tm < m) {
                a_scales[tm] = a_scale[(row + tm) + (bk / BLOCK_K) * m];
            }
        }
        
        // 预取B矩阵的缩放因子
        #pragma unroll
        for (int tn = 0; tn < THREAD_N; tn++) {
            if (col + tn < n) {
                b_scales[tn] = b_scale[((col + tn) / BLOCK_K) + (bk / BLOCK_K) * sn];
            }
        }
        
        // 计算当前块的矩阵乘法
        for (int k_offset = 0; k_offset < BLOCK_K; k_offset += THREAD_K) {
            // 加载A矩阵数据
            float a_frag[THREAD_M][THREAD_K];
            #pragma unroll
            for (int tm = 0; tm < THREAD_M; tm++) {
                #pragma unroll
                for (int tk = 0; tk < THREAD_K; tk++) {
                    if (row + tm < m && bk + k_offset + tk < k) {
                        a_frag[tm][tk] = (float)a[(row + tm) + (bk + k_offset + tk) * m];
                    } else {
                        a_frag[tm][tk] = 0.0f;
                    }
                }
            }
            
            // 加载B矩阵数据
            float b_frag[THREAD_N][THREAD_K];
            #pragma unroll
            for (int tn = 0; tn < THREAD_N; tn++) {
                #pragma unroll
                for (int tk = 0; tk < THREAD_K; tk++) {
                    if (col + tn < n && bk + k_offset + tk < k) {
                        b_frag[tn][tk] = (float)b[(col + tn) + (bk + k_offset + tk) * n];
                    } else {
                        b_frag[tn][tk] = 0.0f;
                    }
                }
            }
            
            // 计算矩阵乘法
            #pragma unroll
            for (int tm = 0; tm < THREAD_M; tm++) {
                #pragma unroll
                for (int tn = 0; tn < THREAD_N; tn++) {
                    #pragma unroll
                    for (int tk = 0; tk < THREAD_K; tk++) {
                        acc[tm][tn] += a_frag[tm][tk] * a_scales[tm] * b_frag[tn][tk] * b_scales[tn];
                    }
                }
            }
        }
    }
    
    // 将结果写回全局内存
    #pragma unroll
    for (int tm = 0; tm < THREAD_M; tm++) {
        if (row + tm < m) {
            #pragma unroll
            for (int tn = 0; tn < THREAD_N; tn++) {
                if (col + tn < n) {
                    c[(row + tm) * n + (col + tn)] = (__hip_bfloat16)acc[tm][tn];
                }
            }
        }
    }
}

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 grid((m + BLOCK_M - 1) / BLOCK_M, (n + BLOCK_N - 1) / BLOCK_N);
    dim3 block((BLOCK_M / THREAD_M) * (BLOCK_N / THREAD_N));
    
    // 启动内核
    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/tmph12lzhdw
Using /home/runner/.cache/torch_extensions/py310_cpu as PyTorch extensions root...
[92mSuccessfully preprocessed all matching files.[0m
Detected CUDA files, patching ldflags
Emitting ninja build file /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/build.ninja...
Building extension module fp8_mm...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
Loading extension module fp8_mm...

/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: 1
ninja: no work to do.
