#!POPCORN leaderboard amd-fp8-mm


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/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

constexpr int BLOCKSIZE     = 128;     // still your K-chunk
constexpr int TILE_HEIGHT   =   8;     // tile height M (rows of C per block)
constexpr int TILE_WIDTH    =   8;     // tile width  N (cols of C per block)

__global__ void custom_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*          __restrict__ c,
    int m, int n, int k
) {
    int tx = threadIdx.x;    // 0..7
    int ty = threadIdx.y;    // 0..7

    int row = blockIdx.x * TILE_HEIGHT + tx;  
    int col = blockIdx.y * TILE_WIDTH + ty;  

    if (row >= m || col >= n) return;

    int numBlocksInK = k / BLOCKSIZE;   // assume k % BLOCKSIZE == 0
    float acc = 0.0f;

    int colBlock = col / BLOCKSIZE;  

    for (int chunk = 0; chunk < numBlocksInK; ++chunk) {
        int base_k = chunk * BLOCKSIZE;

        float sa = as[row * numBlocksInK + chunk];
        float sb = bs[colBlock * numBlocksInK + chunk];

        // in each chunk, do a simple scalar dot over 128 elements
        for (int t = 0; t < BLOCKSIZE; ++t) {
            float a_v = float(a[row * k + base_k + t]) * sa;
            float b_v = float(b[col * k + base_k + t]) * sb;
            acc += a_v * b_v;
        }
    }

    c[row * n + col] = __hip_bfloat16(acc);
}

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);
    dim3 block(TILE_HEIGHT, TILE_WIDTH);
    dim3 grid((m + TILE_HEIGHT - 1) / TILE_HEIGHT,
              (n + TILE_WIDTH - 1) / TILE_WIDTH);

    custom_kernel<<<grid, block>>>(
        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
    );
    // C10_CUDA_CHECK(cudaGetLastError());
}
"""

import os
os.environ["CXX"] = "clang++"

# Compile the kernel
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

    # Convert FROM column-major → true C row-major
    # .contiguous() always returns a C-contiguous copy
    a_rm      = a.contiguous()         # M×K
    b_rm      = b.contiguous()         # N×K
    as_rm     = a_scale.contiguous()   # M×(K/128)
    bs_rm     = b_scale.contiguous()   # (N/128)×(K/128)

    # Launch our kernel
    module.fp8_mm(a_rm, b_rm, as_rm, bs_rm, c)
    return c

python eval.py leaderboard /tmp/tmp_sg_bn6s
Using /home/runner/.cache/torch_extensions/py310_cpu as PyTorch extensions root...
[92mSuccessfully preprocessed all matching files.[0m
Detected CUDA files, patching ldflags
Emitting ninja build file /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/build.ninja...
Building extension module fp8_mm...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
Loading extension module fp8_mm...

/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]
/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]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
[1/1] clang++ main.o hip.cuda.o -shared -L/home/runner/.local/lib/python3.10/site-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
