#!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 64          // ← 减半，寄存器更少

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

/* -------------------------------- kernel */
__global__ void fp8_mm_fp32fma(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 ty  = threadIdx.y;            // 0..31
    int tx4 = threadIdx.x;            // 0..7  (4 cols)
    int row = blockIdx.y * TM + ty;
    int col0 = blockIdx.x * TN + tx4 * 4;
    if (row >= M) return;

    /* 双缓冲 shared */
    __shared__ __hip_fp8_e4m3_fnuz As[2][TM][TK];
    __shared__ __hip_fp8_e4m3_fnuz Bs[2][TK][TN];

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

    int bank = 0;
    /* 预取 tile-0 */
    for(int kk=tx4; kk<TK; kk+=8){
        int k = kk;
        As[bank][ty][kk] = (k < K)? A[row + k*M] : __hip_fp8_e4m3_fnuz{0};
    }
    for(int kk=ty; kk<TK; kk+=TM){
        int k = kk;
        int c0 = col0;
        Bs[bank][kk][tx4*4+0] = (c0   < N && k<K)? B[c0   + k*N]: __hip_fp8_e4m3_fnuz{0};
        Bs[bank][kk][tx4*4+1] = (c0+1 < N && k<K)? B[c0+1 + k*N]: __hip_fp8_e4m3_fnuz{0};
        Bs[bank][kk][tx4*4+2] = (c0+2 < N && k<K)? B[c0+2 + k*N]: __hip_fp8_e4m3_fnuz{0};
        Bs[bank][kk][tx4*4+3] = (c0+3 < N && k<K)? B[c0+3 + k*N]: __hip_fp8_e4m3_fnuz{0};
    }
    __syncthreads();

    for(int kb=0; kb<K; kb+=TK){
        /* 缩放 */
        int   kgrp = kb >> 7;                 // 128-K group index
        float sA  = AS[row + kgrp * M];
        int   kblk = kgrp;
        int cblk0 =  col0      /128;
        int cblk1 = (col0 + 1) /128;
        int cblk2 = (col0 + 2) /128;
        int cblk3 = (col0 + 3) /128;
        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;

        /* 计算 */
        #pragma unroll
        for(int kk=0; kk<TK; ++kk){
            float a = f8(As[bank][ty][kk]);
            acc0 += a * f8(Bs[bank][kk][tx4*4+0]) * sb0;
            acc1 += a * f8(Bs[bank][kk][tx4*4+1]) * sb1;
            acc2 += a * f8(Bs[bank][kk][tx4*4+2]) * sb2;
            acc3 += a * f8(Bs[bank][kk][tx4*4+3]) * sb3;
        }

        /* 预取下一 tile 到另一 bank */
        int next_kb   = kb + TK;
        int next_bank = bank ^ 1;
        __syncthreads();
        if(next_kb < K){
            for(int kk=tx4; kk<TK; kk+=8){
                int k = next_kb + kk;
                As[next_bank][ty][kk] = (k < K)? A[row + k*M] : __hip_fp8_e4m3_fnuz{0};
            }
            for(int kk=ty; kk<TK; kk+=TM){
                int k = next_kb + kk;
                int c0 = col0;
                Bs[next_bank][kk][tx4*4+0] = (c0   < N && k<K)? B[c0   + k*N]: __hip_fp8_e4m3_fnuz{0};
                Bs[next_bank][kk][tx4*4+1] = (c0+1 < N && k<K)? B[c0+1 + k*N]: __hip_fp8_e4m3_fnuz{0};
                Bs[next_bank][kk][tx4*4+2] = (c0+2 < N && k<K)? B[c0+2 + k*N]: __hip_fp8_e4m3_fnuz{0};
                Bs[next_bank][kk][tx4*4+3] = (c0+3 < N && k<K)? B[c0+3 + k*N]: __hip_fp8_e4m3_fnuz{0};
            }
        }
        __syncthreads();
        bank = next_bank;
    }

    __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;

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

    hipLaunchKernelGGL(fp8_mm_fp32fma, 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);

    hipError_t err = hipGetLastError();
    if (err != hipSuccess)
        AT_ERROR("kernel launch failed: ", hipGetErrorString(err));

    err = hipDeviceSynchronize();
    if (err != hipSuccess)
        AT_ERROR("kernel runtime error: ", hipGetErrorString(err));

}
"""

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

def custom_kernel(data: input_t) -> output_t:
    a,b,ascl,bscl,c = data
    module.fp8_mm(a,b,ascl,bscl,c)
    return c

python eval.py test /tmp/tmpc8u6ef6f
[92mSuccessfully preprocessed all matching files.[0m
