# 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
// K_SHARED_STEP = 16 is chosen because shared memory load uses ty for sA's k-index
// and tx for sB's k-index, matching BDIM_N and BDIM_M respectively.
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 memory for tiles of A and B
    // sA stores a BDIM_M x K_SHARED_STEP tile of A. sA[local_M_idx][local_K_idx]
    // sB stores a BDIM_N x K_SHARED_STEP tile of B. sB[local_N_idx][local_K_idx]
    __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;

    // Calculate the number of N-dimension tiles for b_scale indexing
    // (N + SCALE_BLOCK_K - 1) / SCALE_BLOCK_K ensures correct handling if N is not a multiple of SCALE_BLOCK_K
    // Though problem states N is divisible by 64, it might not be by 128.
    int b_scale_num_n_tiles = (N + SCALE_BLOCK_K - 1) / SCALE_BLOCK_K;

    // Outer loop over K in chunks of SCALE_BLOCK_K (e.g., 128 elements)
    for (int k_block_base = 0; k_block_base < K; k_block_base += SCALE_BLOCK_K) {
        // Determine scales for this K-block.
        // These scales are constant for all K_SHARED_STEP iterations within this K-block.
        float scale_a = 1.0f; // Default scale if out of bounds
        float scale_b = 1.0f; // Default scale if out of bounds

        // Load scale for A, with boundary check for M dimension
        if (global_c_row < M) {
             scale_a = a_scale_ptr[global_c_row + (k_block_base / SCALE_BLOCK_K) * M];
        }
        // Load scale for B, with boundary check for N dimension
        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; // Accumulator for the current K-block (128 elements)

        // Inner loop: iterate through the current SCALE_BLOCK_K segment using smaller K_SHARED_STEP tiles
        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
            // sA[tx][ty] stores A[global_c_row, current_k_tile_base + ty]
            // (tx is local M-index, ty is local K-index for this load pattern)
            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); // Pad with zero if out of bounds
            }

            // Load B tile segment into sB
            // sB[ty][tx] stores B[global_c_col, current_k_tile_base + tx]
            // (ty is local N-index, tx is local K-index for this load pattern)
             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); // Pad with zero if out of bounds
            }
            __syncthreads(); // Ensure all threads complete loading sA and sB

            // Perform matrix multiplication for this K_SHARED_STEP tile using shared memory data
            // Each thread computes one element of the C tile.
            // sA[tx][k_s] gives A(global_c_row, k_val)
            // sB[ty][k_s] gives B(global_c_col, k_val)
            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(); // Ensure all threads complete computation before next tile load
        }
        // Accumulate result for this K-block, applying scales
        total_acc_fp32 += current_k_block_sum_fp32 * scale_a * scale_b;
    }

    // Write final accumulated result to C matrix, with boundary check
    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) {
    // Get matrix dimensions from input tensors
    // a_tensor: M x K
    // b_tensor: N x K
    // c_tensor: M x N (output)
    int M_dim = a_tensor.size(0);
    int K_dim = a_tensor.size(1);
    int N_dim = b_tensor.size(0);
    // int K_dim_b = b_tensor.size(1); // Should be equal to K_dim

    // Define thread block and grid dimensions for kernel launch
    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);

    // Launch the HIP kernel
    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
    );

    // Optional: Check for kernel launch errors (useful for debugging)
    // hipError_t err = hipGetLastError();
    // if (err != hipSuccess) {
    //    printf("HIP Kernel Launch Error: %s\\n", hipGetErrorString(err));
    // }
}
"""

import os
# Ensure clang++ is used as the HIP compiler driver by PyTorch's C++ extensions.
# For ROCm, hipcc is the typical compiler frontend. Setting CXX to clang++
# is a common practice if hipcc isn't picked up automatically or if specific clang++ flags are needed.
os.environ["CXX"] = "clang++" # As per the original template

# Load the HIP kernel inline using PyTorch's C++ extension utilities
# The function name in the 'functions' list must match the C++ wrapper function name.
module = load_inline(
    name='fp8_mm_hip_optimized', # Unique name for the loaded module
    cpp_sources=[CPP_WRAPPER],   # C++ wrapper code
    cuda_sources=[HIP_SRC],      # HIP kernel code (passed as 'cuda_sources')
    functions=['fp8_mm_hip'],    # List of C++ functions to make available in the module
    verbose=True,                # Enable verbose output during compilation
    # Extra flags for the HIP compiler (hipcc)
    # --offload-arch=gfx942 targets AMD MI300X GPU
    # -std=c++20 enables C++20 standard features, if used.
    extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20"]
)

# 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
    # Call the loaded HIP kernel via the C++ wrapper function
    module.fp8_mm_hip(a, b, a_scale, b_scale, c)
    return c

# Example usage block (typically not part of the final submission file, for testing)
if __name__ == '__main__':
    # This section can be used for local testing of the kernel.
    # It would require importing `generate_input` from `reference.py` or a similar utility.
    #
    # Example:
    # from reference import generate_input, ref_kernel # Assuming these are available
    #
    # M, N, K, SEED = 128, 128, 256, 42
    # test_input_data = generate_input(M, N, K, SEED) # Generates (a, b, a_scale, b_scale, c_output_buffer)
    # a_val, b_val, a_s_val, b_s_val, c_val_custom = test_input_data
    #
    # # Ensure c_val_custom is a fresh tensor for the custom kernel output
    # c_val_custom = torch.zeros_like(c_val_custom)
    # custom_kernel_input = (a_val, b_val, a_s_val, b_s_val, c_val_custom)
    #
    # print(f"Running custom kernel for M={M}, N={N}, K={K}...")
    # output_tensor_custom = custom_kernel(custom_kernel_input)
    # print("Custom kernel execution complete.")
    #
    # # For verification, run the reference kernel
    # # c_val_ref = torch.zeros_like(c_val_custom)
    # # ref_kernel_input = (a_val, b_val, a_s_val, b_s_val, c_val_ref)
    # # output_tensor_ref = ref_kernel(ref_kernel_input)
    # # print("Reference kernel execution complete.")
    #
    # # Compare results (adjust tolerances as needed for FP8/BF16)
    # # if torch.allclose(output_tensor_custom, output_tensor_ref, atol=1e-1, rtol=1e-1):
    # #     print("SUCCESS: Custom kernel output matches reference output.")
    # # else:
    # #     print("FAILURE: Custom kernel output DOES NOT match reference output.")
    # #     print("Max difference:", torch.max(torch.abs(output_tensor_custom - output_tensor_ref)))

    pass
python submission.py
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized/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 -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/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 -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 -fno-gpu-rdc -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized/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_hip_optimized.so

python eval.py benchmark /tmp/tmp_k83ro6n
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized/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.
