# 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 = 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) {
                   
    // 定义共享内存块
    constexpr const int TILE_DIM = 32;
    __shared__ float shared_a[TILE_DIM][TILE_DIM];
    __shared__ float shared_b[TILE_DIM][TILE_DIM];
    
    int cx = threadIdx.x + blockDim.x * blockIdx.x;
    int cy = threadIdx.y + blockDim.y * blockIdx.y;
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    
    int sn = (n + BLOCK - 1) / BLOCK;
    
    float result = 0;
    // split loop into an outer loop over different blocks, and an inner loop within one block.
    // we can assume k % BLOCK == 0.
    for(int i = 0; i < k; i += BLOCK) {
        // block results accumulates the inner product across a single block.
        // within each block, scales are constant, so we can lift the scaling 
        // outside of the inner loop.
        float block_result = 0;
        
        // 循环处理一个BLOCK内的数据，每次加载一个TILE_DIM大小的块到共享内存
        for(int tile_start = i; tile_start < i + BLOCK; tile_start += TILE_DIM) {
            // 加载A矩阵子块到共享内存
            if(cx < m && (tile_start + ty) < k) {
                shared_a[tx][ty] = (float)a[cx + (tile_start + ty) * m];
            } else {
                shared_a[tx][ty] = 0.0f;
            }
            
            // 加载B矩阵子块到共享内存
            if(cy < n && (tile_start + tx) < k) {
                shared_b[tx][ty] = (float)b[cy + (tile_start + tx) * n];
            } else {
                shared_b[tx][ty] = 0.0f;
            }
            
            // 确保所有线程都完成数据加载
            __syncthreads();
            
            // 计算当前子块的内积
            for(int j = 0; j < TILE_DIM && (tile_start + j) < (i + BLOCK); j++) {
                block_result += shared_a[tx][j] * shared_b[j][ty];
            }
            
            // 确保所有线程完成计算再加载下一个子块
            __syncthreads();
        }
        
        // before we can go to the next block, scale the result of the current block
        // and accumulate to final result
        // note the different indexing into as and bs
        if(cx < m && cy < n) {
            result += block_result * as[cx + i/BLOCK * m] * bs[cy/BLOCK + i/BLOCK * sn];
        }
    }
    
    // finally, write the result as bf16
    if(cx < m && cy < n) {
        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);
    
    dim3 threadsPerBlock(32, 32);  // 32x32线程块
    dim3 numBlocks((m + threadsPerBlock.x - 1) / threadsPerBlock.x, 
                  (n + threadsPerBlock.y - 1) / threadsPerBlock.y);
    
    custom_kernel<<<numBlocks, threadsPerBlock, 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
    );
    //C10_CUDA_CHECK(cudaGetLastError());
}
"""

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/tmpwpelskt6
[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.
