#!POPCORN leaderboard amd-fp8-mm
#!POPCORN gpu MI300

import os

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

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 WARP_SIZE = 64;
constexpr const int BLOCK = 128;
using fp8 = __hip_fp8_e4m3_fnuz;
using bf16 = __hip_bfloat16;
using fp32x4 = __attribute__((__vector_size__(4 * sizeof(float)))) float;

__host__ __device__ constexpr int cdiv(int a, int b) { return (a + b - 1) / b; }

template <int BLOCK_M, int BLOCK_N, int WARP_M, int WARP_N>
__global__ void custom_kernel(
    const fp8* __restrict__ A,  // (K, M)
    const fp8* __restrict__ B,  // (K, N)
    const float* __restrict__ As,  // (K/128, M)
    const float* __restrict__ Bs,  // (K/128, N/128)
    bf16* __restrict__ C,  // (M, N)
    int M,
    int N,
    int K
) {
    constexpr int BLOCK_K = 32;
    static_assert(BLOCK % BLOCK_K == 0);
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int grid_width = cdiv(N, BLOCK_N);
    const int bid_m = bid / grid_width;
    const int bid_n = bid % grid_width;

    constexpr int NUM_WARP_M = BLOCK_M / WARP_M;
    constexpr int NUM_WARP_N = BLOCK_N / WARP_N;
    constexpr int TB_SIZE = NUM_WARP_M * NUM_WARP_N * WARP_SIZE;
    const int lane_id = tid % WARP_SIZE;
    const int warp_id = tid / WARP_SIZE;
    const int warp_id_m = warp_id / NUM_WARP_N;
    const int warp_id_n = warp_id % NUM_WARP_N;

    const int offm = bid_m * BLOCK_M;
    const int offn = bid_n * BLOCK_N;
    A += offm;
    B += offn;
    As += offm + warp_id_m * WARP_M;
    Bs += (offn + warp_id_n * WARP_N) / BLOCK;
    C += (offm + warp_id_m * WARP_M) * N + (offn + warp_id_n * WARP_N);

    fp32x4 acc1[WARP_M / 16][WARP_N / 16] = {};
    fp32x4 acc2[WARP_M / 16][WARP_N / 16];
    __shared__ fp8 A_shared[BLOCK_M * BLOCK_K];
    __shared__ fp8 B_shared[BLOCK_N * BLOCK_K];

    for(int i = 0; i < K; i += BLOCK) {
        // reset acc2
        for (int mi = 0; mi < WARP_M / 16; ++mi)
            for (int ni = 0; ni < WARP_N / 16; ++ni)
                for (int item = 0; item < 4; ++item)
                    acc2[mi][ni][item] = 0.0f;

        for (int ii = 0; ii < BLOCK; ii += BLOCK_K) {
            // TODO: load with wider bit-width
            for (int idx = tid; idx < BLOCK_K * BLOCK_M; idx += TB_SIZE) {
                const int row = idx / BLOCK_M;
                const int col = idx % BLOCK_M;
                // transpose A, so that we can read contiguous-K later
                // TODO: some kind of swizzle layout
                A_shared[col * BLOCK_K + row] = A[row * M + col];
            }
            for (int idx = tid; idx < BLOCK_K * BLOCK_N; idx += TB_SIZE) {
                const int row = idx / BLOCK_N;
                const int col = idx % BLOCK_N;
                B_shared[col * BLOCK_K + row] = B[row * N + col];
            }
            __syncthreads();

            for (int mi = 0; mi < WARP_M / 16; ++mi)
                for (int ni = 0; ni < WARP_N / 16; ++ni) {
                    // NOTE: read 64-bit -> bank conflict
                    long A_reg, B_reg;
                    {
                        const int row = warp_id_m * WARP_M + mi * 16 + (lane_id % 16);
                        const int col = (lane_id / 16) * 8;
                        A_reg = reinterpret_cast<long*>(A_shared + row * BLOCK_K + col)[0];
                    }
                    {
                        const int row = warp_id_n * WARP_N + ni * 16 + (lane_id % 16);
                        const int col = (lane_id / 16) * 8;
                        B_reg = reinterpret_cast<long*>(B_shared + row * BLOCK_K + col)[0];
                    }
                    acc2[mi][ni] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(A_reg, B_reg, acc2[mi][ni], 0, 0, 0);
                }
            __syncthreads();

            A += BLOCK_K * M;
            B += BLOCK_K * N;
        }

        const float b_scale = Bs[(i / BLOCK) * cdiv(N, BLOCK)];
        for (int mi = 0; mi < WARP_M / 16; ++mi) {
            const int row = i / BLOCK;
            const int col = mi * 16 + (lane_id / 16) * 4;
            fp32x4 scale4 = reinterpret_cast<const fp32x4*>(As + row * M + col)[0];
            for (int item = 0; item < 4; ++item)
                scale4[item] *= b_scale;

            for (int ni = 0; ni < WARP_N / 16; ++ni) {
                for (int item = 0; item < 4; ++item) {
                    acc1[mi][ni][item] += acc2[mi][ni][item] * scale4[item];
                }
            }
        }
    }

    for (int mi = 0; mi < WARP_M / 16; ++mi)
        for (int ni = 0; ni < WARP_N / 16; ++ni)
            for (int item = 0; item < 4; ++item) {
                const int row = mi * 16 + (lane_id / 16) * 4 + item;
                const int col = ni * 16 + (lane_id % 16);
                C[row * N + col] = (bf16)acc1[mi][ni][item];
            }
}

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

    // __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8
    // https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-mi300-cdna3-instruction-set-architecture.pdf
    // input layout
    //   K_L = 32 / (64 / 16) = 8
    //   logical layout A [m,k] = [16,32]
    //   physical layout [(k/8) * 16 + m, k%8]
    //
    //   |--8--|--8--|--8--|--8--|
    //   | L0  | L16 | L32 | L48 |
    //   | L1  | L17 | L33 | L49 |
    //   | ... |
    //   | L15 | L31 | L47 | L63 |
    //
    // output layout
    //   H = 4
    //   B_I = cdiv(64, 16 * 16 / 4) = 1
    //   M_I = (64 / 1) / 16 = 4
    //   G = 16 / (4 * 4) = 1
    //   logical layout C [m,n] = [16,16]
    //   physical layout [(m/4) * 16 + n, m%4]
    //
    // 4 | L0  | L1  | ... | L15
    // 4 | L16 | L17 |     | L31
    // 4 | L32 | L33 |     | L47
    // 4 | L48 | L49 |     | L63

    constexpr int BLOCK_M = 64;
    constexpr int BLOCK_N = 64;
    constexpr int WARP_M = 32;
    constexpr int WARP_N = 32;

    const int grid_size = cdiv(M, BLOCK_M) * cdiv(N, BLOCK_N);
    const int tb_size = (BLOCK_M / WARP_M) * (BLOCK_N / WARP_N) * WARP_SIZE; // 256

    custom_kernel<BLOCK_M, BLOCK_N, WARP_M, WARP_N><<<grid_size, tb_size, 0, 0>>>(
        (fp8*)a.data_ptr(),
        (fp8*)b.data_ptr(),
        as.data_ptr<float>(),
        bs.data_ptr<float>(),
        (bf16*)c.data_ptr(),
        M,
        N,
        K
    );
    //C10_CUDA_CHECK(cudaGetLastError());
}
"""

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",
        "-DHIP_FAST_MATH",
    ],
)

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/tmpduevsq4v
[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.
