#!POPCORN leaderboard amd-fp8-mm

# This is a submission template for popcorn leaderboard 'amd-fp8-mm'.
# Your task is as follows:
# > 
# > You will implement a custom fp8-blockwise matmul kernel optimized for MI300.
# > You will be given single-precision scaling factors for your matrices.
# > The shapes of all outer and inner dimensions of tensors are from DeepSeek-R1.
# > To be explicit, you will be given a tuple of tensors:
# > ```
# > (a, b, a_scale, b_scale, c)
# > ```
# > where `a` and `b` are the input matrices, `a_scale` and `b_scale` are the scaling factors for `a` and `b` respectively,
# > and `c` is the output matrix:
# > * `a` is M x K in column-major order in e4m3fnuz
# > * `b` is N x K in column-major order in e4m3fnuz
# > * `a_scale` is M x K // 128 in column-major order in fp32
# > * `b_scale` is N // 128 x K // 128 in column-major order in fp32
# > * `c` is M x N in ROW-major order in bf16
# > 
# > Matrix sizes `m` and `n` are divisible by 64, `k` is divisible by 128.
# > 
# > The ranking criteria is the geometric mean of the benchmark results.
# > 
# > For the grand price, your kernel will be evaluated against the speed of light analysis
# > and the solution closest to the speed of light will be awarded the grand price.
# > ```
# > The speed of light analysis is:
# >  M       N       K     time[us]
# > 1024    1536    7168      8.63
# > 1024    4608    7168     25.89
# > 6144    1536    7168     51.78
# > 6144    4608    7168    155.30
# > 1024    7168     256      3.17
# > 6144    7168     256     17.27
# > ```
# The deadline for this leaderboard is 2025-05-27 00:00:00+00:00

# You can automatically route this file to specific GPUs by adding a line
# `#!POPCORN gpus <GPUs>` to the header of this file.
# Happy hacking!

# 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/hip_runtime.h>
#include <hip/hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>
#include <ATen/ATen.h>
#include <ATen/hip/HIPContext.h>
#include <torch/extension.h>

// ------------ FP8 (e4m3fnuz) to float32 helper ------------
__device__ __forceinline__ float fp8_to_fp32(uint8_t v) {
#if defined(__has_builtin)
 #if __has_builtin(__builtin_amdgcn_cvt_f32_fp8)
    return __builtin_amdgcn_cvt_f32_fp8((int)v, /*fnuz=*/1);
 #endif
#endif
    const uint32_t sign = (v & 0x80u) << 24;
    const uint32_t e4   = (v >> 3) & 0x0Fu;
    const uint32_t m3   =  v & 0x07u;
    if (e4 == 0)     return __uint_as_float(sign);
    if (e4 == 0x0Fu) return __uint_as_float(sign | 0x7fc00000u);
    const uint32_t exp32 = (e4 + (127 - 7)) << 23;
    const uint32_t man32 = m3 << 20;
    return __uint_as_float(sign | exp32 | man32);
}

// ---------------- Reference GEMM kernel -------------------
constexpr int BLK = 16;
extern "C" __global__
void reference_gemm(const __hip_fp8_e4m3_fnuz* __restrict__ A,
                    const __hip_fp8_e4m3_fnuz* __restrict__ B,
                    const float*               __restrict__ Ascale,
                    const float*               __restrict__ Bscale,
                    __hip_bfloat16*            __restrict__ C,
                    int M, int N, int K) {
    int row = blockIdx.y * BLK + threadIdx.y;
    int col = blockIdx.x * BLK + threadIdx.x;
    if (row >= M || col >= N) return;

    const int lda = M;
    const int ldb = N;
    float acc = 0.f;

    int kChunks = K / 128;
    int nBlocks = (N + 127) / 128;

    for (int kc = 0; kc < kChunks; ++kc) {
        float a_s = Ascale[row + kc * M];
        float b_s = Bscale[(col / 128) + kc * nBlocks];
#pragma unroll 4
        for (int kk = 0; kk < 128; ++kk) {
            int kIdx = kc * 128 + kk;
            __hip_fp8_e4m3_fnuz av = A[row + (size_t)kIdx * lda];
            __hip_fp8_e4m3_fnuz bv = B[col + (size_t)kIdx * ldb];
            float a_v = static_cast<float>(av) * a_s;
            float b_v = static_cast<float>(bv) * b_s;
            acc += a_v * b_v;
        }
    }
    C[(size_t)row * N + col] = __float2bfloat16(acc);
}

// ---------------- Host wrapper ----------------------------
void fp8_mm(torch::Tensor a, torch::Tensor b,
            torch::Tensor as, torch::Tensor bs,
            torch::Tensor c) {

    int M = a.size(0), K = a.size(1), N = b.size(0);

    dim3 grid((N + BLK - 1) / BLK, (M + BLK - 1) / BLK);
    dim3 block(BLK, BLK);
    hipStream_t stream = at::hip::getCurrentHIPStream().stream();
    reference_gemm<<<grid, block, 0, stream>>>(
        reinterpret_cast<const __hip_fp8_e4m3_fnuz*>(a.data_ptr()),
        reinterpret_cast<const __hip_fp8_e4m3_fnuz*>(b.data_ptr()),
        as.data_ptr<float>(), bs.data_ptr<float>(),
        reinterpret_cast<__hip_bfloat16*>(c.data_ptr<at::BFloat16>()),
        M, N, K);
}
"""

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: 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 
[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/tmppjqzcuvg
[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.
