# 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
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 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) {
    // Compute global indices for this thread's output element
    int cx = blockIdx.x * blockDim.x + threadIdx.x;  // row index (M dimension)
    int cy = blockIdx.y * blockDim.y + threadIdx.y;  // column index (N dimension)
    bool valid = (cx < m && cy < n);  // flag if this thread corresponds to a valid C element

    // Shared memory tiles for A and B (tile size: 16x128 for each, matching blockDim and BLOCK)
    __shared__ __hip_fp8_e4m3_fnuz A_tile[16][BLOCK];
    __shared__ __hip_fp8_e4m3_fnuz B_tile[16][BLOCK];

    // Precompute number of N-blocks for scaling (N dimension grouped by 128)
    int sn = (n + BLOCK - 1) / BLOCK;

    // FP32 accumulator for this thread's result (to accumulate high-precision sum)
    float result_accum = 0.0f;

    // Loop over K dimension in BLOCK-sized (128) chunks
    for (int p = 0; p < k; p += BLOCK) {
        // Load a 16x128 submatrix of A into shared memory (coalesced by grouping threads by A row)
        if (cx < m) {
            #pragma unroll
            for (int j = 0; j < BLOCK; j += 16) {  // each iteration covers 16 columns
                int colIdx = p + j + threadIdx.y;
                if (colIdx < k) {
                    // Each thread with same threadIdx.x loads one element of A for 16 different columns
                    A_tile[threadIdx.x][j + threadIdx.y] = a[colIdx * m + cx];
                }
            }
        }
        // Load a 16x128 submatrix of B into shared memory (coalesced by grouping threads by B row)
        if (cy < n) {
            #pragma unroll
            for (int j = 0; j < BLOCK; j += 16) {
                int colIdx = p + j + threadIdx.x;
                if (colIdx < k) {
                    // Each thread with same threadIdx.y loads one element of B for 16 different columns
                    B_tile[threadIdx.y][j + threadIdx.x] = b[colIdx * n + cy];
                }
            }
        }
        __syncthreads();  // ensure the whole A_tile and B_tile are loaded

        // Compute partial dot-product for the current 128-length block
        float block_sum = 0.0f;
        if (valid) {
            #pragma unroll 4
            for (int ii = 0; ii < BLOCK; ++ii) {
                // Convert FP8 values to float for computation:contentReference[oaicite:0]{index=0}
                float a_val = (float) A_tile[threadIdx.x][ii];
                float b_val = (float) B_tile[threadIdx.y][ii];
                block_sum += a_val * b_val;  // accumulate product
            }
            // Apply block-wise scaling factors (constant per 128-element block):contentReference[oaicite:1]{index=1}
            float a_scale = as[cx + (p / BLOCK) * m];
            float b_scale = bs[(cy / BLOCK) + (p / BLOCK) * sn];
            result_accum += block_sum * a_scale * b_scale;
        }
        __syncthreads();  // sync before loading the next K-block into shared memory
    }

    // Store the result in BF16 (only if within valid output bounds)
    if (valid) {
        c[cx * n + cy] = (__hip_bfloat16) result_accum;
    }
}


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);
    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++"
os.environ["PYTORCH_ROCM_ARCH"] = "gfx942"

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: 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++ 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/tmpf3mee_18
[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.
