#!POPCORN leaderboard amd-fp8-mm

# This is a submission template for popcorn leaderboard 'amd-fp8-mm'.
# Your task is as follows:
# >
# > You will implement a custom fp8-blockwise matmul kernel optimized for MI300.
# > You will be given single-precision scaling factors for your matrices.
# > The shapes of all outer and inner dimensions of tensors are from DeepSeek-R1.
# > To be explicit, you will be given a tuple of tensors:
# > ```
# > (a, b, a_scale, b_scale, c)
# > ```
# > where `a` and `b` are the input matrices, `a_scale` and `b_scale` are the scaling factors for `a` and `b` respectively,
# > and `c` is the output matrix:
# > * `a` is M x K in column-major order in e4m3fnuz
# > * `b` is N x K in column-major order in e4m3fnuz
# > * `a_scale` is M x K // 128 in column-major order in fp32
# > * `b_scale` is N // 128 x K // 128 in column-major order in fp32
# > * `c` is M x N in ROW-major order in bf16
# >
# > Matrix sizes `m` and `n` are divisible by 64, `k` is divisible by 128.
# >
# > The ranking criteria is the geometric mean of the benchmark results.
# >
# > For the grand price, your kernel will be evaluated against the speed of light analysis
# > and the solution closest to the speed of light will be awarded the grand price.
# > ```
# > The speed of light analysis is:
# >  M       N       K     time[us]
# > 1024    1536    7168      8.63
# > 1024    4608    7168     25.89
# > 6144    1536    7168     51.78
# > 6144    4608    7168    155.30
# > 1024    7168     256      3.17
# > 6144    7168     256     17.27
# > ```
# The deadline for this leaderboard is 2025-06-02 00:00:00+00:00

# You can automatically route this file to specific GPUs by adding a line
# `#!POPCORN gpus <GPUs>` to the header of this file.
# Happy hacking!

import os

os.environ["PYTORCH_ROCM_ARCH"] = "gfx942"

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);
"""

# ─────────────────────────────────────────────────────────────────────────────
# Optimized HIP kernel using 16×16 thread-blocks + 128-wide K-tiling in shared memory
# ─────────────────────────────────────────────────────────────────────────────
CUDA_SRC = """
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

constexpr int TILE_M = 16;        // tile size in M
constexpr int TILE_N = 16;        // tile size in N
constexpr int BLOCK_K = 128;      // tile size in K

// Each 16×16 block of threads cooperatively loads a 16×128 tile of A and B into shared memory.
// Then each thread performs a dot-product over 128 values in shared memory, applies scales, accumulates into 'result'.
__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
) {
    // Compute this thread's global row/col in C
    int bx = blockIdx.x;  // block index in M
    int by = blockIdx.y;  // block index in N
    int tx = threadIdx.x; // [0..15]
    int ty = threadIdx.y; // [0..15]
    int row = bx * TILE_M + tx;    // 0 ≤ row < m
    int col = by * TILE_N + ty;    // 0 ≤ col < n

    // Allocate shared memory for one 16×128 tile of A and one 16×128 tile of B
    __shared__ __hip_fp8_e4m3_fnuz shA[TILE_M][BLOCK_K];
    __shared__ __hip_fp8_e4m3_fnuz shB[TILE_N][BLOCK_K];

    // 'sn' = how many K-tiles along the K dimension
    int num_k_tiles = k / BLOCK_K;    // k is divisible by 128 by assumption
    int sn = (n + BLOCK_K - 1) / BLOCK_K; // used when indexing bs

    float result = 0.0f; // accumulation register

    // Early exit if this thread is outside the M×N range
    if (row >= m || col >= n) {
        return;
    }

    // Loop over each 128-wide K-tile
    for (int tile = 0; tile < num_k_tiles; ++tile) {
        int k_base = tile * BLOCK_K;         // start of this 128-wide chunk in K
        int a_scale_offset = tile * m;       // for indexing into as[]
        int b_scale_offset = tile * sn;      // for indexing into bs[]

        // -------------------------------------------------------------
        // 1) Cooperative load of A-tile (16×128) and B-tile (16×128) into shared memory
        // -------------------------------------------------------------
        // There are 16×16 = 256 threads. The 16×128 = 2048 elements of each tile should be loaded
        // exactly once per tile. So each thread will load 2048/256 = 8 elements of A + 8 elements of B.
        int tid = ty * TILE_M + tx; // unique [0..255] within block

        // Each thread will iterate j=0..7, computing:
        //   idx = tid + j*256,  idx ∈ [0..2047], then:
        //   rowA = idx % TILE_M,  colA = idx / TILE_M
        //   rowB = idx % TILE_N,  colB = idx / TILE_N
        for (int j = 0; j < (BLOCK_K * TILE_M) / (TILE_M * TILE_N); ++j) {
            // But since (BLOCK_K*TILE_M)/(TILE_M*TILE_N) = (128*16)/(16*16) = 8, loop j<8
            int idx = tid + j * (TILE_M * TILE_N); // [0..2047]
            int rowA = idx % TILE_M;                // [0..15]
            int colA = idx / TILE_M;                // [0..127]  (col within K-block)
            int rowB = idx % TILE_N;                // same as rowA
            int colB = idx / TILE_N;                // same as colA

            // Global indices for A and B:
            //   A  is M×K in column‐major: A[rowA_global + (k_base+colA)*m]
            //   B  is N×K in column‐major: B[rowB_global + (k_base+colB)*n]
            int rowA_global = bx * TILE_M + rowA;  // [0..m-1]
            int rowB_global = by * TILE_N + rowB;  // [0..n-1]

            // Only load if within bounds (rowA_global < m, rowB_global < n)
            if (rowA_global < m) {
                shA[rowA][colA] = a[rowA_global + (k_base + colA) * m];
            } else {
                // if out of bounds in M, write zero (so it won't affect accumulation)
                shA[rowA][colA] = static_cast<__hip_fp8_e4m3_fnuz>(0);
            }
            if (rowB_global < n) {
                shB[rowB][colB] = b[rowB_global + (k_base + colB) * n];
            } else {
                shB[rowB][colB] = static_cast<__hip_fp8_e4m3_fnuz>(0);
            }
        }

        __syncthreads();

        // -------------------------------------------------------------
        // 2) Each thread performs a 128-wide dot‐product in shared memory
        //    for A[row, k_base..k_base+127] × B[col, k_base..k_base+127]
        // -------------------------------------------------------------
        float partial_sum = 0.0f;
        for (int kk = 0; kk < BLOCK_K; ++kk) {
            float av = static_cast<float>( shA[tx][kk] );   // A[row][k_base+kk]
            float bv = static_cast<float>( shB[ty][kk] );   // B[col][k_base+kk]
            partial_sum += av * bv;
        }

        // -------------------------------------------------------------
        // 3) Apply the FP32 scales for this 128-tile, accumulate into 'result'
        // -------------------------------------------------------------
        // a_scale index: as[row + (tile * m)]
        // b_scale index: bs[(col / BLOCK_K) + tile * sn]
        float scaleA = as[row + a_scale_offset]; 
        int col_block = col / BLOCK_K; 
        float scaleB = bs[col_block + b_scale_offset];

        result += partial_sum * scaleA * scaleB;

        __syncthreads();
    }

    // -------------------------------------------------------------
    // 4) Finally, convert to BF16 and write into C (row-major M×N)
    // -------------------------------------------------------------
    // Each thread writes c[row * n + col].
    c[row * n + col] = static_cast<__hip_bfloat16>(result);
}

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

    // Launch grid of ((m+15)/16, (n+15)/16) blocks, each 16×16 threads
    dim3 grid( (m + TILE_M - 1) / TILE_M, 
               (n + TILE_N - 1) / TILE_N );
    dim3 block(TILE_M, TILE_N);

    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
    );
    // (Optional) check for launch errors:
    // hipError_t err = hipGetLastError(); 
    // if (err != hipSuccess) { /* handle error */ }
}
"""

import os

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

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
    module.fp8_mm(a, b, a_scale, b_scale, c)
    return c

python submission.py
[92mSuccessfully preprocessed all matching files.[0m

/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]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [ok]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
[1/3] clang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=fp8_mm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/main.cpp -o main.o -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC
[2/3] /opt/rocm-6.3.1/bin/hipcc  -DWITH_HIP -DTORCH_EXTENSION_NAME=fp8_mm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC -DCUDA_HAS_FP16=1 -D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1 -DHIP_ENABLE_WARP_SYNC_BUILTINS=1 --offload-arch=gfx942 -fno-gpu-rdc --offload-arch=gfx942 -std=c++20 -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
[3/3] clang++ hip.cuda.o main.o -shared -L/usr/local/lib/python3.10/dist-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

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

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


Total number of replaced kernel launches: 1
ninja: no work to do.
