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

__global__ void __launch_bounds__(1024, 8) balance_new_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 block covers 32x32 elem
    dim3 block_size(64, 1) -> 1 thread calculates 16 elem

    Perform B @ A.T = C.T, so we can store scale As in a VGPR of each thread in a warp
    B: NxK col-maj -- Bs (128x128 scaling): N/128 x K/128 col-maj
    A.T: KxM row-maj -- As (1x128 scaling): K/ 128 x M row-maj
    C.T: NxM col-maj
    
    M, N, K are divisible by 128

    Kernel to use when M, N is balanced
    Use wmma 32x32x16

    tileB[128][128]
    tileA[128][112]

    SMEM usage per block: 32KB

    */  

    int sn = (N + 128 - 1) / 128;

    __shared__ __hip_fp8_e4m3_fnuz tileB[128][136]; // col-maj = row-maj transposed 
    __shared__ __hip_fp8_e4m3_fnuz tileA[128][120]; // row-maj
    int ldb = 136;
    int lda = 120;

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

    cg::thread_block this_tb = cg::this_thread_block();
    auto group64 = cg::tiled_partition<64>(this_tb); // 1024 threads -> 64 x 16
    int t64x = group64.thread_rank();
    int t64y = group64.meta_group_rank(); 

    for (int kbase = 0; kbase < K; kbase += K_JUMP) {
        int warp_idx_y = threadIdx.y;    
        int warp_idx_x = threadIdx.x / WARP_SIZE;  

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

        // TODO: replace += 16 with += group64.meta_group_size()

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

        tileB[t64x][112 + t64y] = B_ptr[(kbase + 112 + t64y) * N + blockIdx.y * WMMA_N * NUM_WARP_Y + t64x];
        tileB[t64x + 64][112 + t64y] = B_ptr[(kbase + 112 + t64y) * N + blockIdx.y * WMMA_N * NUM_WARP_Y + t64x + 64];

        float tileAs = As_ptr[(kbase / 128) * M + col_base_C + t64x % WMMA_N];
        float tileBs = Bs_ptr[(kbase / 128) * sn + (blockIdx.y * NUM_WARP_Y * WMMA_M) / 128];
        float scale = tileAs * tileBs;

        // float scale = tileAs * tileBs; // this gives large error, due to non-associative of IEEE 754 ?
        __syncthreads();

        // MFMA

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

        auto frag_a = 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>();

  

        #pragma unroll 1
        for (int kmini = 0; kmini < K_JUMP - 16; kmini += WMMA_K) {
            rocwmma::load_matrix_sync(frag_a, &tileB[0][0] + (warp_idx_y * WMMA_M) * ldb + kmini, ldb);            
            rocwmma::load_matrix_sync(frag_b, &tileA[0][0] + kmini + lda * (warp_idx_x * WMMA_N), lda);
            rocwmma::mma_sync(frag_acc_block, frag_a, frag_b, frag_acc_block);                                
        }

        // ---------------------------------- LEFT OVER -------------------------------------------
        // left over 32 rows
  
        auto frag_b_left = rocwmma::fragment<rocwmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, __hip_fp8_e4m3_fnuz, rocwmma::row_major>(); 
        rocwmma::load_matrix_sync(frag_b_left, A_ptr + (kbase + 112) * M + col_base_C, M);

        // auto transposed_frag = rocwmma::applyTranspose(frag_b_left);     
        rocwmma::load_matrix_sync(frag_a, &tileB[0][0] + (warp_idx_y * WMMA_M) * ldb + 112, ldb);            

        rocwmma::mma_sync(frag_acc_block, frag_a, frag_b_left, frag_acc_block);                                




        __syncthreads();
        // Mult scalers

        // TODO: use frag_acc_master.num_elements or 4 to reduce register pressure
        #pragma unroll
        for (int i = 0; i < frag_acc_master.num_elements; ++i){
            // frag_acc_master.x[i] += tileAs * frag_acc_block.x[i] * tileBs;
            // frag_acc_master.x[i] = fmaf(tileAs * frag_acc_block.x[i], tileBs, frag_acc_master.x[i]);
            frag_acc_master.x[i] = fmaf(scale, frag_acc_block.x[i], frag_acc_master.x[i]);
        }
    
        // __syncthreads();

    }// end for (int kbase = 0; kbase < K; kbase += 128) 
    int warp_idx_y = threadIdx.y;
    int row_base_C = (blockIdx.y * NUM_WARP_Y + warp_idx_y) * WMMA_M;
    if (row_base_C >= N) return;

    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 >= M) return;    
    
    // Convert fp32 to bf16 and store
    // TODO: change WMMA_PLACEHOLDER from 8 to 16
    auto frag_result = rocwmma::fragment<rocwmma::accumulator, WMMA_M, WMMA_N, WMMA_PLACEHOLDER, rocwmma::bfloat16_t>(); //Placeholder frag
    
    #pragma unroll
    for (int i = 0; i < frag_acc_master.num_elements; ++i) {
        frag_result.x[i] = (rocwmma::bfloat16_t)frag_acc_master.x[i];
    }
    
    rocwmma::store_matrix_sync(C_ptr + row_base_C + N * col_base_C, frag_result, N, rocwmma::mem_col_major);

}

#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) balance_old_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
    18.5 KD of LDS
    */
    
    
    
    __shared__ __hip_fp8_e4m3_fnuz tileA[128][72]; // row-major
    __shared__ __hip_fp8_e4m3_fnuz tileB[128][72]; // col-major, = row-major transposed
    int lda = 72;
    int ldb = 72;
    __shared__ float tileAs[128];

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

    // int lda = M;
    // int ldb = N;
    cg::thread_block this_tb = cg::this_thread_block();
    auto group64 = cg::tiled_partition<64>(this_tb); // 1024 threads -> 64 x 16
    int t64x = group64.thread_rank();
    int t64y = group64.meta_group_rank(); 
    
    for (int kbase = 0; kbase < K; kbase += K_JUMP) {

        // Load SMEM: tileA, tileB, tileAs, tileBs
        // TRY: test tiled partition 64   

        // ----------------------------- FIRST HALF --------------------------
        #pragma unroll
        for (int kmini_load = 0; kmini_load < 64; kmini_load += 16){  // 16 is num group
            tileA[t64x][kmini_load + t64y] = A_ptr[blockIdx.y * WMMA_M * NUM_WARP_Y + t64x + M * (kbase + kmini_load + t64y)];
            tileA[t64x + 64][kmini_load + t64y] = A_ptr[blockIdx.y * WMMA_M * NUM_WARP_Y + t64x + 64 + M * (kbase + kmini_load + t64y)];
            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; 

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


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

        auto frag_a = 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 1
        for (int kmini = 0; kmini < 64; kmini += WMMA_K) {           
            rocwmma::load_matrix_sync(frag_a, &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, frag_a, frag_b, frag_acc_block);
        } // end for kmini

        __syncthreads();

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

        __syncthreads();

        #pragma unroll
        for (int kmini = 0; kmini < 64; kmini += WMMA_K) {           
            rocwmma::load_matrix_sync(frag_a, &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, frag_a, frag_b, frag_acc_block);
        } // end for kmini
        // Apply scaler


        // __syncthreads();

        int place = warp_idx_y * WMMA_M + 4 * (t64x / 32);;

        // i = 0
        frag_acc_master.x[0]  += tileAs[place + 0]  * frag_acc_block.x[0];
        frag_acc_master.x[1]  += tileAs[place + 1]  * frag_acc_block.x[1];
        frag_acc_master.x[2]  += tileAs[place + 2]  * frag_acc_block.x[2];
        frag_acc_master.x[3]  += tileAs[place + 3]  * frag_acc_block.x[3];

        // i = 1
        frag_acc_master.x[4]  += tileAs[place + 8]  * frag_acc_block.x[4];
        frag_acc_master.x[5]  += tileAs[place + 9]  * frag_acc_block.x[5];
        frag_acc_master.x[6]  += tileAs[place + 10] * frag_acc_block.x[6];
        frag_acc_master.x[7]  += tileAs[place + 11] * frag_acc_block.x[7];

        // i = 2
        frag_acc_master.x[8]  += tileAs[place + 16] * frag_acc_block.x[8];
        frag_acc_master.x[9]  += tileAs[place + 17] * frag_acc_block.x[9];
        frag_acc_master.x[10] += tileAs[place + 18] * frag_acc_block.x[10];
        frag_acc_master.x[11] += tileAs[place + 19] * frag_acc_block.x[11];

        // i = 3
        frag_acc_master.x[12] += tileAs[place + 24] * frag_acc_block.x[12];
        frag_acc_master.x[13] += tileAs[place + 25] * frag_acc_block.x[13];
        frag_acc_master.x[14] += tileAs[place + 26] * frag_acc_block.x[14];
        frag_acc_master.x[15] += tileAs[place + 27] * frag_acc_block.x[15];

        __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 + warp_idx_y) * WMMA_M;
    if (row_base_C >= M) return;

    // Convert fp32 to bf16 and store
    auto frag_result = rocwmma::fragment<rocwmma::accumulator, WMMA_M, WMMA_N, WMMA_PLACEHOLDER, rocwmma::bfloat16_t>(); //Placeholder frag
    
    #pragma unroll
    for (int i = 0; i < frag_acc_master.num_elements; ++i) {
        frag_result.x[i] = (rocwmma::bfloat16_t)frag_acc_master.x[i];
    }

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

    //int warp_idx_y = threadIdx.y;
    //int row_base_C = (blockIdx.y * NUM_WARP_Y + warp_idx_y) * WMMA_M;
      
    rocwmma::store_matrix_sync(C_ptr + row_base_C * N + col_base_C, frag_result, N, rocwmma::mem_row_major);
}

#define K_JUMP 256
#define WMMA_M 16
#define WMMA_N 16
#define WMMA_K 32
#define NUM_WARP_X 4
#define NUM_WARP_Y 4
#define WARP_SIZE 64
#define WMMA_PLACEHOLDER 8

__global__ void __launch_bounds__(1024, 8) k_large_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) {
    /*
    Perform B @ A.T = C.T, so we can store scale As in a VGPR of each thread in a warp
    B: NxK col-maj -- Bs (128x128 scaling): N/128 x K/128 col-maj
    A.T: KxM row-maj -- As (1x128 scaling): K/ 128 x M row-maj
    C.T: NxM col-maj
    M, N are divisible by 64, K is divisible by 256

    SMEM with padding per block:
    64 * (256 + 8) / 1024 + 64 * (224 + 8) / 1024 = 16.5 + 14.5 = 31 KB per block

    */
    int sn = (N + 128 - 1) / 128;

    __shared__ __hip_fp8_e4m3_fnuz tileA[64][256]; // col-maj = row-maj transposed 
    __shared__ __hip_fp8_e4m3_fnuz tileB[64][256]; // row-maj

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

    cg::thread_block this_tb = cg::this_thread_block();

    for (int kbase = 0; kbase < K; kbase += K_JUMP) {
        int warp_idx_y = threadIdx.y;
        int row_base_C = (blockIdx.y * NUM_WARP_Y + warp_idx_y) * WMMA_M;
    
        int warp_idx_x = threadIdx.x / WARP_SIZE;    
        int col_base_C = (blockIdx.x * NUM_WARP_X + warp_idx_x) * WMMA_N;

        int t64x = this_tb.thread_rank() % 64;
        int t64y = this_tb.thread_rank() / 64;

        #pragma unroll 4
        for (int kmini = 0; kmini < 128; kmini += 16){ // 16 is num group
            int K_idx_shifty = (kmini + t64y + 8 * (t64x % 16)) % 128;            

            tileB[t64x][K_idx_shifty] = B_ptr[(kbase + kmini + t64y) * N + blockIdx.y * WMMA_N * NUM_WARP_Y + t64x];   
            tileB[t64x][128 + K_idx_shifty] = B_ptr[(128 + kbase + kmini + t64y) * N + blockIdx.y * WMMA_N * NUM_WARP_Y + t64x];
        }   
        #pragma unroll
        for (int kmini = 0; kmini < 128; kmini += 16){ // 16 is num group
            int K_idx_shifty = (kmini + t64y + 8 * (t64x % 16)) % 128;            
            tileA[t64x][K_idx_shifty] = A_ptr[blockIdx.x * WMMA_M * NUM_WARP_X + t64x + M * (kbase + kmini + t64y)];       
       
            tileA[t64x][128 + K_idx_shifty] = A_ptr[blockIdx.x * WMMA_M * NUM_WARP_X + t64x + M * (128 + kbase + kmini + t64y)];
            
        } 

        float tileAs1 = As_ptr[(kbase / 128) * M + col_base_C + t64x % WMMA_N];
        float tileBs1 = Bs_ptr[(kbase / 128) * sn + (blockIdx.y * NUM_WARP_Y * WMMA_M) / 128];
        float scale1 = tileAs1 * tileBs1;
        float tileAs2 = As_ptr[(1 + kbase / 128) * M + col_base_C + t64x % WMMA_N];
        float tileBs2 = Bs_ptr[(1 + kbase / 128) * sn + (blockIdx.y * NUM_WARP_Y * WMMA_M) / 128];
        float scale2 = tileAs2 * tileBs2;

        __syncthreads(); 
        // __threadfence_block();
        // MFMA

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

        auto frag_a = 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>();

        #pragma unroll
        for (int kmini = 0; kmini < 128; kmini += WMMA_K){
            int tileB_row = warp_idx_y * WMMA_M + t64x % 16;
            int K_idx_base = kmini + (t64x / 16) * 8;
            int K_idx_shiftyB = (K_idx_base + 8 * tileB_row) % 128;
            // memcpy(&frag_a.x[0], &tileB[tileB_row][K_idx_shiftyB], sizeof(long));
            memcpy(&frag_a.x[0], &tileB[0][0] + tileB_row * 256 + K_idx_shiftyB, sizeof(long)); 

            int tileA_col = warp_idx_x * WMMA_N + t64x % 16;
            // int K_idx_shiftyA = (K_idx_base + 8 * tileA_col) % 128;
            // memcpy(&frag_b.x[0], &tileA[tileA_col][K_idx_shiftyB], sizeof(long));
            memcpy(&frag_b.x[0], &tileA[0][0] + tileA_col * 256 + K_idx_shiftyB, sizeof(long));
            rocwmma::mma_sync(frag_acc_block, frag_a, frag_b, frag_acc_block);
        }


        // Mult scalers
        #pragma unroll
        for (int i = 0; i < frag_acc_master.num_elements; ++i){
            frag_acc_master.x[i] = fmaf(scale1, frag_acc_block.x[i], frag_acc_master.x[i]);            
        }

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

        #pragma unroll
        for (int kmini = 0; kmini < 128; kmini += WMMA_K){
            int tileB_row = warp_idx_y * WMMA_M + t64x % 16;
            int K_idx_base = kmini + (t64x / 16) * 8;
            int K_idx_shiftyB = (K_idx_base + 8 * tileB_row) % 128;
            // memcpy(&frag_a.x[0], &tileB[tileB_row][128 + K_idx_shiftyB], sizeof(long));
            memcpy(&frag_a.x[0], &tileB[0][0] + tileB_row * 256 + 128 + K_idx_shiftyB, sizeof(long)); 

            int tileA_col = warp_idx_x * WMMA_N + t64x % 16;
            // int K_idx_shiftyA = (K_idx_base + 8 * tileA_col) % 128;
            // memcpy(&frag_b.x[0], &tileA[tileA_col][128 + K_idx_shiftyB], sizeof(long));
            memcpy(&frag_b.x[0], &tileA[0][0] + tileA_col * 256 + 128 + K_idx_shiftyB, sizeof(long));

            rocwmma::mma_sync(frag_acc_block2, frag_a, frag_b, frag_acc_block2);
        } 



        // Mult scalers
        #pragma unroll
        for (int i = 0; i < frag_acc_master.num_elements; ++i){
            frag_acc_master.x[i] = fmaf(scale2, frag_acc_block2.x[i], frag_acc_master.x[i]);            
        }


    }// end for (int kbase = 0; kbase < K; kbase += 128) 
    int warp_idx_y = threadIdx.y;
    int row_base_C = (blockIdx.y * NUM_WARP_Y + warp_idx_y) * WMMA_M;
    if (row_base_C >= N) return;

    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 >= M) return; 
        
    // Convert fp32 to bf16 and store
    // TODO: change WMMA_PLACEHOLDER from 8 to 16
    auto frag_result = rocwmma::fragment<rocwmma::accumulator, WMMA_M, WMMA_N, WMMA_PLACEHOLDER, rocwmma::bfloat16_t>(); //Placeholder frag
    
    #pragma unroll
    for (int i = 0; i < frag_acc_master.num_elements; ++i) {
        frag_result.x[i] = (rocwmma::bfloat16_t)frag_acc_master.x[i];
    }
       
    rocwmma::store_matrix_sync(C_ptr + row_base_C + N * col_base_C, frag_result, N, rocwmma::mem_col_major);
}

constexpr int K_VALUE = 7168;
constexpr int M_VALUE = 6144;
constexpr int N_VALUE = 4608;
constexpr int N_VALUE2 = 1536;
constexpr int THRESHOLD = 1024 * 576; // 589,824

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);
    
    if (K != K_VALUE) {
        dim3 grid_size((N + num_warp_x * 32 - 1) / (num_warp_x * 32), 
                    (M + num_warp_y * 32 - 1) / (num_warp_y * 32));
        balance_old_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);
        return;
    }

    if (M == M_VALUE && (N == N_VALUE || N == N_VALUE2)) {
        dim3 grid_size((N + num_warp_x * 32 - 1) / (num_warp_x * 32), 
                    (M + num_warp_y * 32 - 1) / (num_warp_y * 32));
        balance_old_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);
        return;
    }
    
    if (M * N <= THRESHOLD) {
        dim3 grid_size((M + num_warp_x * 16 - 1) / (num_warp_x * 16), 
                    (N + num_warp_y * 16 - 1) / (num_warp_y * 16));
        k_large_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);
                          
    }

    else {
        dim3 grid_size((M + num_warp_x * 32 - 1) / (num_warp_x * 32), 
                    (N + num_warp_y * 32 - 1) / (num_warp_y * 32));
        balance_new_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);
        return;    
    }
}
"""

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

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 eval.py leaderboard /tmp/tmpskcoatxm
[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: 4
ninja: no work to do.
