#!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 int BLOCK_SIZE_M = 64;
constexpr int BLOCK_SIZE_N = 64;
constexpr int BLOCK_SIZE_K = 128;

constexpr int TM = 4;
constexpr int TN = 4;

__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)
{
    __shared__ __hip_fp8_e4m3_fnuz As[BLOCK_SIZE_M][BLOCK_SIZE_K];
    __shared__ __hip_fp8_e4m3_fnuz Bs[BLOCK_SIZE_K][BLOCK_SIZE_N];
    __shared__ float scaleA_s[BLOCK_SIZE_M]; 
    __shared__ float scaleB_s[BLOCK_SIZE_N]; 

    const int tidx = threadIdx.x; 
    const int tidy = threadIdx.y; 

    const int block_row = blockIdx.x * BLOCK_SIZE_M;
    const int block_col = blockIdx.y * BLOCK_SIZE_N;

    const int row_base = block_row + tidx * TM; 
    const int col_base = block_col + tidy * TN;  

    float acc[TM][TN] = {0.0f};

    for (int tile_k = 0; tile_k < k; tile_k += BLOCK_SIZE_K){
        for (int i = tidy; i < BLOCK_SIZE_M; i += blockDim.y)
            for (int j = tidx; j < BLOCK_SIZE_K; j += blockDim.x){
                int g_row = block_row + i;
                int g_col = tile_k + j;
                As[i][j] = (g_row < m && g_col < k) ? A[g_row + g_col * m] : __hip_fp8_e4m3_fnuz(0.0f);
            }

        for (int i = tidy; i < BLOCK_SIZE_K; i += blockDim.y)
            for (int j = tidx; j < BLOCK_SIZE_N; j += blockDim.x){
                int g_row = block_col + j;  
                int g_col = tile_k + i;    
                Bs[i][j] = (g_row < n && g_col < k) ? B[g_row + g_col * n] : __hip_fp8_e4m3_fnuz(0.0f);
            }

        for (int i = tidy; i < BLOCK_SIZE_M; i += blockDim.y){
            int g_row = block_row + i;
            if (g_row < m){
                int k_blk = tile_k / BLOCK_SIZE_K;
                int scale_idx = g_row + k_blk * m;
                scaleA_s[i] = scaleA[scale_idx];
            }
        }

        const int n_blk_tot = (n + 127) / 128;
        for (int j = tidx; j < BLOCK_SIZE_N; j += blockDim.x){
            int g_col = block_col + j;
            if (g_col < n)
            {
                int k_blk = tile_k / BLOCK_SIZE_K;
                int n_blk = g_col / 128;
                int scale_idx = n_blk + k_blk * n_blk_tot;
                scaleB_s[j] = scaleB[scale_idx];
            }
        }

        __syncthreads();

        for (int kk = 0; kk < BLOCK_SIZE_K; ++kk){
            float a_frag[TM];
            #pragma unroll
            for (int ii = 0; ii < TM; ++ii){
                int r = tidx * TM + ii; 
                a_frag[ii] = static_cast<float>(As[r][kk]) * scaleA_s[r];
            }

            float b_frag[TN];
            #pragma unroll
            for (int jj = 0; jj < TN; ++jj){
                int c = tidy * TN + jj;  
                b_frag[jj] = static_cast<float>(Bs[kk][c]) * scaleB_s[c];
            }

            #pragma unroll
            for (int ii = 0; ii < TM; ++ii)
                #pragma unroll
                for (int jj = 0; jj < TN; ++jj)
                    acc[ii][jj] += a_frag[ii] * b_frag[jj];
        }

        __syncthreads();
    } 

    #pragma unroll
    for (int ii = 0; ii < TM; ++ii){
        int g_row = row_base + ii;
        if (g_row < m){
            #pragma unroll
            for (int jj = 0; jj < TN; ++jj){
                int g_col = col_base + jj;
                if (g_col < n)
                    C[g_row * n + g_col] = __hip_bfloat16(acc[ii][jj]);
            }
        }
    }
}

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(16, 16);
    dim3 grid((m + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M, 
              (n + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N);
    
    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/tmps47f853v
[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.
