# This script provides a template for using load_inline to run a HIP kernel for
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t # Assuming task.py defines these

# Updated CPP_WRAPPER to reflect the function signature in HIP_SRC
CPP_WRAPPER = """
void fp8_mm_hip(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c);
"""

# Renamed CUDA_SRC to HIP_SRC for clarity and updated with optimized kernel
HIP_SRC = """
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

// Scaling block size for K dimension (as defined by problem, e.g., 128)
constexpr const int SCALE_BLOCK_K = 128;

// Thread block dimensions for C tile (output-stationary)
constexpr const int BDIM_M = 16; // Tile dimension for M
constexpr const int BDIM_N = 16; // Tile dimension for N

// K-dimension step for shared memory tiling
constexpr const int K_SHARED_STEP = 16;

__global__ void custom_hip_kernel(
    const __hip_fp8_e4m3_fnuz* __restrict__ a_ptr, // M x K, col-major
    const __hip_fp8_e4m3_fnuz* __restrict__ b_ptr, // N x K, col-major
    const float* __restrict__ a_scale_ptr,      // M x K/SCALE_BLOCK_K, col-major
    const float* __restrict__ b_scale_ptr,      // N/SCALE_BLOCK_K x K/SCALE_BLOCK_K, col-major
    __hip_bfloat16* __restrict__ c_ptr,         // M x N, row-major
    int M, int N, int K) {

    // Thread indices within the block
    int tx = threadIdx.x; // 0..BDIM_M-1, corresponds to M-dimension index in C tile
    int ty = threadIdx.y; // 0..BDIM_N-1, corresponds to N-dimension index in C tile

    // Global row and column in C that this thread computes
    int global_c_row = blockIdx.x * BDIM_M + tx;
    int global_c_col = blockIdx.y * BDIM_N + ty;

    __shared__ __hip_fp8_e4m3_fnuz sA[BDIM_M][K_SHARED_STEP];
    __shared__ __hip_fp8_e4m3_fnuz sB[BDIM_N][K_SHARED_STEP];

    float total_acc_fp32 = 0.0f;

    int b_scale_num_n_tiles = (N + SCALE_BLOCK_K - 1) / SCALE_BLOCK_K;

    for (int k_block_base = 0; k_block_base < K; k_block_base += SCALE_BLOCK_K) {
        float scale_a = 1.0f;
        float scale_b = 1.0f;

        if (global_c_row < M) {
             scale_a = a_scale_ptr[global_c_row + (k_block_base / SCALE_BLOCK_K) * M];
        }
        if (global_c_col < N) {
            scale_b = b_scale_ptr[(global_c_col / SCALE_BLOCK_K) + (k_block_base / SCALE_BLOCK_K) * b_scale_num_n_tiles];
        }

        float current_k_block_sum_fp32 = 0.0f;

        for (int k_shared_offset = 0; k_shared_offset < SCALE_BLOCK_K; k_shared_offset += K_SHARED_STEP) {
            int current_k_tile_base = k_block_base + k_shared_offset;

            // Load A tile segment into sA
            if (global_c_row < M && (current_k_tile_base + ty) < K) {
                sA[tx][ty] = a_ptr[global_c_row + (current_k_tile_base + ty) * M];
            } else {
                sA[tx][ty] = __hip_fp8_e4m3_fnuz(0.0f);
            }

            // Load B tile segment into sB
             if (global_c_col < N && (current_k_tile_base + tx) < K) {
                sB[ty][tx] = b_ptr[global_c_col + (current_k_tile_base + tx) * N];
            } else {
                sB[ty][tx] = __hip_fp8_e4m3_fnuz(0.0f);
            }
            __syncthreads();

            // Apply pragma unroll to this loop to potentially reduce loop overhead
            #pragma unroll
            for (int k_s = 0; k_s < K_SHARED_STEP; ++k_s) {
                current_k_block_sum_fp32 += static_cast<float>(sA[tx][k_s]) * static_cast<float>(sB[ty][k_s]);
            }
            __syncthreads();
        }
        total_acc_fp32 += current_k_block_sum_fp32 * scale_a * scale_b;
    }

    if (global_c_row < M && global_c_col < N) {
        c_ptr[global_c_row * N + global_c_col] = static_cast<__hip_bfloat16>(total_acc_fp32);
    }
}

// C++ Wrapper function called by PyTorch
void fp8_mm_hip(torch::Tensor a_tensor, torch::Tensor b_tensor,
                torch::Tensor a_scale_tensor, torch::Tensor b_scale_tensor,
                torch::Tensor c_tensor) {
    int M_dim = a_tensor.size(0);
    int K_dim = a_tensor.size(1);
    int N_dim = b_tensor.size(0);

    dim3 threads_per_block(BDIM_M, BDIM_N);
    dim3 num_blocks((M_dim + BDIM_M - 1) / BDIM_M, (N_dim + BDIM_N - 1) / BDIM_N);

    hipLaunchKernelGGL(custom_hip_kernel, num_blocks, threads_per_block, 0, 0,
        reinterpret_cast<const __hip_fp8_e4m3_fnuz*>(a_tensor.data_ptr()),
        reinterpret_cast<const __hip_fp8_e4m3_fnuz*>(b_tensor.data_ptr()),
        a_scale_tensor.data_ptr<float>(),
        b_scale_tensor.data_ptr<float>(),
        reinterpret_cast<__hip_bfloat16*>(c_tensor.data_ptr()),
        M_dim, N_dim, K_dim
    );
}
"""

import os
os.environ["CXX"] = "clang++"

# Load the HIP kernel inline
module = load_inline(
    name='fp8_mm_hip_optimized_v2', # Changed module name to avoid conflicts
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[HIP_SRC],
    functions=['fp8_mm_hip'],
    verbose=True,
    # Added -O3 for explicitness, though it's often a default for release builds.
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20", "-O3"]
)

# The Python function that will be called by the evaluation script
def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    module.fp8_mm_hip(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_hip_optimized_v2/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v2/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v2/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v2/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_hip_optimized_v2 -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_hip_optimized_v2/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_hip_optimized_v2 -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 -std=c++20 -O3 -fno-gpu-rdc -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v2/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_hip_optimized_v2.so

python eval.py test /tmp/tmpegzeffxf
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v2/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v2/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v2/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v2/hip.hip [skipped, already hipified]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 0
ninja: no work to do.
