# 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/hip_runtime.h>
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>
#include <ATen/ATen.h>

#ifndef BLOCK
#define BLOCK 128            // depth of one scale slice
#endif

// tile sizes (original values)
#define TILE_M 128
#define TILE_N 128
#define TILE_K  16
#define WM 8                 // rows  per thread
#define WN 8                 // cols  per thread

// -------------------------------------------------------------------------
// Helpers
// -------------------------------------------------------------------------
__device__ inline float fp8_to_float(__hip_fp8_e4m3_fnuz v) { return (float)v; }
__device__ inline __hip_bfloat16 float_to_bf16(float v)      { return __float2bfloat16(v); }

// -------------------------------------------------------------------------
// Fused kernel
// -------------------------------------------------------------------------
__global__ void fp8_gemm_fused(const __hip_fp8_e4m3_fnuz* __restrict__ A_fp8,
                               const __hip_fp8_e4m3_fnuz* __restrict__ B_fp8,
                               const float*                __restrict__ AS,   // M × (K/128)
                               const float*                __restrict__ BS,   // (N/128) × (K/128)
                               __hip_bfloat16*             __restrict__ C,    // M × N (row-major)
                               int  M, int N, int K, int SN /* = N/128 */)
{
    //----------------------------------------------------------------------
    // Thread / block coordinates
    //----------------------------------------------------------------------
    int tx   = threadIdx.x;                 // 0..15
    int ty   = threadIdx.y;                 // 0..15
    int brow = blockIdx.x;                  // C-tile row  index
    int bcol = blockIdx.y;                  // C-tile col  index

    int row0 = brow * TILE_M + ty * WM;     // first row this thread owns
    int col0 = bcol * TILE_N + tx * WN;     // first col this thread owns

    //----------------------------------------------------------------------
    // Accumulator for this 8×8 micro-tile (float)
    //----------------------------------------------------------------------
    float acc[WM][WN] = {0.0f};

    //----------------------------------------------------------------------
    // Shared memory panels
    //----------------------------------------------------------------------
    extern __shared__ float smem[];
    float* As = smem;                                // 128 × 16  floats
    float* Bs = smem + TILE_M * TILE_K;              // 16  × 128 floats

    //----------------------------------------------------------------------
    // Loop over 128-column “scale slices” of K
    //----------------------------------------------------------------------
    for (int kSlice = 0; kSlice < K; kSlice += BLOCK)
    {
        int slice_id = kSlice / BLOCK;               // p in the math
        // each thread needs the row-scale for its 8 rows
        float sA[WM];
        #pragma unroll
        for (int ii = 0; ii < WM; ++ii)
            sA[ii] = AS[(row0 + ii) + slice_id * M]; // column-major

        // each thread needs the col-scale for its 8 cols
        float sB[WN];
        #pragma unroll
        for (int jj = 0; jj < WN; ++jj)
            sB[jj] = BS[(col0 + jj) / 128 + slice_id * SN];

        // zero the slice accumulator
        float blk[WM][WN] = {0.0f};

        //---------------- loop over the 128 columns in 16-col panels ------
        for (int kk = kSlice; kk < kSlice + BLOCK; kk += TILE_K)
        {
            // ---- 1. cooperative load & de-quant/scaling of A panel -----
            #pragma unroll
            for (int ii = 0; ii < WM; ++ii) {
                int a_row = row0 + ii;               // guaranteed < M
                int a_col = kk + tx;                 // 16 cols per warp
                float a_val = fp8_to_float(A_fp8[a_row + a_col * M])
                               * sA[ii];             // apply row-scale
                As[(ty*WM + ii) * TILE_K + tx] = a_val;
            }

            // ---- 2. cooperative load & de-quant/scaling of B panel -----
            #pragma unroll
            for (int jj = 0; jj < WN; ++jj) {
                int b_row = kk + ty;                 // 16 rows per warp
                int b_col = col0 + jj;               // one of 8 cols
                float b_val = fp8_to_float(B_fp8[b_col + b_row * N])
                               * sB[jj];             // apply col-scale
                Bs[ty * TILE_N + tx * WN + jj] = b_val;
            }

            __syncthreads();

            // ---- 3. 16-deep inner-K FMAs into blk ----------------------
            #pragma unroll
            for (int t = 0; t < TILE_K; ++t)
            {
                float a_frag[WM];
                #pragma unroll
                for (int ii = 0; ii < WM; ++ii)
                    a_frag[ii] = As[(ty*WM + ii) * TILE_K + t];

                float b_frag[WN];
                #pragma unroll
                for (int jj = 0; jj < WN; ++jj)
                    b_frag[jj] = Bs[t * TILE_N + tx * WN + jj];

                #pragma unroll
                for (int ii = 0; ii < WM; ++ii)
                    #pragma unroll
                    for (int jj = 0; jj < WN; ++jj)
                        blk[ii][jj] += a_frag[ii] * b_frag[jj];
            }
            __syncthreads();
        } // end 128-wide slice

        // ---- 4. add this slice’s contribution to the final accumulator --
        #pragma unroll
        for (int ii = 0; ii < WM; ++ii)
            #pragma unroll
            for (int jj = 0; jj < WN; ++jj)
                acc[ii][jj] += blk[ii][jj];          // already scaled
    }

    //----------------------------------------------------------------------
    // 5. Store 8×8 micro-tile to row-major C (BF16)
    //----------------------------------------------------------------------
    #pragma unroll
    for (int ii = 0; ii < WM; ++ii) {
        int c_row = row0 + ii;
        if (c_row >= M) break;                       // tail block safety
        #pragma unroll
        for (int jj = 0; jj < WN; ++jj) {
            int c_col = col0 + jj;
            if (c_col >= N) continue;
            C[c_row * N + c_col] = float_to_bf16(acc[ii][jj]);
        }
    }
}

// -------------------------------------------------------------------------
// Host helper (called from Python)
// -------------------------------------------------------------------------
void fp8_mm(torch::Tensor a_fp8,
            torch::Tensor b_fp8,
            torch::Tensor a_scale,
            torch::Tensor b_scale,
            torch::Tensor c_bf16)
{
    int M  = a_fp8.size(0);
    int K  = a_fp8.size(1);
    int N  = b_fp8.size(0);
    int SN = (N + 127) >> 7;              // exact by assumption

    dim3 threads(16, 16);            // 256 threads (8 warps)
    dim3 grid((M + TILE_M - 1) / TILE_M,
              (N + TILE_N - 1) / TILE_N);

    std::size_t shmem =
        (TILE_M * TILE_K + TILE_K * TILE_N) * sizeof(float); // 128×16 + 16×128

    hipLaunchKernelGGL(fp8_gemm_fused,
        grid, threads, shmem, 0,
        reinterpret_cast<const __hip_fp8_e4m3_fnuz*>(a_fp8.data_ptr()),
        reinterpret_cast<const __hip_fp8_e4m3_fnuz*>(b_fp8.data_ptr()),
        a_scale.data_ptr<float>(),
        b_scale.data_ptr<float>(),
        reinterpret_cast<__hip_bfloat16*>(c_bf16.data_ptr<at::BFloat16>()),
        M, N, K, SN);

    // Optional:
    // C10_HIP_CHECK( hipGetLastError() );
}

"""

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: 0
[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 benchmark /tmp/tmplqrvzqif
[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: 0
[1/1] 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
