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

# ----------------------------------------------------------------------
#  1) 只编译一次
# ----------------------------------------------------------------------
os.environ.setdefault("CXX", "clang++")

CPP_SRC = r"""
void fp8_mm(torch::Tensor a, torch::Tensor b,
            torch::Tensor as, torch::Tensor bs,
            torch::Tensor c);
"""

# ----------------------------------------------------------------------
#  2) HIP kernel：FP8 存 uint8_t，运行时解码
# ----------------------------------------------------------------------
HIP_SRC = r"""
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_fp8.h>     // 只拿枚举，用不到对象
#include <hip/amd_detail/amd_hip_bf16.h>

#define TM 32
#define TN 32
#define TK 128

/* ---- ① fp8-bits → float32 ------------------------------------------------*/
__device__ __forceinline__
float fp8_to_f32(uint8_t v)
{
#if defined(__has_builtin)
    /* ──① 优先用最精确指令 (ROCm≥6.4/gfx11) ─────────────────────────── */
#   if __has_builtin(__builtin_amdgcn_cvt_f32_fp8_e4m3)
        return __builtin_amdgcn_cvt_f32_fp8_e4m3(v, /*sat=*/false);
    /* ──② 部分 6.3 toolchain 只有通用版本 ──────────────────────────── */
    /* ──② 通用指令：所有 ROCm 版本都有的 2 参重载 ───────────────── */
#   elif __has_builtin(__builtin_amdgcn_cvt_f32_fp8)
        return __builtin_amdgcn_cvt_f32_fp8(v, /*sat=*/false);
#   endif
#endif
    /* ──③ 没有硬件指令就用软件解码（任何架构都能过编译）──────────── */
    uint32_t s = (uint32_t)(v & 0x80u) << 24;      // sign
    int      e = (v >> 3) & 0x0f;                  // exponent
    uint32_t m =  v & 0x07u;                       // mantissa
    if(e == 0){
        if(m == 0)  return __uint_as_float(s);     // ±0
        do{ m <<= 1; --e; } while((m & 0x08u) == 0);
        m &= 0x07u;
    }
    uint32_t exp32 = (uint32_t)(e + (127 - 7)) << 23;
    uint32_t man32 = m << 20;                      // 23-3
    return __uint_as_float(s | exp32 | man32);
}

/* ---- ② kernel ------------------------------------------------------------*/
extern "C"
__global__ void fp8_gemm_vec4(
        const uint8_t* __restrict__ A,   // [M, K] col-major
        const uint8_t* __restrict__ B,   // [N, K] row-major
        const float*   __restrict__ AS,  // [M, K//128]
        const float*   __restrict__ BS,  // [N//128, K//128]
        __hip_bfloat16*      C,          // [M, N] row-major (bf16)
        int  M, int N, int K,
        int  scale_n)
{
    int ty  = threadIdx.y;          // 0..31
    int tx4 = threadIdx.x;          // 0..7  (4 cols/thread)

    int row  = blockIdx.y * TM + ty;
    int col0 = blockIdx.x * TN + tx4 * 4;

    if (row >= M) return;

    __shared__ uint8_t As[TM][TK];
    __shared__ uint8_t Bs[TK][TN];

    float acc0 = 0, acc1 = 0, acc2 = 0, acc3 = 0;

    for(int kb = 0; kb < K; kb += TK)
    {
        /* ---- load A (col-major) to shared ---- */
        #pragma unroll
        for(int kk = tx4; kk < TK; kk += 8){
            int k = kb + kk;
            As[ty][kk] = (k < K) ? A[row + k * M] : 0u;
        }

        /* ---- load B (row-major) ---- */
        #pragma unroll
        for(int kk = ty; kk < TK; kk += TM){
            int k = kb + kk;
            int c0 = col0;
            if(k < K){
                Bs[kk][tx4*4+0] = (c0   < N)? B[c0   + k * N] : 0u;
                Bs[kk][tx4*4+1] = (c0+1 < N)? B[c0+1 + k * N] : 0u;
                Bs[kk][tx4*4+2] = (c0+2 < N)? B[c0+2 + k * N] : 0u;
                Bs[kk][tx4*4+3] = (c0+3 < N)? B[c0+3 + k * N] : 0u;
            }
        }
        __syncthreads();

        /* ---- scaling ---- */
        float sA  = AS[row + (kb/TK) * M];
        int   kblk = kb / TK;

        int cblk0 = col0     / TK;
        int cblk1 = (col0+1) / TK;
        int cblk2 = (col0+2) / TK;
        int cblk3 = (col0+3) / TK;

        float sb0 = (col0   < N)? sA * BS[cblk0 + kblk * scale_n] : 0.f;
        float sb1 = (col0+1 < N)? sA * BS[cblk1 + kblk * scale_n] : 0.f;
        float sb2 = (col0+2 < N)? sA * BS[cblk2 + kblk * scale_n] : 0.f;
        float sb3 = (col0+3 < N)? sA * BS[cblk3 + kblk * scale_n] : 0.f;

        /* ---- inner product ---- */
        #pragma unroll
        for(int kk = 0; kk < TK; ++kk){
            float a  = fp8_to_f32(As[ty][kk]);
            float b0 = fp8_to_f32(Bs[kk][tx4*4+0]);
            float b1 = fp8_to_f32(Bs[kk][tx4*4+1]);
            float b2 = fp8_to_f32(Bs[kk][tx4*4+2]);
            float b3 = fp8_to_f32(Bs[kk][tx4*4+3]);

            acc0 += a * b0 * sb0;
            acc1 += a * b1 * sb1;
            acc2 += a * b2 * sb2;
            acc3 += a * b3 * sb3;
        }
        __syncthreads();
    }

    __hip_bfloat16* dst = C + row * N + col0;
    if(col0   < N) dst[0] = (__hip_bfloat16)acc0;
    if(col0+1 < N) dst[1] = (__hip_bfloat16)acc1;
    if(col0+2 < N) dst[2] = (__hip_bfloat16)acc2;
    if(col0+3 < N) dst[3] = (__hip_bfloat16)acc3;
}

/* ---- ③ 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), N = b.size(0), K = a.size(1);
    int scale_n = (N + 127) / 128;

    dim3 block(8, 32);                        // 256 threads
    dim3 grid((N + TN - 1) / TN,
              (M + TM - 1) / TM);

    hipLaunchKernelGGL(fp8_gemm_vec4, grid, block, 0, 0,
        reinterpret_cast<const uint8_t*>(a.data_ptr()),      // ★
        reinterpret_cast<const uint8_t*>(b.data_ptr()),      // ★
        as.data_ptr<float>(),
        bs.data_ptr<float>(),
        reinterpret_cast<__hip_bfloat16*>(c.data_ptr()),
        M, N, K, scale_n);

    hipDeviceSynchronize();   // error check 由 PyTorch 负责
}
"""

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

# ----------------------------------------------------------------------
#  3) Popcorn 入口：确保张量在 GPU，避免多余拷贝
# ----------------------------------------------------------------------
def custom_kernel(data: input_t) -> output_t:
    # task.py 里生成的是 CPU tensor；这里一次性搬到 HIP 侧
    a, b, a_scale, b_scale, c = [t.to("cuda", non_blocking=True) for t in data]

    module.fp8_mm(a, b, a_scale, b_scale, c)
    return c

python eval.py benchmark /tmp/tmpfo20ks3g
[92mSuccessfully preprocessed all matching files.[0m
