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

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

CPP_WRAPPER = """
void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c);
"""

CUDA_SRC = """
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>
#include <hip/hip_cooperative_groups.h>
#include <rocwmma/rocwmma.hpp>
namespace cg = cooperative_groups;

#define K_JUMP 128
#define WMMA_M 32
#define WMMA_N 32
#define WMMA_K 16
#define NUM_WARP_X 4
#define NUM_WARP_Y 4
#define WARP_SIZE 64
#define WMMA_PLACEHOLDER 4
#define OPTION 0

__global__ void __launch_bounds__(1024, 8) custom_kernel(const __hip_fp8_e4m3_fnuz* A_ptr, const __hip_fp8_e4m3_fnuz* B_ptr, const float* As_ptr, const float* Bs_ptr, rocwmma::bfloat16_t* C_ptr, int M, int N, int K) {
    /*
    A: M x K col-maj  --- As (1x128 scaling): M x K/128 col-maj
    B: K x N row-maj  --- Bs (128x128 scaling): K/128 x N/128 row-maj
    C: M x N row-maj
    M, N is divisible by 64 (some are only divisible by 32)
    K is divisible by 128
    */
    
    __shared__ __hip_fp8_e4m3_fnuz tileA[256][136]; // row-major
    __shared__ __hip_fp8_e4m3_fnuz tileB[128][136]; // col-major, = row-major transposed
    int lda = 136;
    int ldb = 136;
    __shared__ float tileAs[256];

    auto frag_acc_master_top = rocwmma::fragment<rocwmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float>();
    rocwmma::fill_fragment(frag_acc_master_top, 0.f);

    auto frag_acc_master_bottom = rocwmma::fragment<rocwmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float>();
    rocwmma::fill_fragment(frag_acc_master_bottom, 0.f);

    // int lda = M;
    // int ldb = N;
    cg::thread_block this_tb = cg::this_thread_block();
    
    for (int kbase = 0; kbase < K; kbase += K_JUMP) {

        // Load SMEM: tileA, tileB, tileAs, tileBs
        // TRY: test tiled partition 64
        // int warp_idx_y = threadIdx.y;
        // int warp_idx_x = threadIdx.x / WARP_SIZE;      

        // ---------------------- FIRST HALF -------------------------


        // Try 2
        #pragma unroll // 1024 / 256 = 4
        for (int kmini_load = 0; kmini_load < 64; kmini_load += 4){
            tileA[threadIdx.x][kmini_load + threadIdx.y] = A_ptr[blockIdx.y * WMMA_M * NUM_WARP_Y * 2 + threadIdx.x + M * (kbase + kmini_load + threadIdx.y)];
        } // end for kmini_load < K_JUMP
        
        // auto group64 = cg::tiled_partition<64>(this_tb); // 1024 threads -> 64 x 16 
        // int t64x = group64.thread_rank();
        // int t64y = group64.meta_group_rank(); 

        // #pragma unroll
        // for (int kmini_load = 0; kmini_load < 64; kmini_load += 16){  // 16 is num group
        //     tileB[t64x][kmini_load + t64y] = B_ptr[(kbase + kmini_load + t64y) * N + blockIdx.x * WMMA_N * NUM_WARP_X + t64x];
        //     tileB[t64x + 64][kmini_load + t64y] = B_ptr[(kbase + kmini_load + t64y) * N + blockIdx.x * WMMA_N * NUM_WARP_X + t64x + 64];
        // } // end for kmini_load < K_JUMP

        int t128x = this_tb.thread_rank() % 128;
        int t128y = this_tb.thread_rank() / 128; // 0 -> 7

        #pragma unroll
        for (int kmini_load = 0; kmini_load < 64; kmini_load += 8){  // 16 is num group
            tileB[t128x][kmini_load + t128y] = B_ptr[(kbase + kmini_load + t128y) * N + blockIdx.x * WMMA_N * NUM_WARP_X + t128x];
        } // end for kmini_load < K_JUMP

        if (threadIdx.y == 0){
            int sn = (N + 128 - 1) / 128;
            float tileBs = Bs_ptr[(kbase / 128) * sn + ((blockIdx.x * NUM_WARP_X * WMMA_N) / 128)];
            tileAs[threadIdx.x] = As_ptr[blockIdx.y * WMMA_M * NUM_WARP_Y * 2 + threadIdx.x + M * (kbase / 128)] * tileBs;
        }

        // float tileBs = Bs_ptr[(kbase / 128) * sn + ((blockIdx.x * NUM_WARP_X * WMMA_N) / 128)];
        // tileAs[threadIdx.x] = As_ptr[blockIdx.y * WMMA_M * NUM_WARP_Y + threadIdx.x + M * (kbase / 128)] * tileBs;


        __syncthreads();


        // MFMA on SMEM
        auto frag_acc_block_top = rocwmma::fragment<rocwmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float>() ;
        rocwmma::fill_fragment(frag_acc_block_top, 0.f);

        auto frag_acc_block_bottom = rocwmma::fragment<rocwmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float>() ;
        rocwmma::fill_fragment(frag_acc_block_bottom, 0.f);
        auto frag_a_top = rocwmma::fragment<rocwmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, __hip_fp8_e4m3_fnuz, rocwmma::row_major>();
        auto frag_a_bottom = rocwmma::fragment<rocwmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, __hip_fp8_e4m3_fnuz, rocwmma::row_major>();
        auto frag_b = rocwmma::fragment<rocwmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, __hip_fp8_e4m3_fnuz, rocwmma::col_major>();   

        int warp_idx_y = threadIdx.y;
        int warp_idx_x = threadIdx.x / WARP_SIZE; 
            
        #pragma unroll 2
        for (int kmini = 0; kmini < 64; kmini += WMMA_K) {           
            rocwmma::load_matrix_sync(frag_a_top, &tileA[0][0] + (warp_idx_y * WMMA_M) * lda + kmini, lda);    
            rocwmma::load_matrix_sync(frag_b, &tileB[0][0] + kmini + ldb * (warp_idx_x * WMMA_N), ldb);
            rocwmma::mma_sync(frag_acc_block_top, frag_a_top, frag_b, frag_acc_block_top); 

            rocwmma::load_matrix_sync(frag_a_bottom, &tileA[0][0] + (warp_idx_y * WMMA_M + 128) * lda + kmini, lda); 
            rocwmma::mma_sync(frag_acc_block_bottom, frag_a_bottom, frag_b, frag_acc_block_bottom);            
        } // end for kmini

        // __syncthreads();

        // ---------------------- SECOND HALF -------------------------

        #pragma unroll
        for (int kmini_load = 64; kmini_load < K_JUMP; kmini_load += 4){
            tileA[threadIdx.x][kmini_load + threadIdx.y] = A_ptr[blockIdx.y * WMMA_M * NUM_WARP_Y * 2 + threadIdx.x + M * (kbase + kmini_load + threadIdx.y)];
        } // end for kmini_load < K_JUMP
        
        // #pragma unroll
        // for (int kmini_load = 64; kmini_load < K_JUMP; kmini_load += 16){  // 16 is num group
        //     tileB[t64x][kmini_load + t64y] = B_ptr[(kbase + kmini_load + t64y) * N + blockIdx.x * WMMA_N * NUM_WARP_X + t64x];
        //     tileB[t64x + 64][kmini_load + t64y] = B_ptr[(kbase + kmini_load + t64y) * N + blockIdx.x * WMMA_N * NUM_WARP_X + t64x + 64];
        // } // end for kmini_load < K_JUMP

        #pragma unroll
        for (int kmini_load = 64; kmini_load < K_JUMP; kmini_load += 8){  // 16 is num group
            tileB[t128x][kmini_load + t128y] = B_ptr[(kbase + kmini_load + t128y) * N + blockIdx.x * WMMA_N * NUM_WARP_X + t128x];
        } // end for kmini_load < K_JUMP

        __syncthreads();

        #pragma unroll 2
        for (int kmini = 64; kmini < K_JUMP; kmini += WMMA_K) {           
            rocwmma::load_matrix_sync(frag_a_top, &tileA[0][0] + (warp_idx_y * WMMA_M) * lda + kmini, lda);    
            rocwmma::load_matrix_sync(frag_b, &tileB[0][0] + kmini + ldb * (warp_idx_x * WMMA_N), ldb);
            rocwmma::mma_sync(frag_acc_block_top, frag_a_top, frag_b, frag_acc_block_top); 

            rocwmma::load_matrix_sync(frag_a_bottom, &tileA[0][0] + (warp_idx_y * WMMA_M + 128) * lda + kmini, lda); 
            rocwmma::mma_sync(frag_acc_block_bottom, frag_a_bottom, frag_b, frag_acc_block_bottom);            
        } // end for kmini


        // Apply scaler
        // int col_base_C = (blockIdx.x * NUM_WARP_X + warp_idx_x) * WMMA_N;
        // int row_base_C = (blockIdx.y * NUM_WARP_Y * 2 + warp_idx_y) * WMMA_M;

        int t64x = this_tb.thread_rank() % 64;
        int place = warp_idx_y * WMMA_M + 4 * (t64x / 32);

        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                int index = 4 * i + j;
                frag_acc_master_top.x[index] += tileAs[place + 8 * i + j] * frag_acc_block_top.x[index];
            }
        }

        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                int index = 4 * i + j;
                frag_acc_master_bottom.x[index] += tileAs[place + 128 + 8 * i + j] * frag_acc_block_bottom.x[index];
            }
        }

        // __syncthreads(); // We need this synthreads()
    } // end for kbase

    int warp_idx_x = threadIdx.x / WARP_SIZE;    
    int col_base_C = (blockIdx.x * NUM_WARP_X + warp_idx_x) * WMMA_N;
    if (col_base_C >= N) return;    

    int warp_idx_y = threadIdx.y;
    int row_base_C = (blockIdx.y * NUM_WARP_Y * 2 + warp_idx_y) * WMMA_M;
    if (row_base_C >= M) return;

    // Convert fp32 to bf16 and store
    auto frag_result_top = rocwmma::fragment<rocwmma::accumulator, WMMA_M, WMMA_N, WMMA_PLACEHOLDER, rocwmma::bfloat16_t>();

    #pragma unroll
    for (int i = 0; i < frag_result_top.num_elements; ++i) {
        frag_result_top.x[i] = (rocwmma::bfloat16_t)frag_acc_master_top.x[i];
    }
    rocwmma::store_matrix_sync(C_ptr + row_base_C * N + col_base_C, frag_result_top, N, rocwmma::mem_row_major);

    if (row_base_C + 128 >= M) return;

    auto frag_result_bottom = rocwmma::fragment<rocwmma::accumulator, WMMA_M, WMMA_N, WMMA_PLACEHOLDER, rocwmma::bfloat16_t>();

    #pragma unroll
    for (int i = 0; i < frag_result_top.num_elements; ++i) {
        frag_result_bottom.x[i] = (rocwmma::bfloat16_t)frag_acc_master_bottom.x[i];
    }   
    
    rocwmma::store_matrix_sync(C_ptr + (row_base_C + 128) * N + col_base_C, frag_result_bottom, N, rocwmma::mem_row_major);
}

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

    int num_warp_x = NUM_WARP_X;
    int num_warp_y = NUM_WARP_Y;
    dim3 block_size((WARP_SIZE * num_warp_x), num_warp_y);
    dim3 grid_size((N + num_warp_x * WMMA_N - 1) / (num_warp_x * WMMA_N), 
                   (M + num_warp_y * WMMA_M * 2 - 1) / (num_warp_y * WMMA_M * 2));

    custom_kernel<<<grid_size, block_size, 0, 0>>> ((__hip_fp8_e4m3_fnuz*)a.data_ptr(), (__hip_fp8_e4m3_fnuz*)b.data_ptr(), 
    as.data_ptr<float>(), bs.data_ptr<float>(), (rocwmma::bfloat16_t*)c.data_ptr(), M, N, K);
    //C10_CUDA_CHECK(cudaGetLastError());
}
"""

import os
os.environ["CXX"] = "clang++"
import torch
block_shape = 128


# big one, used when N is >= 3072
module_big = load_inline(
    name='fp8_mm',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['fp8_mm'],
    verbose=True,
    extra_cuda_cflags=["-std=c++20", "-ffast-math"],    
)

def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    module_big.fp8_mm(a, b, a_scale, b_scale, c) 
    return c

python submission.py
[92mSuccessfully preprocessed all matching files.[0m

/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]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip [ok]
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:xnack- -fno-gpu-rdc -std=c++20 -ffast-math -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:23:44: warning: failed to meet occupancy target given by 'amdgpu-waves-per-eu' in '_Z13custom_kernelPK19__hip_fp8_e4m3_fnuzS1_PKfS3_P12hip_bfloat16iii': desired occupancy was 8, final occupancy is 4 [-Wpass-failed]
   23 | __global__ void __launch_bounds__(1024, 8) custom_kernel(const __hip_fp8_e4m3_fnuz* A_ptr, const __hip_fp8_e4m3_fnuz* B_ptr, const float* As_ptr, const float* Bs_ptr, rocwmma::bfloat16_t* C_ptr, int M, int N, int K) {
      |                                            ^
1 warning generated when compiling for gfx942.
[3/3] clang++ hip.cuda.o main.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 benchmark /tmp/tmpfqq7pq51
[92mSuccessfully preprocessed all matching files.[0m

/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]
/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]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 1
[1/1] 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
