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

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 = """
template <int FP8_PER_LOAD, int FP32_PER_LOAD, int SCALE_SIZE, int MAX_K_CHUNKS>
__global__ void __launch_bounds__(128, 8) kernel(
    const fp8* __restrict a,
    const fp8* __restrict b,
    const fp32* __restrict a_scale,
    const fp32* __restrict b_scale,
    bf16* __restrict c,
    const int m,
    const int n,
    const int k
) {
int threadIdx_x = threadIdx.x;

int row0 = threadIdx_x / FP8_PER_LOAD;
int col0 = threadIdx_x % FP8_PER_LOAD;
int row1 = (threadIdx_x + blockDim.x) / FP8_PER_LOAD;
int col1 = (threadIdx_x + blockDim.x) % FP8_PER_LOAD;

int row_start = blockIdx.x * FP8_PER_LOAD;
int col_start = blockIdx.y * FP8_PER_LOAD;

fp32 acc0 = 0.0f;
fp32 acc1 = 0.0f;
    
    const int k_chunks = k / SCALE_SIZE;
    __shared__ fp32 a_scale_share[FP8_PER_LOAD * MAX_K_CHUNKS];

    // We know that the highest k_chunk is 7168 / 128 = 56 and we have 128 threads
    // We need 16 FP32 values per k_chunk, and each thread can load 4 FP32 elements per instruction
    // We therefore at most need to emit 56 * 16 / 4 = 224 instructions
    // we therefore let each thread load 2 instructions given that it would be more than 224 instructions
    if (threadIdx_x < k_chunks * 2) {
        const int jump = threadIdx_x / 2; // Every 2 threads we need to wrap around
        const int offset_scale_a = threadIdx_x % 2 * (2 * FP32_PER_LOAD); // The first 2 x 4 elements have already been loaded if it is uneven 
        const int memory_offset_scale_a = row_start + jump * m + offset_scale_a;

        const float4* ptr = reinterpret_cast<const float4*>(a_scale + memory_offset_scale_a);
        float4 f1 = ptr[0];
        float4 f2 = ptr[1];

        int store = threadIdx_x * FP32_PER_LOAD * 2;
        a_scale_share[store + 0] = f1.x;
        a_scale_share[store + 1] = f1.y;
        a_scale_share[store + 2] = f1.z;
        a_scale_share[store + 3] = f1.w;
        a_scale_share[store + 4] = f2.x;
        a_scale_share[store + 5] = f2.y;
        a_scale_share[store + 6] = f2.z;
        a_scale_share[store + 7] = f2.w;
    }

    __shared__ fp32 b_scale_share[MAX_K_CHUNKS];
    if (threadIdx_x < k_chunks) {
        b_scale_share[threadIdx_x] = b_scale[col_start / SCALE_SIZE + threadIdx_x * ceil_div(n, SCALE_SIZE)]; // Each thread loads a single value for a k_chunk
    }
    __syncthreads();

    // Shared-memory tiles
    __shared__ fp32 a_share[FP8_PER_LOAD * SCALE_SIZE]; // Maybe it is possible to keep either the a or b chunk in registers instead of having them in shared memory, given that it won't be used by other blocks.
    __shared__ fp32 b_share[FP8_PER_LOAD * SCALE_SIZE];

    for (int kb = 0; kb < k_chunks; ++kb) {
        int kidx = kb * SCALE_SIZE + threadIdx_x;
        uint4 a_v = reinterpret_cast<const uint4*>(a + kidx * m + row_start)[0];

        for(int w = 0; w < 4; ++w) {
            uint32_t packed = ((uint32_t*)&a_v)[w];

            for(int l = 0; l < 4; ++l) {
                uint8_t raw = (packed >> (l*8)) & 0xFF;

                __half h = __hip_cvt_fp8_to_halfraw(raw, __HIP_E4M3_FNUZ);
                float  f = __half2float(h);                                  

                int laneIndex = w*4 + l;
                int idx       = laneIndex * SCALE_SIZE + threadIdx_x;
                a_share[idx]  = f * a_scale_share[kb * FP8_PER_LOAD + laneIndex];
            }
        }

        uint4 b_v = reinterpret_cast<const uint4*>(b + kidx * n + col_start)[0];
        for(int w = 0; w < 4; ++w) {
            uint32_t packed = ((uint32_t*)&b_v)[w];
            for(int l = 0; l < 4; ++l) {
                uint8_t raw = (packed >> (l*8)) & 0xFF;
                __half h = __hip_cvt_fp8_to_halfraw(raw, __HIP_E4M3_FNUZ);
                float  f = __half2float(h);
                int laneIndex = w*4 + l;
                int idx       = laneIndex * SCALE_SIZE + threadIdx_x;
                b_share[idx]  = f * b_scale_share[kb];
            }
        }
        __syncthreads();

        for (int s = 0; s < SCALE_SIZE; ++s) {
            acc0 += a_share[row0 * SCALE_SIZE + s]
                  * b_share[col0 * SCALE_SIZE + s];
            acc1 += a_share[row1 * SCALE_SIZE + s]
                  * b_share[col1 * SCALE_SIZE + s];
        }
        __syncthreads();
    }

    // Write-back two BF16 outputs
    int r0 = row_start + row0;
    int c0_ = col_start + col0;
    if (r0 < m && c0_ < n) {
        c[r0 * n + c0_] = static_cast<bf16>(acc0);
    }
    int r1 = row_start + row1;
    int c1_ = col_start + col1;
    if (r1 < m && c1_ < n) {
        c[r1 * n + c1_] = static_cast<bf16>(acc1);
    }
}

#define FP8_PER_LOAD 16
#define FP32_PER_LOAD 4
#define SCALE_SIZE 128
#define MAX_K_CHUNKS 56

void run(
    const uintptr_t a_ptr,
    const uintptr_t b_ptr,
    const uintptr_t a_scale_ptr,
    const uintptr_t b_scale_ptr,
    const uintptr_t c_ptr,
    const int m,
    const int n,
    const int k
) {
    const auto* a = reinterpret_cast<const fp8*>(a_ptr);
    const auto* b = reinterpret_cast<const fp8*>(b_ptr);
    const auto* a_scale = reinterpret_cast<const fp32*>(a_scale_ptr);
    const auto* b_scale = reinterpret_cast<const fp32*>(b_scale_ptr);
    auto* c = reinterpret_cast<bf16*>(c_ptr);

    // Advise and prefetch scale arrays to device L2 cache
    const int k_chunks = k / SCALE_SIZE;
    size_t a_SCALE_SIZE = size_t(m) * size_t(k_chunks) * sizeof(fp32);
    size_t b_SCALE_SIZE = size_t((n + SCALE_SIZE -1)/SCALE_SIZE) * size_t(k_chunks) * sizeof(fp32);

    // int device;
    // HIP_CHECK(hipGetDevice(&device));
    // HIP_CHECK(hipMemAdvise(a_scale, a_SCALE_SIZE, hipMemAdviseSetReadMostly, device));
    // HIP_CHECK(hipMemAdvise(b_scale, b_SCALE_SIZE, hipMemAdviseSetReadMostly, device));
    // HIP_CHECK(hipMemPrefetchAsync(a_scale, a_SCALE_SIZE, device, 0));
    // HIP_CHECK(hipMemPrefetchAsync(b_scale, b_SCALE_SIZE, device, 0));
    
    const int grid_dim_m = ceil_div(m, FP8_PER_LOAD);
    const int grid_dim_n = ceil_div(n, FP8_PER_LOAD);
    
    const int block_dim_x = SCALE_SIZE;

    dim3 gridDim(grid_dim_m, grid_dim_n);
    dim3 blockDim(block_dim_x);

    kernel<FP8_PER_LOAD, FP32_PER_LOAD , SCALE_SIZE, MAX_K_CHUNKS>
        <<<gridDim, blockDim, 0>>>(
            a,
            b,
            a_scale,
            b_scale,
            c,
            m,
            n,
            k
        );
    HIP_CHECK(hipGetLastError());
}
"""

prefix = """
#if !defined(__HIP_PLATFORM_AMD__)
  #define __HIP_PLATFORM_AMD__
#endif
#include <hip/hip_runtime.h>
#include <hip/hip_fp8.h>
#include <hip/hip_fp16.h>
#include <hip/hip_bfloat16.h>
#include <hip/hip_vector_types.h>
#include <pybind11/pybind11.h>
#include <iostream>

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

#undef __HIP_PLATFORM_AMD__

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


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

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

full_cpp = prefix + kernel_cpp + postfix

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

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

    # # Print a 4x4 slice of c for inspection
    # print("c[0:4, 0:4]:")
    # arr = c[:4, :4].to(torch.float32).cpu().numpy()
    # for row in arr:
    #     print(" ".join(f"{v:9.3f}" for v in row))

    return c

python eval.py benchmark /tmp/tmprc9s9yqt
[92mSuccessfully preprocessed all matching files.[0m
executing on AMD Instinct MI300X
