# 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 BLOCKM = 64;
constexpr const int BLOCKN = 64;
constexpr const int BLOCK = 128;

using fp8_raw_t = uint8_t;
using fp8x8_t  = fp8_raw_t __attribute((ext_vector_type(8)));

using fp32x2_t = float __attribute((ext_vector_type(2)));
using fp32x4_t = float __attribute((ext_vector_type(4)));

template <typename Y, typename X>
__device__ constexpr Y bit_cast(const X& x)
{
    return __builtin_bit_cast(Y, x);
}

__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) {
                   
    // Your implementation here
    int tid = threadIdx.x;
    int lane_id = threadIdx.x % 64;
    int warp_id = threadIdx.x / 64;
    int lane_mod8 = lane_id % 8;
    int lane_div8 = lane_id / 8;
    int lane_mod16 = lane_id % 16;
    int lane_div16 = lane_id / 16;
    int warp_mod2 = warp_id % 2;
    int warp_div2 = warp_id / 2;

    int A_ldg_idx = (tid % 4) * 16 + (tid / 4) * m + BLOCKM * blockIdx.x;
    int B_ldg_idx = (tid % 4) * 16 + (tid / 4) * n + BLOCKN * blockIdx.y;
    
    int A_sts_idx = (tid % 4) * 16 + (tid / 4) * BLOCKM;
    int B_sts_idx = (tid % 4) * 16 + (tid / 4) * BLOCKN;
        
    int A_lds_idx = lane_mod16 + warp_div2 * 16 + lane_div16 * 8 * 64;
    int B_lds_idx = lane_mod16 + warp_mod2 * 16 + lane_div16 * 8 * 64;
    
    int sn = (n + BLOCK - 1) / BLOCK;

    as += BLOCKM * blockIdx.x;
    bs += BLOCKN * blockIdx.y / BLOCK;
    c +=  BLOCKN * blockIdx.y + BLOCKM * blockIdx.x * n;
    // smem
    __shared__ char smem[((BLOCKM + BLOCKN) * BLOCK) * sizeof(__hip_fp8_e4m3_fnuz) + (1 + BLOCKM) * sizeof(float)];
    __hip_fp8_e4m3_fnuz *smemA = reinterpret_cast<__hip_fp8_e4m3_fnuz *>(smem);
    __hip_fp8_e4m3_fnuz *smemB = reinterpret_cast<__hip_fp8_e4m3_fnuz *>(smem + BLOCKM * BLOCK * sizeof(__hip_fp8_e4m3_fnuz));
    float *smem_scales_a = reinterpret_cast<float *>(smem + (BLOCKM + BLOCKN) * BLOCK * sizeof(__hip_fp8_e4m3_fnuz));
    float *smem_scales_b = reinterpret_cast<float *>(smem + ((BLOCKM + BLOCKN) * BLOCK) * sizeof(__hip_fp8_e4m3_fnuz) + BLOCKM * sizeof(float));

    fp32x4_t acc[2][2];
    for(int i = 0; i < 2; i++) {
        for(int j = 0; j < 2; j++) {
            acc[i][j] = {0.0f};
        }
    }

    fp32x4_t A_ldg_reg[2];
    fp32x4_t B_ldg_reg[2];

    __hip_fp8_e4m3_fnuz A_lds_reg[2][8];
    __hip_fp8_e4m3_fnuz B_lds_reg[2][8];
    float A_scale_lds_reg[2][4];
    float B_scale_lds_reg;

    for(int itile = 0; itile < k; itile += BLOCK) {

        fp32x4_t block_acc[2][2];
        for(int i = 0; i < 2; i++) {
            for(int j = 0; j < 2; j++) {
                block_acc[i][j] = {0.0f};
            }
        }
        // ldg
        for(int i = 0; i < 2; i++) {
            A_ldg_reg[i] = *(fp32x4_t *)(&a[A_ldg_idx + (itile + i * 64) * m]);
            B_ldg_reg[i] = *(fp32x4_t *)(&b[B_ldg_idx + (itile + i * 64) * n]);
        }
        // sts
        for(int i = 0; i < 2; i++) {
            *(fp32x4_t *)(&smemA[A_sts_idx + (i * 64) * 64]) = A_ldg_reg[i];
            *(fp32x4_t *)(&smemB[B_sts_idx + (i * 64) * 64]) = B_ldg_reg[i];
        }
        if(warp_id < 1){
            smem_scales_a[tid] = as[tid];
        }
        if(tid == 0){
            smem_scales_b[0] = bs[0];
        }
        __syncthreads();

        for(int i = 0; i < 2; i++) {
            for(int k = 0; k < 4; k++) {
                A_scale_lds_reg[i][k] = smem_scales_a[k + lane_div16 * 4 + warp_div2 * 16 + 32 * i];
            }
        }
        B_scale_lds_reg = smem_scales_b[0];

        for(int k_step = 0; k_step < BLOCK; k_step+=32) {

            for(int i = 0; i < 2; i++) {
                for(int k = 0; k < 8; k++) {
                    A_lds_reg[i][k] = smemA[A_lds_idx + 32 * i + (k + k_step) * 64];
                }
            }
            for(int j = 0; j < 2; j++) {
                for(int k = 0; k < 8; k++) {
                    B_lds_reg[j][k] = smemB[B_lds_idx + 32 * j + (k + k_step) * 64];
                }
            }

            for(int i = 0; i < 2; i++) {
                for(int j = 0; j < 2; j++) {
                    block_acc[i][j] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
                        bit_cast<long>(A_lds_reg[i]), bit_cast<long>(B_lds_reg[j]), block_acc[i][j], 0, 0, 0);
                }
            }
            
        }
        
        for(int i = 0; i < 2; ++i) {
            for(int j = 0; j < 2; ++j) {
                for(int k = 0; k < 4; ++k) {
                    acc[i][j][k] += block_acc[i][j][k] * A_scale_lds_reg[i][k] * B_scale_lds_reg;
                }
            }
        }
        as += m;
        bs += sn;
        __syncthreads();
    }
    
    // finally, write the result as bf16
    for(int i = 0; i < 2; ++i) {
        for(int j = 0; j < 2; ++j) {
            for(int k = 0; k < 4; ++k) {
                int C_stg_idx = (lane_mod16 + warp_mod2 * 16 + 32 * j) 
                + (k + lane_div16 * 4 + warp_div2 * 16 + 32 * i) * n;
                c[C_stg_idx] = (__hip_bfloat16)acc[i][j][k];
            }
        }
    }
}

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);
    custom_kernel<<<dim3((m+BLOCKM-1)/BLOCKM, (n+BLOCKN-1)/BLOCKN), dim3(256), 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 leaderboard /tmp/tmpzvugvtqy
[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
[1/1] clang++ main.o hip.cuda.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o fp8_mm.so
