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

using f32x16 = __attribute__( (__vector_size__(16 * sizeof(float)) )) float;

constexpr const int BLOCK = 128;

__device__ static uint64_t pack_u16x8_to_u64(const uint16_t* src)
{
    uint16_t dst[] = { src[1], src[3], src[5], src[7] };
    return *reinterpret_cast<uint64_t*>(dst);
}

__global__ void custom_kernel(const uint64_t* cache, const float* as, const float* bs, 
                   __hip_bfloat16* c, int m, int n, int k, int col, int offset) {
    const uint64_t* a = cache;
    const uint64_t* b = cache + offset;
    int fx = threadIdx.x >> 6 & 1;
    int fy = threadIdx.x >> 7 & 1;
    int cx = 32 * (2 * blockIdx.x | fx);
    int cy = 32 * (2 * blockIdx.y | fy);
    if(cx >= m || cy >= n) return;
    int tx = threadIdx.x & 31;
    int ty = threadIdx.x >> 5 & 1;
    int sn = (n + BLOCK - 1) >> 7;
    bs += cy >> 7;
    as += cx + tx;
    a += blockIdx.x * col | (threadIdx.x & 127);
    b += blockIdx.y * col | (threadIdx.x & 63) | (fy << 6);
    f32x16 result = { };
    for(int i = 0; i < (k + 127) >> 7; ++i) {
        const auto* ap = a + (i << 10);
        const auto* bp = b + (i << 10);
        f32x16 block_result = { };
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            block_result = __builtin_amdgcn_mfma_f32_32x32x16_fp8_fp8(bp[j * 0x80], ap[j * 0x80], block_result, 0, 0, 0);
        }
        result += block_result * bs[i * sn] * as[i * m];
    }
    const auto* src = reinterpret_cast<const uint16_t(*)[8]>(&result);
    uint64_t* dst = reinterpret_cast<uint64_t*>(c + (tx + cx) * n + cy + ty * 4);
    dst[0] = pack_u16x8_to_u64(src[0]);
    dst[2] = pack_u16x8_to_u64(src[1]);
    dst[4] = pack_u16x8_to_u64(src[2]);
    dst[6] = pack_u16x8_to_u64(src[3]);
}

__global__ void reorder_kernel(const __hip_fp8_e4m3_fnuz* src, const __hip_fp8_e4m3_fnuz* src1, uint64_t* dst, int m, int n, int k, int col, int next) {
    next = blockIdx.x >= next ? next : 0;
    m = next ? n : m;
    src = next ? src1 : src;
    int cx = (((blockIdx.x - next) * 2 + (threadIdx.x >> 6)) << 5) + (threadIdx.x & 31);
    int cy = (blockIdx.y * 2 + (threadIdx.x >> 5 & 1)) * 8;
    uint64_t value = 0;
    if (cx < m && cy < k) {
        src += cy * m + cx;
        __hip_fp8_e4m3_fnuz buffer[8] = { src[0], src[m], src[2 * m], src[3 * m],
        src[4 * m], src[5 * m], src[6 * m], src[7 * m]};
        value = *reinterpret_cast<uint64_t*>(buffer);
    }
    dst[blockIdx.x * col + blockIdx.y * 128 + threadIdx.x] = value;
}


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 next = (m + 127) / 128 * 2;
    int row = (n + 127) / 128 * 2 + next;
    int col = (k + 127) / 128 * 1024;
    auto options = torch::TensorOptions().dtype(torch::kUInt64).device(torch::kCUDA);
    torch::Tensor cache = torch::empty(row * col, options);
    reorder_kernel<<<dim3(row, col / 128), dim3(128), 0, 0>>>((__hip_fp8_e4m3_fnuz*)a.data_ptr(),  (__hip_fp8_e4m3_fnuz*)b.data_ptr(),
        cache.data_ptr<uint64_t>(), m, n, k, col, next);
    custom_kernel<<<dim3((m+63)/64, (n+63)/64), dim3(256), 0, 0>>> (cache.data_ptr<uint64_t>(),
        as.data_ptr<float>(), bs.data_ptr<float>(), (__hip_bfloat16*)c.data_ptr(), m, n, k, col, next * col);
}
"""

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", "-O3"],
)


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: 2
[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 -O3 -fno-gpu-rdc -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
[3/3] 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

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