#!POPCORN leaderboard amd-fp8-mm
import os, torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

os.environ.setdefault("CXX", "clang++")       # hipcc backend

# ---------- C++/HIP sources ----------
CPP_SRC = r"""
void fp8_mm(torch::Tensor a, torch::Tensor b,
            torch::Tensor as, torch::Tensor bs,
            torch::Tensor c);
"""

HIP_SRC = r"""
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

constexpr int BLOCK_M = 16;     // blockDim.x
constexpr int BLOCK_N = 16;     // blockDim.y
constexpr int BLOCK_K = 128;    // tile depth / scaling granularity

__device__ inline float fp8_to_f32(__hip_fp8_e4m3_fnuz v) {
    return static_cast<float>(v);
}

__global__ void kernel(const __hip_fp8_e4m3_fnuz* __restrict__ A,
                       const __hip_fp8_e4m3_fnuz* __restrict__ B,
                       const float*  __restrict__ AS,   // α
                       const float*  __restrict__ BS,   // β
                       __hip_bfloat16*             C,
                       int M, int N, int K)
{
    /* --- global coordinates --- */
    int row = blockIdx.x * BLOCK_M + threadIdx.x;   // 0 … M-1
    int col = blockIdx.y * BLOCK_N + threadIdx.y;   // 0 … N-1
    if (row >= M || col >= N) return;

    const int scale_n = (N + BLOCK_K - 1) / BLOCK_K;   // = N / 128

    /* shared tiles: A 16×128, B 128×16 */
    __shared__ __hip_fp8_e4m3_fnuz As[BLOCK_M][BLOCK_K];
    __shared__ __hip_fp8_e4m3_fnuz Bs[BLOCK_K][BLOCK_N];

    float acc = 0.f;

    /* --- loop over K tiles --- */
    for (int kb = 0; kb < K; kb += BLOCK_K) {

        /* ===== 1. load A -> shared (16×128) ===== */
        #pragma unroll
        for (int kk = threadIdx.y; kk < BLOCK_K; kk += BLOCK_N) {
            int k_global = kb + kk;
            As[threadIdx.x][kk] =
                (row < M) ? A[row + k_global * M] : __hip_fp8_e4m3_fnuz{0};
        }

        /* ===== 2. load B -> shared (128×16) ===== */
        #pragma unroll
        for (int kk = threadIdx.x; kk < BLOCK_K; kk += BLOCK_M) {
            int k_global = kb + kk;
            Bs[kk][threadIdx.y] =
                (col < N) ? B[col + k_global * N] : __hip_fp8_e4m3_fnuz{0};
        }

        __syncthreads();

        /* ===== 3. α-β scaling constants for this tile ===== */
        float sA    = AS[row + (kb / BLOCK_K) * M];               // α(row, kblk)
        float sB    = BS[(col / BLOCK_K) + (kb / BLOCK_K) * scale_n]; // β(colBlk,kblk)
        float sBoth = sA * sB;

        /* ===== 4. compute inner product from shared ===== */
        #pragma unroll
        for (int kk = 0; kk < BLOCK_K; ++kk) {
            float a = fp8_to_f32(As[threadIdx.x][kk]);
            float b = fp8_to_f32(Bs[kk][threadIdx.y]);
            acc += a * b * sBoth;
        }
        __syncthreads();          // reuse shared mem next iteration
    }

    /* --- write C (row-major) --- */
    C[row * N + col] = (__hip_bfloat16)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);
    int N = b.size(0);
    int K = a.size(1);

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

    hipLaunchKernelGGL(kernel, grid, block, 0, 0,
        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()),
        M, N, K);
}
"""

# ---------- build inline module ----------
module = load_inline(
    name         = "fp8_mm_step2",
    cpp_sources  = [CPP_SRC],
    cuda_sources = [HIP_SRC],
    functions    = ["fp8_mm"],
    verbose      = False,
    extra_cuda_cflags=["--offload-arch=gfx942", "-O3", "-std=c++20"],
)

# ---------- python entry ----------
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/tmpipw1sxfe
[92mSuccessfully preprocessed all matching files.[0m
