# 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 mm_kernel_v3(const float* A, const float* B, float* C, int M, int K, int N) {
__global__ void mm_kernel_v3(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 K_TILE_SIZE = 8;

    //int tileRow = blockIdx.y;
    //int tileCol = blockIdx.x;

    int sn = (N + BLOCK - 1) / BLOCK;

    int thread_id = threadIdx.x % 64;
    int thread_row = 2*(thread_id / 16) + (thread_id % 2);
    int thread_col = (thread_id / 2) % 8;
    int warp_id = threadIdx.x / 64;
    int warp_row = warp_id / 2;
    int warp_col = warp_id % 2;

    //int C_tile_row_begin = blockIdx.y * 128;
    //int C_tile_col_begin = blockIdx.x * 128;
    //int A_tile_col_begin = C_tile_row_begin;
    //int B_tile_col_begin = C_tile_col_begin;

    float A_frag[8];
    float B_frag[8];
    float C_frag[8][8];

    __shared__ float A_tile[8][128];
    __shared__ float B_tile[8][128];

    #pragma unroll
    for(int i=0;i<8;i++){
        #pragma unroll
        for(int j=0;j<8;j++){
            C_frag[i][j] = 0;
        }
    }

    for(int k_iter = K/K_TILE_SIZE - 1; k_iter >= 0; k_iter --){
        #pragma unroll
        for(int i=threadIdx.x;i<8*128;i+=256){
            int A_row = k_iter * K_TILE_SIZE + i/128;
            int A_col = blockIdx.y * 128 + i % 128;
            if(A_col < M){
                A_tile[i/128][i%128] = (float)A[A_row * M + A_col] * as[(A_row / 128)*M + A_col];
            }
            else{
                A_tile[i/128][i%128] = 0;
            }
        }
        #pragma unroll
        for(int i=threadIdx.x;i<8*128;i+=256){
            int B_row = k_iter * K_TILE_SIZE + i/128;
            int B_col = blockIdx.x * 128 + i % 128;
            if(B_col < N){
                B_tile[i/128][i%128] = (float)B[B_row * N + B_col]*bs[(B_row / 128) * sn + B_col/128]; // bs[cy/BLOCK + i/BLOCK * sn];
            }
            else{
                B_tile[i/128][i%128] = 0;
            }
        }
        __syncthreads();
        #pragma unroll
        for(int k_step = 0; k_step < K_TILE_SIZE; k_step ++){
            #pragma unroll
            for(int i=0;i<8;i++){
                A_frag[i] = A_tile[k_step][warp_row * 64 + thread_row * 8 + i];
            }
            #pragma unroll
            for(int i=0;i<8;i++){
                B_frag[i] = B_tile[k_step][warp_col * 64 + thread_col * 8 + i];
            }
            #pragma unroll
            for(int i=0;i<8;i++){
                #pragma unroll
                for(int j=0;j<8;j++){
                    C_frag[i][j] += A_frag[i] * B_frag[j];
                }
            }
        }
        __syncthreads();
    }
    #pragma unroll
    for(int i=0;i<8;i++){
        #pragma unroll
        for(int j=0;j<8;j++){
            int C_row = blockIdx.y * 128 + warp_row * 64 + thread_row * 8 + i;
            int C_col = blockIdx.x * 128 + warp_col * 64 + thread_col * 8 + j;
            if(C_row < M && C_col < N){
                C[C_row * N + C_col] = (__hip_bfloat16)C_frag[i][j];
            }
        }
    }
}

// C(M, N) = A(M, K) * B(N, K)^T
// A B col major
// C row major
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(256);
    dim3 blocksPerGrid((n + 128 - 1) / 128,
                       (m + 128 - 1) / 128);

    mm_kernel_v3<<<blocksPerGrid, threadsPerBlock>>>((__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/tmpndbeusl5
[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.
