#!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-05-27 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!

# This script provides a template for using load_inline to run a HIP kernel for
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 const int BlockSize = 32;

__global__ void custom_kernel(
    const __hip_fp8_e4m3_fnuz* A, 
    const __hip_fp8_e4m3_fnuz* B, 
    const float* scaleA, 
    const float* scaleB, 
    __hip_bfloat16* C, 
    int m, int n, int k
) {
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int bx = blockIdx.x;
    const int by = blockIdx.y;
    
    const unsigned int row = ty + blockDim.y * by;
    const unsigned int col = tx + blockDim.x * bx;
        
    if (row >= m || col >= n) return;

    const int num_b128 = (n + 127) / 128;

    __shared__ float smem_A[BlockSize][BlockSize];
    __shared__ float smem_B[BlockSize][BlockSize];
    
    float result = 0.0f;

    for (int tile = 0; tile < k; tile += BlockSize) {

        const int tileA_idx = tile + tx;
        const int tileB_idx = tile + ty;

        smem_A[ty][tx] = (row < m && (tileA_idx) < k) ? (float)A[row + (tileA_idx) * m] : 0.0f;
        smem_B[ty][tx] = (col < n && (tileB_idx) < k) ? (float)B[col + (tileB_idx) * n] : 0.0f;
        
        __syncthreads();  

        float block_result = 0.0f;
        
        #pragma unroll
        for (int i = 0; i < BlockSize; ++i) {
            block_result += smem_A[ty][i] * smem_B[i][tx];
        }
        
        const int k_grp = tile / 128;
        const int col_grp = col / 128;
        const float sA = scaleA[row + k_grp * m];
        const float sB = scaleB[col_grp + k_grp * num_b128];

        result += block_result * sA * sB;
        
        __syncthreads();  
    }
    
    if (row < m && col < n) {
        C[row * n + col] = (__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);
    
    dim3 block(BlockSize, BlockSize);         
    dim3 grid((n + BlockSize - 1) / BlockSize,    
              (m + BlockSize - 1) / BlockSize);      

    custom_kernel<<<grid, block>>>(
        (const __hip_fp8_e4m3_fnuz*)(a.data_ptr()),
        (const __hip_fp8_e4m3_fnuz*)(b.data_ptr()),
        as.data_ptr<float>(),
        bs.data_ptr<float>(),
        (__hip_bfloat16*)(c.data_ptr()),
        m, n, k);
    //C10_CUDA_CHECK(cudaGetLastError());
}
"""

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 eval.py leaderboard /tmp/tmpgzoon9wm
[92mSuccessfully preprocessed all matching files.[0m

/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
ninja: no work to do.
