#!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++")

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>

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

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

__global__ void fp8_gemm_vec4(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,
                              int scale_n, int scale_k)
{
    int ty  = threadIdx.y;   // 0..31
    int tx4 = threadIdx.x;   // 0..7   (4 cols per thread)

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

    if (row >= M) return;

    __shared__ __hip_fp8_e4m3_fnuz As[TM][TK];
    __shared__ __hip_fp8_e4m3_fnuz BsTile[TK][TN];   // renamed to avoid clash

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

    for (int kb = 0; kb < K; kb += TK) {

        /* ---- load A to shared ---- */
        #pragma unroll
        /* ---- load A tile 32×128 (完整) ---- */
        for (int kk = tx4; kk < TK; kk += 8) {          // ★ 步长 8 覆盖全部列
            int k = kb + kk;
            As[ty][kk] = (k < K) ? A[row + k * M] : __hip_fp8_e4m3_fnuz{0};
        }

        /* ---- load B to shared ---- */
        #pragma unroll
        for (int kk = ty; kk < TK; kk += TM) {
            int k = kb + kk;
            int c0 = col0;
            if (k < K) {
                BsTile[kk][tx4*4+0] = (c0     < N)? B[c0     + k * N] : __hip_fp8_e4m3_fnuz{0};
                BsTile[kk][tx4*4+1] = (c0+1   < N)? B[c0 + 1 + k * N] : __hip_fp8_e4m3_fnuz{0};
                BsTile[kk][tx4*4+2] = (c0+2   < N)? B[c0 + 2 + k * N] : __hip_fp8_e4m3_fnuz{0};
                BsTile[kk][tx4*4+3] = (c0+3   < N)? B[c0 + 3 + k * N] : __hip_fp8_e4m3_fnuz{0};
            }
        }
        __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;

        /* β(rowBlk, kBlk) ，列主序展开：rowBlk + kBlk * scale_n */
        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 = f8(As[ty][kk]);
            acc0 += a * f8(BsTile[kk][tx4*4+0]) * sb0;
            acc1 += a * f8(BsTile[kk][tx4*4+1]) * sb1;
            acc2 += a * f8(BsTile[kk][tx4*4+2]) * sb2;
            acc3 += a * f8(BsTile[kk][tx4*4+3]) * sb3;
        }
        __syncthreads();
    }

    __hip_bfloat16* dst = reinterpret_cast<__hip_bfloat16*>(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;
    int scale_k = (K + 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 __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, scale_n, scale_k);
}
"""

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"],
)

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/tmp3xeiertf
[92mSuccessfully preprocessed all matching files.[0m
