# This script provides a template for using load_inline to run a HIP kernel for
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 = """
#ifdef __NVCC__ // compiling with nvcc

#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <mma.h>

// using __half to be able to debug on A100,
// which does not support __nv_float8_e4m3
using float8 = __half;
using bfloat16 = __nv_bfloat16;

using nvcuda::wmma::accumulator;
using nvcuda::wmma::col_major;
using nvcuda::wmma::matrix_a;
using nvcuda::wmma::matrix_b;
using nvcuda::wmma::row_major;
using nvcuda::wmma::fragment;
using nvcuda::wmma::fill_fragment;
using nvcuda::wmma::load_matrix_sync;
using nvcuda::wmma::mma_sync;
using nvcuda::wmma::store_matrix_sync;
using nvcuda::wmma::mem_row_major;

constexpr uint warp_size = 32;
constexpr int wmma_size_m = 16;
constexpr int wmma_size_n = 16;
constexpr int wmma_size_k = 16;

#else // compiling with hip

#include <hip/hip_ext.h>
#include <hip/hip_runtime.h>

#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

#include <rocwmma/rocwmma.hpp>
#include <rocwmma/rocwmma_coop.hpp>
#include <rocwmma/rocwmma_transforms.hpp>

using float8 = __hip_fp8_e4m3_fnuz;
using bfloat16 = __hip_bfloat16;

using rocwmma::accumulator;
using rocwmma::col_major;
using rocwmma::matrix_a;
using rocwmma::matrix_b;
using rocwmma::row_major;
using rocwmma::fragment;
using rocwmma::fill_fragment;
using rocwmma::load_matrix_sync;
using rocwmma::mma_sync;
using rocwmma::store_matrix_sync;
using rocwmma::mem_row_major;

constexpr uint warp_size = 64; // CDNA has warp size of 64
constexpr int wmma_size_m = 16;
constexpr int wmma_size_n = 16;
constexpr int wmma_size_k = 32;

#endif

constexpr int num_warps_m = 4;
constexpr int num_warps_n = 2;
constexpr int thread_block_size_x = num_warps_m * warp_size;
constexpr int thread_block_size_y = num_warps_n;

// CTA tile sizes
constexpr int tile_size_m = num_warps_m * wmma_size_m;
constexpr int tile_size_n = num_warps_n * wmma_size_n;
constexpr int quant_tile_size = 128;

template<typename T>
__host__ __device__ 
inline T cdiv(T num, T den) {
    return (num + den - 1) / den;
}

template<typename DType>
__device__ __forceinline__ DType* index_colmajor(DType* data, int row, int col, int nrows) {
    return data + row + col * nrows;
}

template<typename DType>
__device__ __forceinline__ DType* index_rowmajor(DType* data, int row, int col, int ncols) {
    return data + row * ncols + col;
}

__global__ void custom_kernel(
    const float8* a,
    const float8* bt, 
    const float* a_scale, 
    const float* bt_scale, 
    bfloat16* c,
    const int m, 
    const int n, 
    const int k
) {

    // CTAs are allocated 1d, so they can be rearranged.
    const int num_cta_x = cdiv(m, tile_size_m);
    const int num_cta_y = cdiv(n, tile_size_n);

    const int block_idx_x = blockIdx.x % num_cta_x;
    const int block_idx_y = blockIdx.x / num_cta_x;

    // threads are allocated 2d, since they're naturally grouped by warp already.
    const int num_threads_per_block = blockDim.x * blockDim.y; 
    const int tid = threadIdx.x * blockDim.y + threadIdx.y;

    const int warp_row = (threadIdx.x / warp_size) * wmma_size_m;
    const int warp_col = threadIdx.y * wmma_size_n;

    // CTA tiling
    const int c_row = block_idx_x * tile_size_m;
    const int c_col = block_idx_y * tile_size_n;

    // move a, b, c pointers to starting positions for tile
    a = index_colmajor(a, c_row, 0, m);
    bt = index_colmajor(bt, c_col, 0, n);
    c = index_rowmajor(c, c_row, c_col, n);

    // move as and bs pointers to starting positions.
    const int scale_n = (n + quant_tile_size - 1) / quant_tile_size;
    a_scale = index_colmajor(a_scale, c_row, 0, m);
    bt_scale = index_colmajor(bt_scale, c_col / quant_tile_size, 0, scale_n);

    // smem takes entire quantized tiles
    __shared__ float8 a_smem[tile_size_m * quant_tile_size];
    __shared__ float8 b_smem[tile_size_n * quant_tile_size];

    // shared memory for scaling tiles.
    __shared__ float quant_tile_smem[tile_size_m * tile_size_n];

    // shared memory that we write accumulate output onto
    // this is explicitly zerod first.
    __shared__ float accum_smem[tile_size_m * tile_size_n];
    for (int i = 0; i < tile_size_m * tile_size_n; i += num_threads_per_block) {
        int idx = tid + i;
        if (idx < tile_size_m * tile_size_n) {
            accum_smem[idx] = 0;
        }
    }
    
    for (int quant_tile_idx = 0; quant_tile_idx < k; quant_tile_idx += quant_tile_size) {

        // we have enough threads to cover a row, so step across quant tile columns.
        // the row indices should correspond to contiguous threads.
        // NOTE: this assumes that we have enough threads to cover the smem rows.
        // TODO: need zero-padding for weird sized problems
        const int a_load_k_size = num_threads_per_block / tile_size_m;
        const int a_inner_row = tid % tile_size_m;
        const int a_inner_col = tid / tile_size_m;
        for (int load_offset = 0; load_offset < quant_tile_size; load_offset += a_load_k_size) {
            const int row_idx = a_inner_row;
            const int col_idx = a_inner_col + load_offset;
            if (c_row + row_idx < m) {
                a_smem[row_idx + col_idx * tile_size_m] = a[row_idx + (col_idx + quant_tile_idx) * m];
            } else {
                a_smem[row_idx + col_idx * tile_size_m] = 0.0;
            }
        }
        
        const int b_load_k_size = num_threads_per_block / tile_size_n;
        const int b_inner_row = tid % tile_size_n;
        const int b_inner_col = tid / tile_size_n;
        for (int load_offset = 0; load_offset < quant_tile_size; load_offset += b_load_k_size) {
            const int row_idx = b_inner_row;
            const int col_idx = b_inner_col + load_offset;
            // b_smem is column major tile of bt, transposed.
            if (c_col + row_idx < n) {
                b_smem[col_idx + row_idx * quant_tile_size] = bt[row_idx + (col_idx + quant_tile_idx) * n];
            } else {
                b_smem[col_idx + row_idx * quant_tile_size] = 0.0;
            }
        }

        __syncthreads();

        auto frag_a = fragment<matrix_a, wmma_size_m, wmma_size_n, wmma_size_k, float8, col_major>();
        auto frag_b = fragment<matrix_b, wmma_size_m, wmma_size_n, wmma_size_k, float8, col_major>();
        auto frag_accum = fragment<accumulator, wmma_size_m, wmma_size_n, wmma_size_k, float>();
        fill_fragment(frag_accum, 0.0f);

        for (int wmma_tile_idx = 0; wmma_tile_idx < quant_tile_size; wmma_tile_idx += wmma_size_k) {
            float8* a_smem_wmma_tile = index_colmajor(a_smem, warp_row, wmma_tile_idx, tile_size_m);
            float8* b_smem_wmma_tile = index_colmajor(b_smem, wmma_tile_idx, warp_col, quant_tile_size);

            load_matrix_sync(frag_a, a_smem_wmma_tile, tile_size_m);
            load_matrix_sync(frag_b, b_smem_wmma_tile, quant_tile_size);
            mma_sync(frag_accum, frag_a, frag_b, frag_accum);
        }

        // copy the accumulated quant tile into smem array, so we can scale it.
        // I think writing into shared memory is the recommended approach, since
        // it's not specified how the data is mapped to each thread.
        float* quant_tile_smem_tile = index_rowmajor(quant_tile_smem, warp_row, warp_col, tile_size_n);
        store_matrix_sync(quant_tile_smem_tile, frag_accum, tile_size_n, mem_row_major);

        __syncthreads();
        
        const int accum_scale_load_k_size = num_threads_per_block / tile_size_m;
        const int accum_scale_inner_row = tid % tile_size_m;
        const int accum_scale_inner_col = tid / tile_size_m;

        // apply scales to accum.
        // all threads should have the same b scale, since they're all in the same 128x128 block.
        const float bs = *index_colmajor(bt_scale, 0, quant_tile_idx / quant_tile_size, scale_n);
        const float as = *index_colmajor(a_scale, accum_scale_inner_row, quant_tile_idx / quant_tile_size, m);

        __syncthreads();

        // add scaled tile onto the accumulator.
        for (int j = 0; j < tile_size_n; j += accum_scale_load_k_size) {
            const int row_idx = accum_scale_inner_row;
            const int col_idx = accum_scale_inner_col + j;
            accum_smem[row_idx * tile_size_n + col_idx] += quant_tile_smem[row_idx * tile_size_n + col_idx] * as * bs;
        }

        __syncthreads();
    }

    // this assumes num_threads > tile_size_m AND it is a multiple...
    const int c_write_n_size = num_threads_per_block / tile_size_m;
    const int c_inner_row = tid % tile_size_m;
    const int c_inner_col = tid / tile_size_m;
    for (int j = 0; j < tile_size_n; j += c_write_n_size) {
        const int row_idx = c_inner_row;
        const int col_idx = c_inner_col + j;
        if (c_row + row_idx < m && c_col + col_idx < n) {
            c[row_idx * n + col_idx] = 
                static_cast<bfloat16>(*index_rowmajor(accum_smem, row_idx, col_idx, tile_size_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);

    const int grid_dim_x = cdiv(m, tile_size_m);
    const int grid_dim_y = cdiv(n, tile_size_n);

    custom_kernel<<<
        dim3(grid_dim_x * grid_dim_y),
        dim3(thread_block_size_x, thread_block_size_y),
        0, 
        0
    >>>
    (
        (float8*)a.data_ptr(), 
        (float8*)b.data_ptr(), 
        as.data_ptr<float>(), 
        bs.data_ptr<float>(), 
        (bfloat16*)c.data_ptr(),
        m, 
        n, 
        k
    );
    //C10_CUDA_CHECK(cudaGetLastError());
}
"""

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

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


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