#!POPCORN leaderboard amd-identity
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_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;
}

constexpr const int BLOCK = 128; 
__global__ void kernel(
    const fp8* a,
    const fp8* b,
    const fp32* as,
    const fp32* bs,
    bf16* c,
    int m,
    int n,
    int k
) {    
    // Block that we want to compute in this ThreadBlock
    const int cRow = blockIdx.x;
    const int cCol = blockIdx.y;
    if(cRow >= m || cCol >= n) return;

    int sn = ceil_div(n, BLOCK);

    // allocate buffer for current block in fast shared mem
    // shared mem is shared between all threads in a block
    __shared__ fp8 Amem[BLOCK][BLOCK];
    __shared__ fp8 Bmem[BLOCK][BLOCK];

    // inner row and col that we access in this thread
    const int threadCol = threadIdx.x % BLOCK;
    const int threadRow = threadIdx.x / BLOCK;
    
    // advance pointers to the starting positions
    a += cRow * BLOCK * k;                  // row=cRow, col=0
    b += cCol * BLOCK;                      // row=0, col=cCol
    c += cRow * BLOCK * n + cCol * BLOCK;   // row=cRow, col=cCol
    
    fp32 result = 0;
    for(int bkIdx = 0; bkIdx < k; bkIdx += BLOCK) {
        // Have each thread load one of the elements in A & B
        // Make the threadCol (=threadIdx.x) the consecutive index
        // to allow global memory access coalescing
        *Amem[threadRow * BLOCK + threadCol] = a[threadRow * k + threadCol];
        *Bmem[threadRow * BLOCK + threadCol] = b[threadRow * n + threadCol];

        // block threads in this block until cache is fully populated
        __syncthreads();
        a += BLOCK;
        a += BLOCK * n;

        fp32 tile_result = 0;
        #pragma loop unroll
        for (int dotIdx = 0; dotIdx < BLOCK; ++dotIdx) {
            tile_result += (float)*Amem[threadRow * BLOCK + dotIdx] * (float)*Bmem[dotIdx * BLOCK + threadCol];
        }
        // need to sync again at the end, to avoid faster threads
        // fetching the next block into the cache before slower threads are done
        result += tile_result * as[threadRow + bkIdx/BLOCK * m] * bs[threadCol/BLOCK + bkIdx/BLOCK * sn];
        
        __syncthreads();
    }
    // uncomment of all for memory faults!
    // c[threadRow * n + threadCol] = (__hip_bfloat16)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);

    int block_m = 16;
    int block_n = 16;

    assert(m % block_m == 0);
    assert(n % block_n == 0);

    int grid_m = m / block_m;
    int grid_n = n / block_n;

    kernel<<<dim3(grid_m, grid_n), dim3(block_m, block_n), 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"],
    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 // BLOCK],
            b_scale: torch.Tensor[float32] of shape [n // BLOCK, k // BLOCK],
            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,
    # )
    inp, out = data

    return inp
python eval.py test /tmp/tmp1_fh4cce
[92mSuccessfully preprocessed all matching files.[0m
