#!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_fp16.h>
#include <hip/hip_fp8.h>                 // defines __hip_fp8_e4m3_fnuz + cast
#include <hip/amd_detail/amd_hip_bf16.h> // __hip_bfloat16

#include <ATen/ATen.h>
#include <ATen/hip/HIPContext.h>
#include <torch/extension.h>

// ----------- launch parameters (debug not perf) --------------
constexpr int BLOCK_M = 16;
constexpr int BLOCK_N = 16;

// ----------------------- kernel ------------------------------
extern "C" __global__
void custom_kernel(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 * BLOCK_M + threadIdx.y; // 0..M-1
    int col = blockIdx.x * BLOCK_N + threadIdx.x; // 0..N-1
    if (row >= M || col >= N) return;

    const int lda = M; // column-major
    const int ldb = N;

    float acc = 0.f;
        const int kBlocks = K / 128;               // K is guaranteed divisible by 128
    const int nBlocks = (N + 127) / 128;       // **ceil** so tail-cols get a scale row

    for (int kb = 0; kb < kBlocks; ++kb) {
        float a_s = Ascale[row + kb * M];
        float b_s = Bscale[(col / 128) + kb * nBlocks];

#pragma unroll
        for (int kk = 0; kk < 128; ++kk) {
            int k = kb * 128 + kk;
            __hip_fp8_e4m3_fnuz a_fp8 = A[row + (size_t)k * lda];
            __hip_fp8_e4m3_fnuz b_fp8 = B[col + (size_t)k * ldb];
            float a_f32 = static_cast<float>(a_fp8) * a_s;
            float b_f32 = static_cast<float>(b_fp8) * b_s;
            acc += a_f32 * b_f32;
        }
    }
    C[(size_t)row * N + col] = __float2bfloat16(acc);
}

// -------------------- host launcher --------------------------
static void LaunchGemm(const torch::Tensor& a, const torch::Tensor& b,
                       const torch::Tensor& as, const torch::Tensor& bs,
                       torch::Tensor& c)
{
    int M = a.size(0), K = a.size(1), N = b.size(0);
    hipStream_t s = at::hip::getCurrentHIPStream().stream();

    dim3 block(BLOCK_N, BLOCK_M);
    dim3 grid((N + BLOCK_N - 1)/BLOCK_N, (M + BLOCK_M - 1)/BLOCK_M);

    custom_kernel<<<grid, block, 0, s>>>(
        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);
}

// -------------------- PyTorch wrapper ------------------------
void fp8_mm(torch::Tensor a, torch::Tensor b,
            torch::Tensor as, torch::Tensor bs,
            torch::Tensor c)
{
    TORCH_CHECK(as.scalar_type()==at::kFloat && bs.scalar_type()==at::kFloat,
                "scale tensors must be fp32");
    TORCH_CHECK(c.scalar_type()==at::kBFloat16, "C must be bf16");

    TORCH_CHECK(a.is_cuda() && b.is_cuda() && as.is_cuda() && bs.is_cuda() && c.is_cuda(),
                "all tensors must be on same CUDA/HIP device");

    TORCH_CHECK(a.dim()==2 && b.dim()==2 && c.dim()==2);
    int64_t M=a.size(0), K=a.size(1), N=b.size(0);
    TORCH_CHECK(b.size(1)==K && c.size(0)==M && c.size(1)==N,
                "dimension mismatch");
    TORCH_CHECK(K%128==0, "K must be multiple 128");

    LaunchGemm(a,b,as,bs,c);
}

"""

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 benchmark /tmp/tmp_1n_y5or
[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
