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

import os

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

import torch
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;

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

template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
__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
) {
    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;

    const int offm = bid_m * BLOCK_M;
    const int offn = bid_n * BLOCK_N;
    A += offm;
    B += offn;
    As += offm;
    Bs += offn / BLOCK;
    C += offm * N + offn;

    constexpr int THREAD_M = 4;
    constexpr int THREAD_N = 4;
    constexpr int num_thread_tiles_n = BLOCK_N / THREAD_N;  // 16

    // TODO: warp tiling?
    const int thread_offm = (tid / num_thread_tiles_n) * THREAD_M;
    const int thread_offn = (tid % num_thread_tiles_n) * THREAD_N;

    float acc1[THREAD_M][THREAD_N] = {};
    float acc2[THREAD_M][THREAD_N];
    __shared__ float A_shared[BLOCK_K][BLOCK_M];
    __shared__ float B_shared[BLOCK_K][BLOCK_N];

    for (int i = 0; i < K; i += BLOCK) {
        // reset acc2
        for (int m = 0; m < THREAD_M; ++m)
            for (int n = 0; n < THREAD_N; ++n)
                acc2[m][n] = 0.0f;

        for (int ii = 0; ii < BLOCK; ii += BLOCK_K) {
            // load A(BLOCK_K, BLOCK_M) and B(BLOCK_K, BLOCK_N)
            const int row = tid / BLOCK_M;
            const int col = tid % BLOCK_M;
            A_shared[row][col] = (float)A[row * M + col];
            B_shared[row][col] = (float)B[row * N + col];
            __syncthreads();

            for (int k = 0; k < BLOCK_K; ++k) {
                float B_reg[THREAD_N];
                for (int n = 0; n < THREAD_N; ++n)
                    B_reg[n] = B_shared[k][thread_offn + n];

                for (int m = 0; m < THREAD_M; ++m) {
                    float A_reg = A_shared[k][thread_offm + m];
                    for (int n = 0; n < THREAD_N; ++n)
                        acc2[m][n] += A_reg * B_reg[n];
                }
            }
            __syncthreads();

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

        const float b_scale = Bs[(i / BLOCK) * cdiv(N, BLOCK)];
        for (int m = 0; m < THREAD_M; ++m) {
            const float scale = b_scale * As[(i / BLOCK) * M + (thread_offm + m)];
            for (int n = 0; n < THREAD_N; ++n)
                acc1[m][n] += acc2[m][n] * scale;
        }
    }

    for (int m = 0; m < THREAD_M; ++m)
        for (int n = 0; n < THREAD_N; ++n) {
            C[(thread_offm + m) * N + (thread_offn + n)] = (bf16)acc1[m][n];
        }
}

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

    // thread tiling
    // each thread will compute outer product (2,1) x (1,2) = (2,2)
    // threa block will compute outer product (64,4) x (4,64) = (64,64)
    constexpr int BLOCK_M = 64;
    constexpr int BLOCK_N = 64;
    constexpr int BLOCK_K = 4;
    const int grid_size = cdiv(M, BLOCK_M) * cdiv(N, BLOCK_N);
    const int tb_size = (BLOCK_M / 4) * (BLOCK_N / 4);  // 256

    custom_kernel<BLOCK_M, BLOCK_N, BLOCK_K><<<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)
    torch.cuda.synchronize()
    return c

python submission.py
[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 [ok]
/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/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 -std=c++20 -DHIP_FAST_MATH -fno-gpu-rdc -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
[3/3] clang++ main.o hip.cuda.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 test /tmp/tmppc34yf27
[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.
