# 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>

#define TILE_SIZE 16

__global__ void custom_kernel(
    const __hip_fp8_e4m3_fnuz* a, 
    const __hip_fp8_e4m3_fnuz* b, 
    float* a_scale, 
    float* b_scale, 
    __hip_bfloat16* c, 
    int m, 
    int n, 
    int k) {
    __shared__ __hip_fp8_e4m3_fnuz tile_A[TILE_SIZE][TILE_SIZE*4];
    __shared__ __hip_fp8_e4m3_fnuz tile_B[TILE_SIZE][TILE_SIZE*4];

    int tile_row = blockIdx.x * TILE_SIZE*4; //A的列
    int tile_col = blockIdx.y * TILE_SIZE*4; //B的列

    int tx = threadIdx.x;   // 0~TILE_SIZE-1
    int ty = threadIdx.y;   // 0~TILE_SIZE-1

    float zero[4] = {0.0f, 0.0f, 0.0f, 0.0f};
    float result [16] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};

    for (int i = 0; i < k; i += TILE_SIZE){
        int k_start = i;
        if(k_start + ty < k && tile_row + tx*4 < m)
        {
            (__hip_fp8x4_e4m3_fnuz &)tile_A[ty][tx*4] = (__hip_fp8x4_e4m3_fnuz &)a[tile_row + (k_start + ty) * m + (tx*4)];
        }
        else{
            (__hip_fp8x4_e4m3_fnuz &)tile_A[ty][tx*4] = (__hip_fp8x4_e4m3_fnuz &)zero;
        }

        if(k_start + tx < k && tile_col + ty*4 < n)
        {
            (__hip_fp8x4_e4m3_fnuz &)tile_B[tx][ty*4] = (__hip_fp8x4_e4m3_fnuz &)b[tile_col + (k_start + tx) * n + (ty*4)];
        }
        else{
            (__hip_fp8x4_e4m3_fnuz &)tile_B[tx][ty*4] = (__hip_fp8x4_e4m3_fnuz &)zero;

        }

        float a_scale_shared[TILE_SIZE*4];
        float b_scale_shared;

        if(k_start + ty < k && tile_row + tx*4 < m)
        {
            int k_block = ((k_start + tx) / 128) * m ;
            int row = tile_row + tx * 4;
            a_scale_shared[tx*4] = a_scale[k_block + row + 0];
            a_scale_shared[tx*4 + 1] = a_scale[k_block + row + 1];
            a_scale_shared[tx*4 + 2] = a_scale[k_block + row + 2];
            a_scale_shared[tx*4 + 3] = a_scale[k_block + row + 3];
        }
        else{
            a_scale_shared[tx*4] = 0.0f;
            a_scale_shared[tx*4 + 1] = 0.0f;
            a_scale_shared[tx*4 + 2] = 0.0f;
            a_scale_shared[tx*4 + 3] = 0.0f;
        }

        if(k_start + tx < k && tile_col + ty*4 < n)
        {
            int sn = (n + 127) / 128;
            int k_block = k_start / 128;
            int col = tile_col + ty * 4;
            int n_block = col / 128;
            b_scale_shared = b_scale[n_block + k_block * sn];
        }
        else{
            b_scale_shared = 0.0f;
        }

        __syncthreads();
        float block_result [16] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
        float block_a[4] = {0.0f, 0.0f, 0.0f, 0.0f};
        float block_b[4] = {0.0f, 0.0f, 0.0f, 0.0f};
        
        //#pragma unroll TILE_SIZE
        for (int j = 0; j < TILE_SIZE; ++j) {
            block_a[0] = (float)tile_A[j][tx*4 + 0];
            block_a[1] = (float)tile_A[j][tx*4 + 1];
            block_a[2] = (float)tile_A[j][tx*4 + 2];
            block_a[3] = (float)tile_A[j][tx*4 + 3];
            block_b[0] = (float)tile_B[j][ty*4 + 0];
            block_b[1] = (float)tile_B[j][ty*4 + 1];
            block_b[2] = (float)tile_B[j][ty*4 + 2];
            block_b[3] = (float)tile_B[j][ty*4 + 3];
            block_result[0] += block_a[0] * block_b[0];
            block_result[1] += block_a[0] * block_b[1];
            block_result[2] += block_a[0] * block_b[2];
            block_result[3] += block_a[0] * block_b[3];
            block_result[4] += block_a[1] * block_b[0];
            block_result[5] += block_a[1] * block_b[1];
            block_result[6] += block_a[1] * block_b[2];
            block_result[7] += block_a[1] * block_b[3];
            block_result[8] += block_a[2] * block_b[0];
            block_result[9] += block_a[2] * block_b[1];
            block_result[10] += block_a[2] * block_b[2];
            block_result[11] += block_a[2] * block_b[3];
            block_result[12] += block_a[3] * block_b[0];
            block_result[13] += block_a[3] * block_b[1];
            block_result[14] += block_a[3] * block_b[2];
            block_result[15] += block_a[3] * block_b[3];
        }
        result[0] += block_result[0] *a_scale_shared[tx*4]*b_scale_shared;
        result[1] += block_result[1] *a_scale_shared[tx*4]*b_scale_shared;
        result[2] += block_result[2] *a_scale_shared[tx*4]*b_scale_shared;
        result[3] += block_result[3] *a_scale_shared[tx*4]*b_scale_shared;
        result[4] += block_result[4] *a_scale_shared[tx*4 + 1]*b_scale_shared;
        result[5] += block_result[5] *a_scale_shared[tx*4 + 1]*b_scale_shared;
        result[6] += block_result[6] *a_scale_shared[tx*4 + 1]*b_scale_shared;
        result[7] += block_result[7] *a_scale_shared[tx*4 + 1]*b_scale_shared;
        result[8] += block_result[8] *a_scale_shared[tx*4 + 2]*b_scale_shared;
        result[9] += block_result[9] *a_scale_shared[tx*4 + 2]*b_scale_shared;
        result[10] += block_result[10] *a_scale_shared[tx*4 + 2]*b_scale_shared;
        result[11] += block_result[11] *a_scale_shared[tx*4 + 2]*b_scale_shared;
        result[12] += block_result[12] *a_scale_shared[tx*4 + 3]*b_scale_shared;
        result[13] += block_result[13] *a_scale_shared[tx*4 + 3]*b_scale_shared;
        result[14] += block_result[14] *a_scale_shared[tx*4 + 3]*b_scale_shared;
        result[15] += block_result[15] *a_scale_shared[tx*4 + 3]*b_scale_shared;
        __syncthreads();
    }

    int C_row = tile_row + tx * 4;
    int C_col = tile_col + ty * 4;

    if ((C_row + 3 < m) && (C_col + 3 < n)) {
        c[C_row * n + C_col] = static_cast<__hip_bfloat16>(result[0]);
        c[C_row * n + C_col + 1] = static_cast<__hip_bfloat16>(result[1]);
        c[C_row * n + C_col + 2] = static_cast<__hip_bfloat16>(result[2]);
        c[C_row * n + C_col + 3] = static_cast<__hip_bfloat16>(result[3]);
        c[(C_row + 1) * n + C_col + 0] = static_cast<__hip_bfloat16>(result[4]);
        c[(C_row + 1) * n + C_col + 1] = static_cast<__hip_bfloat16>(result[5]);
        c[(C_row + 1) * n + C_col + 2] = static_cast<__hip_bfloat16>(result[6]);
        c[(C_row + 1) * n + C_col + 3] = static_cast<__hip_bfloat16>(result[7]);
        c[(C_row + 2) * n + C_col + 0] = static_cast<__hip_bfloat16>(result[8]);
        c[(C_row + 2) * n + C_col + 1] = static_cast<__hip_bfloat16>(result[9]);
        c[(C_row + 2) * n + C_col + 2] = static_cast<__hip_bfloat16>(result[10]);
        c[(C_row + 2) * n + C_col + 3] = static_cast<__hip_bfloat16>(result[11]);
        c[(C_row + 3) * n + C_col + 0] = static_cast<__hip_bfloat16>(result[12]);
        c[(C_row + 3) * n + C_col + 1] = static_cast<__hip_bfloat16>(result[13]);
        c[(C_row + 3) * n + C_col + 2] = static_cast<__hip_bfloat16>(result[14]);
        c[(C_row + 3) * n + C_col + 3] = static_cast<__hip_bfloat16>(result[15]);
        
    }

}

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 blockDim(TILE_SIZE, TILE_SIZE);
    dim3 gridDim((m + TILE_SIZE*4 - 1) / (TILE_SIZE*4), (n + TILE_SIZE*4 - 1) / (TILE_SIZE*4));

    custom_kernel<<<gridDim, blockDim>>>( (__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/tmpn9imp_sr
[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.
