#!POPCORN leaderboard amd-fp8-mm 
#!POPCORN GPU MI300
# 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 = 32;  // Increased block size for MI300X

__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, const uint M, const uint N, const uint K) {


    // converting 1D block Dim to 2D block Dim
    const uint cx  = BLOCK * blockIdx.x + (threadIdx.x / BLOCK) ;
    const uint cy = BLOCK * blockIdx.y + (threadIdx.x % BLOCK);
    
 //    __shared__ __hip_fp8_e4m3_fnuz shared_a[BLOCK][BLOCK];
 //    __shared__ __hip_fp8_e4m3_fnuz shared_b[BLOCK][BLOCK];
    
    // loadi a, b, as, bs
    
    // TODO: convert this return to proper if to enter the execution
    if(cx >= M || cy >= N) return;
    
    float result = 0;
    
    for (int k = 0; k < K; k+=BLOCK) {
        
        float block_result = 0;
        
 //       shared_a[threadIdx.y][threadIdx.x] = a[threadIdx.y + k + blockIdx.y * blockDim.y];
 //       shared_b[threadIdx.y][threadIdx.x] = b[cy + ((threadIdx.y + BLOCK) * N)];
        
        __syncthreads();
        
        for(int ii = 0; ii < BLOCK; ++ii){
            float av = (float)(a[ (k + ii) * M  + cx]);
            float bv = (float)(b[ (k + ii) * N  + cy]) ;
            block_result +=  av * bv; 
        }
        
        result += block_result * as[( (k / 128) * M + cx ) ] * bs[(  (k/128) * ((N + 128 - 1) /128) + (cy/128)) ];
    }
    
    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);
    
    // Optimized grid and block dimensions for MI300X
    dim3 block(BLOCK * BLOCK);  // 1024 threads per block
    
    dim3 grid((m + BLOCK - 1) / BLOCK, (n + BLOCK - 1) / BLOCK);
    
    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 benchmark /tmp/tmpg0icig_p
[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.
