#!POPCORN leaderboard amd-fp8-mm
from task import input_t, output_t
import torch
from torch.utils.cpp_extension import load_inline
import time
import os
import sys

if "PYTORCH_ROCM_ARCH" not in os.environ:
    os.environ["PYTORCH_ROCM_ARCH"] = "gfx942:xnack-"

kernel_cpp = r"""
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#include <hip/hip_fp8.h>
#include <pybind11/pybind11.h>
#include <iostream>

#define HIP_CHECK(condition)                                                                  @
    do                                                                                        @
    {                                                                                         @
        hipError_t error = condition;                                                         @
        if(error != hipSuccess)                                                               @
        {                                                                                     @
            std::cout << "Error " << hipGetErrorName(error) << '(' << error << ')' << ": "    @
                      << hipGetErrorString(error) << " in " << __func__ << " at " << __FILE__ @
                      << ':' << __LINE__ << '@n';                                             @
            exit(error);                                                                      @
        }                                                                                     @
    }                                                                                         @
    while(false)

using fp8 = __hip_fp8_e4m3_fnuz;
using bf16 = __hip_bfloat16;
using fp32 = float;

template <typename T, typename U>
__host__ __device__
constexpr auto ceil_div(T a, U b) {
    return (a + b - 1) / b;
}

template<unsigned int BlockSize>
__global__ void kernel(
    const fp8* a,
    const fp8* b,
    const fp32* as,
    const fp32* bs,
    bf16* c,
    int m,
    int n,
    int k
) {
    const unsigned int tx = threadIdx.x;
    const unsigned int ty = threadIdx.y;
    const unsigned int bx = blockIdx.x;
    const unsigned int by = blockIdx.y;

    const unsigned int row = bx * blockDim.x + tx;
    const unsigned int col = by * blockDim.y + ty;

    const unsigned int COARSE_BLOCK_SIZE = 128 / BlockSize;
    __shared__ fp8 a_shared[BlockSize][BlockSize * COARSE_BLOCK_SIZE];
    __shared__ fp8 b_shared[BlockSize * COARSE_BLOCK_SIZE][BlockSize];
    __shared__ fp32 a_scale_shared[BlockSize];
    __shared__ fp32 b_scale_shared;
    int sn = ceil_div(n, 128);

    fp32 result = 0;
    for (int i = 0; i < k; i += BlockSize * COARSE_BLOCK_SIZE) {
        for (int cbs=0; cbs<COARSE_BLOCK_SIZE; ++cbs) {
            a_shared[tx][ty + cbs * BlockSize] = a[(i + ty + cbs * BlockSize) * m + bx * blockDim.x + tx];
            b_shared[ty + cbs * BlockSize][tx] = b[(i + ty + cbs * BlockSize) * n + by * blockDim.y + tx];
        }
        if (i%128==0 && ty == 0) {
            a_scale_shared[tx] = as[i / 128 * m + row];
            if (tx == 0) {
                b_scale_shared = bs[i / 128 * sn + col / 128];
            }
        }
        __syncthreads();
        fp32 block_result = 0;

        for (int ii = 0; ii < BlockSize * COARSE_BLOCK_SIZE; ++ii) {
            block_result +=
                fp32(a_shared[tx][ii]) *
                fp32(b_shared[ii][ty]);
        }
        result += block_result * a_scale_shared[tx] * b_scale_shared;
        __syncthreads();
    }
    c[row * n + col] = bf16(result);
}

void run(
    uintptr_t a_ptr,
    uintptr_t b_ptr,
    uintptr_t as_ptr,
    uintptr_t bs_ptr,
    uintptr_t c_ptr,
    int m,
    int n,
    int k
) {
    const auto* a = reinterpret_cast<const fp8*>(a_ptr);
    const auto* b = reinterpret_cast<const fp8*>(b_ptr);
    const auto* as = reinterpret_cast<const fp32*>(as_ptr);
    const auto* bs = reinterpret_cast<const fp32*>(bs_ptr);
    auto* c = reinterpret_cast<bf16*>(c_ptr);

    constexpr unsigned int block_size = 32;
    assert(m % block_size == 0);
    assert(n % block_size == 0);

    int grid_m = m / block_size;
    int grid_n = n / block_size;

    kernel<block_size><<<dim3(grid_m, grid_n), dim3(block_size, block_size), 0>>>(
        a,
        b,
        as,
        bs,
        c,
        m,
        n,
        k
    );
    HIP_CHECK(hipGetLastError());
}

PYBIND11_MODULE(fp8, m) {
  m.def("fp8", &run, "HIP kernel");
}
"""

hip_module = load_inline(
    name="fp8",
    cpp_sources="",
    cuda_sources=kernel_cpp.replace('@', chr(92)),
    with_cuda=True,
    verbose=False,
    extra_cuda_cflags=["-std=c++20"],
)

first = True

def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of block-scale fp8 gemm
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k],
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k],
            a_scale: torch.Tensor[float32] of shape [m, k // 128],
            b_scale: torch.Tensor[float32] of shape [n // 128, k // 128],
            c: torch.Tensor[bfloat16] of shape [m, n]
    Returns:
        Tensor containing output in bf16
    """

    global first
    if first:
        print("executing on", torch.cuda.get_device_name(), file=sys.stderr)
        first = False

    a, b, a_scale, b_scale, c = data

    m, n = c.shape
    k = a.shape[1]

    hip_module.fp8(
        a.data_ptr(),
        b.data_ptr(),
        a_scale.data_ptr(),
        b_scale.data_ptr(),
        c.data_ptr(),
        m,
        n,
        k,
    )

    return c

# if __name__ == "__main__":
#     a = torch.randn(1024, 1024, dtype=torch.float8_e4m3fnuz)
#     b = torch.randn(1024, 1024, dtype=torch.float8_e4m3fnuz)
#     a_scale = torch.randn(1024, 1024 // 128, dtype=torch.float32)
#     b_scale = torch.randn(1024 // 128, 1024 // 128, dtype=torch.float32)
#     c = torch.randn(1024, 1024, dtype=torch.bfloat16)
#     custom_kernel((a, b, a_scale, b_scale, c))
python eval.py benchmark /tmp/tmpye8ia4sm
[92mSuccessfully preprocessed all matching files.[0m
executing on AMD Instinct MI300X
