#!POPCORN leaderboard amd-fp8-mm 
#!POPCORN GPU MI300
# 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
import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'
CPP_WRAPPER = """
void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c);
"""

CUDA_SRC = """
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

constexpr  int COARSENING = 2; // increase coarsening to reduce thread sequential processing
constexpr uint BM = 64, BN = 64 , BK = 8, TM = 8;


__global__ void custom_kernel(const __hip_fp8_e4m3_fnuz* a, const __hip_fp8_e4m3_fnuz* b, const float* as, const float* bs, 
                   __hip_bfloat16* c, const uint M, const uint N, const uint K) {

    // converting 1D block Dim to 2D block Dim

    const uint cCol = blockIdx.x;
    const uint cRow = blockIdx.y ;
    // TODO: ajust this to each thread to load multiple entries and 
    // better exploit the cache sizes
    
    //uint col_no = cCol * BN ;
    //uint row_no = cRow * BM
    const uint threadCol = threadIdx.x % BN;
    const uint threadRow = threadIdx.x / BN;
    
    __shared__ __hip_fp8_e4m3_fnuz shared_a[BM * BK];
    __shared__ __hip_fp8_e4m3_fnuz shared_b[BN * BK];
    __shared__ float shared_as[BM * ( (BK + 127) / 128)]; // 127 = 128 - 1 for padding
    __shared__ float shared_bs[((BN + 127) / 128) * ( (BK + 127) / 128)];

    // loadi a, b, as, bs
    
    // TODO: convert this return to proper if to enter the execution
   

    //int next = 1;
    //int current = 0;
    
    // this place C matrix Row is the row for A and C matrix Col is the row for B which needs to be accessed.
    
    a += cRow * BM;
    b += cCol * BN;
    c += cRow * BM * N + cCol * BN;


    //const uint innerColA = threadIdx.x / BM;
    //const uint innerRowA = threadIdx.x % BM;
    
    //const uint innerRowB = threadIdx.x % BN;
    //const uint innerColB = threadIdx.x / BN;


    float  threadResults[TM] = {0.0};
    float  scales[TM] = {0.0};
    

    for (int k = 0; k < K ; k+=BK) {
        
        // (to access base pointer of this block) + (accessing specific tile of the block) + ( to access the specific value inside this tile)
        
        // TODO: need to correct this indexing
        // ERROR: this is wrong indexing
        
        #pragma unroll 8
        for (int resIdx = 0 ; resIdx < TM; resIdx++){
            // Loading values using threadRow and threadCol. since they are the indices based on which as and bs need to be accessed. not innerRowA and innerColB
            
            //int k_col =  k  / 128 ;
            //int scaleM = cRow * BM + threadRow * TM  + resIdx;
            //int scaleIdxA =  k_col * M + cRow * BM + threadRow * TM  + resIdx;
            //int scaleN  = cCol * BN + threadCol;
            //int scaleIdxB =  ( ( cCol * BN + threadCol) / 128 ) + k_col * ( ( N + 128 - 1 ) / 128 ) ;
            
            
            
            if( cRow * BM + threadRow * TM  + resIdx < M  &&   cCol * BN + threadCol < N ){
                
                scales[resIdx] = as[(k  / 128) * M + cRow * BM + threadRow * TM  + resIdx] * bs[( ( cCol * BN + threadCol) / 128 ) + (k  / 128) * ( ( N + 128 - 1 ) / 128 )];
                
            }
            else{

                scales[resIdx] = 0.0f;

            }

        }
        // (block_ptr) + (kth tile pointer ) + (col,trow) value to load : here we load (col, trow) for memory coalescing
        // due to column major matrixs a & b.
        
        if ( cRow * BM + ( threadIdx.x % BM)  < M ){
            
            shared_a[ ( threadIdx.x % BM) + BM * (threadIdx.x / BM) ] = a[ ( threadIdx.x % BM) + M * (threadIdx.x / BM) ];
            
        }
        else{
            
            shared_a[ ( threadIdx.x % BM) + BM * (threadIdx.x / BM) ] = __hip_fp8_e4m3_fnuz(0.0f);
            
        }
        
        shared_b[ (threadIdx.x % BN) + BN * (threadIdx.x / BN) ] = b[ (threadIdx.x % BN) +  N * (threadIdx.x / BN) ];

        
        __syncthreads();

        a +=  BK * M;
        b +=  BK * N;
        
        float block_results[TM] = {0.0};
        
        #pragma unroll 8
        for(int ii = 0; ii < BK; ++ii){

            // NOTE: here using threadCol to access not innerColB as threadCol provides the col index for B matrix and threadRow provides the row index for A matrix
            float tmpB  = (float)shared_b[BN * ii + threadCol];
            
            for (uint resIdx = 0 ; resIdx < TM; resIdx++){
                // NOTE:  here using threadRow to access not innerRowA as threadrow provides the row index for A matrix and threadCol provides the col index for B matrix
                block_results[resIdx] += (float)shared_a[BM * ii + threadRow * TM + resIdx ] * tmpB;
            }
            
        }
        
        #pragma unroll 8
        for (uint resIdx = 0 ; resIdx < TM; resIdx++){
            threadResults[resIdx] += block_results[resIdx] * scales[resIdx];
        }

        //current ^= 1;
        //next ^= 1;
        __syncthreads();

    }
    
    #pragma unroll
    for (uint resIdx = 0 ; resIdx < TM; resIdx++){
        if (cRow * BM + threadRow * TM + resIdx < M && cCol * BN + threadCol < N){
            c[ (threadRow * TM + resIdx) * N + threadCol ] = threadResults[resIdx];
        }
    }

}

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);
    
    // Optimized grid and block dimensions for MI300X
    dim3 block( ( BM * BN ) / TM);  // 512 threads per block now
    
    dim3 grid( (n + BN - 1) / BN, (m + BM - 1) / BM  );
    
    custom_kernel<<<grid, block, 0, 0>>>(
        (__hip_fp8_e4m3_fnuz*)a.data_ptr(),
        (__hip_fp8_e4m3_fnuz*)b.data_ptr(),
        as.data_ptr<float>(),
        bs.data_ptr<float>(),
        (__hip_bfloat16*)c.data_ptr(),
        m, n, k
    );
}
"""


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
    # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    module.fp8_mm(a, b, a_scale, b_scale, c)
    # prof.export_chrome_trace("resnet18_profile.json")

    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 -fno-gpu-rdc --offload-arch=gfx942 -std=c++20 -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip -o hip.cuda.o 
[3/3] 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

python eval.py test /tmp/tmpzvaozxuz
[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
ninja: no work to do.
