# This script provides a template for using load_inline to run a HIP kernel for
import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

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 = """
#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>
#include <rocwmma/internal/type_traits.hpp>

using namespace rocwmma;

// Check these for compute accumulate from docs
// Types
using InputT   = __hip_fp8_e5m2;
using OutputT  = __hip_bfloat16;
using ComputeT = float32_t;

namespace gfx9Params
{
    enum kernelParams : uint32_t
    {
        ROCWMMA_M = 32u,
        ROCWMMA_N = 32u,
        ROCWMMA_K = 16u,
        BLOCKS_X  = 2u,
        BLOCKS_Y  = 2u,
        TBLOCK_X  = 128u,
        TBLOCK_Y  = 2u,
        WARP_SIZE = Constants::AMDGCN_WAVE_SIZE_64
    };
}

namespace gfx11Params
{
    enum kernelParams : uint32_t
    {
        ROCWMMA_M = 16u,
        ROCWMMA_N = 16u,
        ROCWMMA_K = 16u,
        BLOCKS_X  = 4u,
        BLOCKS_Y  = 2u,
        TBLOCK_X  = 64u,
        TBLOCK_Y  = 4u,
        WARP_SIZE = Constants::AMDGCN_WAVE_SIZE_32
    };
}


// HIP Host functions to determine the gfx architecture
bool isGfx9(){return true;}



#if(ROCWMMA_ARCH_GFX9)
using namespace gfx9Params;
#else
using namespace gfx11Params;
#endif // defined(ROCWMMA_ARCH_GFX9)

uint32_t getWarpSize(){return 64;}

ROCWMMA_KERNEL void __launch_bounds__(256) gemm_rocwmma_d(uint32_t       m,
                                                          uint32_t       n,
                                                          uint32_t       k,
                                                          InputT*  a,
                                                          InputT*  b,
                                                          OutputT* c,
                                                          OutputT*       d,
                                                          uint32_t       lda,
                                                          uint32_t       ldb,
                                                          uint32_t       ldc,
                                                          uint32_t       ldd,
                                                          ComputeT       alpha,
                                                          ComputeT       beta)
{

}


constexpr const int BLOCK = 128;

__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, int m, int n, int k) {
                   
    // Your implementation here
    int cx = threadIdx.x + blockDim.x * blockIdx.x;
    int cy = threadIdx.y + blockDim.y * blockIdx.y;
    if(cx >= m || cy >= n) return;
    
    int sn = (n + BLOCK - 1) / BLOCK;
    
    float result = 0;
    // split loop into an outer loop over different blocks, and an inner loop within one block.
    // we can assume k % BLOCK == 0.
    for(int i = 0; i < k; i += BLOCK) {
        // block results accumulates the inner product across a single block.
        // within each block, scales are constant, so we can lift the scaling 
        // outside of the inner loop.
        float block_result = 0;
        for(int ii = 0; ii < BLOCK; ++ii) {
            // load input matrix elements and convert to float for computations
            float av = (float)a[cx + (i + ii) * m];
            float bv = (float)b[cy + (i + ii) * n];
            block_result += av * bv; 
        }
        
        // before we can go to the next block, scale the result of the current block
        // and accumulate to final result
        // note the different indexing into as and bs
        result += block_result * as[cx + i/BLOCK * m] * bs[cy/BLOCK + i/BLOCK * sn];
    }
    
    // finally, write the result as bf16
    c[cx * n + cy] = (__hip_bfloat16)result;
}

void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c) {
    // Runtime checks for host parameters
    uint32_t hTBLOCK_X    = isGfx9() ? gfx9Params::TBLOCK_X : gfx11Params::TBLOCK_X;
    uint32_t hTBLOCK_Y    = isGfx9() ? gfx9Params::TBLOCK_Y : gfx11Params::TBLOCK_Y;
    uint32_t hBLOCKS_X    = isGfx9() ? gfx9Params::BLOCKS_X : gfx11Params::BLOCKS_X;
    uint32_t hBLOCKS_Y    = isGfx9() ? gfx9Params::BLOCKS_Y : gfx11Params::BLOCKS_Y;
    uint32_t hROCWMMA_M   = isGfx9() ? gfx9Params::ROCWMMA_M : gfx11Params::ROCWMMA_M;
    uint32_t hROCWMMA_N   = isGfx9() ? gfx9Params::ROCWMMA_N : gfx11Params::ROCWMMA_N;
    uint32_t hROCWMMA_K   = isGfx9() ? gfx9Params::ROCWMMA_K : gfx11Params::ROCWMMA_K;
    uint32_t hWARP_TILE_X = hBLOCKS_X * hROCWMMA_M;
    uint32_t hWARP_TILE_Y = hBLOCKS_Y * hROCWMMA_N;

    // Runtime warp calculation (host code needs to query warpsize dynamically)
    auto warpSize = getWarpSize();
    auto macroTileSize
        = rocwmma::make_coord2d(hTBLOCK_X / warpSize * hWARP_TILE_X, hTBLOCK_Y * hWARP_TILE_Y);



    auto m = a.size(0);
    auto n = b.size(0);
    auto k = a.size(1);
    auto d_a =  reinterpret_cast<InputT*>(a.data_ptr());
    auto d_b = reinterpret_cast<InputT*>(b.data_ptr());
    auto d_c = reinterpret_cast<OutputT*>(c.data_ptr());
    auto d_d = reinterpret_cast<OutputT*>(c.data_ptr());
    // check all of these
    auto lda = a.stride(0);
    auto ldb = b.stride(1);
    auto ldc = c.stride(1);
    auto ldd = ldc;
    auto alpha = 1.0f;
    auto beta = 0.0f;

        // Uses 2 lds blocks for prefetch loop (A and B)
    int ldsusage
        = 2u * sizeof(InputT) * (get<0>(macroTileSize) + get<1>(macroTileSize)) * hROCWMMA_K;

    auto blockDim = dim3(hTBLOCK_X, hTBLOCK_Y);
    auto gridDim  = dim3(rocwmma::ceilDiv(m, get<0>(macroTileSize)),
                        rocwmma::ceilDiv(n, get<1>(macroTileSize)));
    gemm_rocwmma_d<<<gridDim, blockDim, ldsusage, 0>>>(
                            m,
                            n,
                            k,
                            d_a,
                            d_b,
                            d_c,
                            d_d,
                            lda,
                            ldb,
                            ldc,
                            ldd,
                            alpha,
                            beta);


    custom_kernel<<<dim3((m+15)/16, (n+15)/16), dim3(16, 16), 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);
    //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 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: 2
[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 
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:124:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  124 |     uint32_t hTBLOCK_X    = isGfx9() ? gfx9Params::TBLOCK_X : gfx11Params::TBLOCK_X;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:125:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  125 |     uint32_t hTBLOCK_Y    = isGfx9() ? gfx9Params::TBLOCK_Y : gfx11Params::TBLOCK_Y;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:126:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  126 |     uint32_t hBLOCKS_X    = isGfx9() ? gfx9Params::BLOCKS_X : gfx11Params::BLOCKS_X;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:127:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  127 |     uint32_t hBLOCKS_Y    = isGfx9() ? gfx9Params::BLOCKS_Y : gfx11Params::BLOCKS_Y;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:128:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  128 |     uint32_t hROCWMMA_M   = isGfx9() ? gfx9Params::ROCWMMA_M : gfx11Params::ROCWMMA_M;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:129:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  129 |     uint32_t hROCWMMA_N   = isGfx9() ? gfx9Params::ROCWMMA_N : gfx11Params::ROCWMMA_N;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:130:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  130 |     uint32_t hROCWMMA_K   = isGfx9() ? gfx9Params::ROCWMMA_K : gfx11Params::ROCWMMA_K;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~~
7 warnings generated when compiling for gfx942.
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:124:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  124 |     uint32_t hTBLOCK_X    = isGfx9() ? gfx9Params::TBLOCK_X : gfx11Params::TBLOCK_X;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:125:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  125 |     uint32_t hTBLOCK_Y    = isGfx9() ? gfx9Params::TBLOCK_Y : gfx11Params::TBLOCK_Y;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:126:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  126 |     uint32_t hBLOCKS_X    = isGfx9() ? gfx9Params::BLOCKS_X : gfx11Params::BLOCKS_X;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:127:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  127 |     uint32_t hBLOCKS_Y    = isGfx9() ? gfx9Params::BLOCKS_Y : gfx11Params::BLOCKS_Y;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:128:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  128 |     uint32_t hROCWMMA_M   = isGfx9() ? gfx9Params::ROCWMMA_M : gfx11Params::ROCWMMA_M;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:129:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  129 |     uint32_t hROCWMMA_N   = isGfx9() ? gfx9Params::ROCWMMA_N : gfx11Params::ROCWMMA_N;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~~
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm/hip.hip:130:38: warning: conditional expression between different enumeration types ('gfx9Params::kernelParams' and 'gfx11Params::kernelParams') is deprecated [-Wdeprecated-enum-compare-conditional]
  130 |     uint32_t hROCWMMA_K   = isGfx9() ? gfx9Params::ROCWMMA_K : gfx11Params::ROCWMMA_K;
      |                                      ^ ~~~~~~~~~~~~~~~~~~~~~   ~~~~~~~~~~~~~~~~~~~~~~
7 warnings generated when compiling for host.
[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/tmpwzip7ihw
[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: 2
ninja: no work to do.
