# template_hip.py
# Inline HIP implementation of blockwise FP8 GEMM with block-scale factors
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
import os

# Use clang++ for HIP compilation
os.environ["CXX"] = "clang++"

# C++ wrapper declaration
CPP_WRAPPER = """
#include <torch/extension.h>
// Declare the GPU entry-point for Python binding
void fp8_mm(
    at::Tensor A,
    at::Tensor B,
    at::Tensor alpha,
    at::Tensor beta,
    at::Tensor C
);
"""

# HIP source: kernel + definition
HIP_SRC = """
#include <torch/extension.h>
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>
constexpr int BLOCK = 128;

// Blockwise FP8 GEMM with per-block scales
__global__ void custom_kernel(
    const __hip_fp8_e4m3_fnuz* A,
    const __hip_fp8_e4m3_fnuz* B,
    const float* alpha,
    const float* beta,
    __hip_bfloat16* C,
    int M, int N, int K
) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    int col = blockIdx.y * blockDim.y + threadIdx.y;
    if (row >= M || col >= N) return;
    int kb = (K + BLOCK - 1) / BLOCK;
    int nb = (N + BLOCK - 1) / BLOCK;
    float acc = 0.0f;
    for (int b = 0; b < kb; ++b) {
        int k0 = b * BLOCK;
        float ascl = alpha[row + b * M];
        float bscl = beta[(col / BLOCK) + b * nb];
        float blk = 0.0f;
        #pragma unroll
        for (int i = 0; i < BLOCK; ++i) {
            int ki = k0 + i;
            if (ki >= K) break;
            float av = (float)A[row + ki * M];
            float bv = (float)B[col + ki * N];
            blk += av * bv;
        }
        acc += blk * ascl * bscl;
    }
    C[row * N + col] = (__hip_bfloat16)acc;
}

// C++ wrapper implementing the fp8_mm symbol
void fp8_mm(
    at::Tensor A,
    at::Tensor B,
    at::Tensor alpha,
    at::Tensor beta,
    at::Tensor C
) {
    int M = A.size(0), N = B.size(0), K = A.size(1);
    dim3 blk(16, 16), grd((M + 15) / 16, (N + 15) / 16);
    custom_kernel<<<grd, blk>>>(
        (const __hip_fp8_e4m3_fnuz*)A.data_ptr(),
        (const __hip_fp8_e4m3_fnuz*)B.data_ptr(),
        alpha.data_ptr<float>(),
        beta.data_ptr<float>(),
        (__hip_bfloat16*)C.data_ptr(),
        M, N, K
    );
    // Note: explicit error-check removed for compatibility
}
"""

# Build with inline compilation
module = load_inline(
    name="fp8_mm",
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[HIP_SRC],
    functions=["fp8_mm"],
    verbose=True,
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20"],
)

# Python entrypoint


def custom_kernel(data: input_t) -> output_t:
    A, B, alpha, beta, C = data
    module.fp8_mm(A, B, alpha, beta, 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: 1
[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 
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:12:17: warning: loop not unrolled: the optimizer was unable to perform the requested transformation; the transformation might be disabled or specified as part of an unsupported transformation ordering [-Wpass-failed=transform-warning]
   12 | __global__ void custom_kernel(
      |                 ^
1 warning generated when compiling for gfx942.
[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 benchmark /tmp/tmpm75xw0g9
[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++ 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
