# 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_M = 32;
constexpr const int BLOCK_N = 32;
constexpr const int BLOCK_K = 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,
    int scale_n)
{
    __shared__ float a_shared[BLOCK_M][BLOCK_K];
    __shared__ float b_shared[BLOCK_N][BLOCK_K];
    __shared__ float as_shared[BLOCK_M];
    __shared__ float bs_shared;

    int m_offset = BLOCK_M * blockIdx.x;
    int n_offset = BLOCK_N * blockIdx.y;

    int cx = m_offset + threadIdx.x;
    int cy = n_offset + threadIdx.y;

    int sn_offset = n_offset / 128;
    float acc = 0.0f;
    for (int i = 0; i < k; i += BLOCK_K)
    {
        int sk_offset = i / 128;        
        as_shared[threadIdx.x] = as[cx + sk_offset * m];
        bs_shared = bs[sn_offset + sk_offset * scale_n];
        
        float block_acc = 0.0f;
        for (int j = 0; j < BLOCK_K; ++j)
        {
            a_shared[threadIdx.x][j] = (float)a[cx + (i + j) * m];
            b_shared[threadIdx.y][j] = (float)b[cy + (i + j) * n];            
            __syncthreads();
                
            block_acc += a_shared[threadIdx.x][j] * b_shared[threadIdx.y][j];
        }
        acc += block_acc * as_shared[threadIdx.x] * bs_shared;
    }
    c[cx * n + cy] = (__hip_bfloat16)acc;
}

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);
    int scale_n = bs.size(0);

    auto threadsPerBlock = dim3(BLOCK_M, BLOCK_N);
    auto blocksPerGrid = dim3(m / BLOCK_M, n / BLOCK_N);
    hipLaunchKernelGGL(
        custom_kernel,
        blocksPerGrid,
        threadsPerBlock,
        0,
        0,
        (const __hip_fp8_e4m3_fnuz *)a.data_ptr(),
        (const __hip_fp8_e4m3_fnuz *)b.data_ptr(),
        (const float *)as.data_ptr(),
        (const float *)bs.data_ptr(),
        (__hip_bfloat16 *)c.data_ptr(),
        m,
        n,
        k,
        scale_n);
}
"""

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 submission.py
[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 [ok]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 0
[1/3] clang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=fp8_mm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -o main.o -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC
[2/3] /opt/rocm-6.3.1/bin/hipcc  -DWITH_HIP -DTORCH_EXTENSION_NAME=fp8_mm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC -DCUDA_HAS_FP16=1 -D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1 -DHIP_ENABLE_WARP_SYNC_BUILTINS=1 --offload-arch=gfx942 -std=c++20 -fno-gpu-rdc -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
[3/3] clang++ hip.cuda.o main.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

python eval.py test /tmp/tmpxgwov1mp
[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: 0
ninja: no work to do.
