import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

CPP_WRAPPER = """
void fp8_mm_hip_optimized_v3_wrapper(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c);
"""

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 K dimension
constexpr int SCALE_BLOCK_K = 128;

// --- Kernel Configuration ---
// Output tile computed by one thread block
constexpr int BLOCK_TILE_M = 16;
constexpr int BLOCK_TILE_N = 16; // Each thread block computes a 16xN_elements_per_thread*THREADS_N_LOGICAL C tile

// Threads in a block (M dimension)
constexpr int THREADS_M = 16;
// Logical threads in N dimension (blockDim.y)
constexpr int THREADS_N_LOGICAL = 8;
// Each logical N-thread handles this many elements in N dimension
constexpr int UNROLL_N_FACTOR = BLOCK_TILE_N / THREADS_N_LOGICAL; // = 16 / 8 = 2

// Shared memory K dimension tiling
constexpr int K_SHARED_STEP = 32; // Increased K tile in shared memory

__global__ void custom_hip_kernel_optimized_v3(
    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
    int tx = threadIdx.x; // 0..THREADS_M-1
    int ty_logical = threadIdx.y; // 0..THREADS_N_LOGICAL-1

    // Global row for this thread's M-dimension responsibility
    int global_c_row = blockIdx.x * BLOCK_TILE_M + tx;

    // Accumulators for UNROLL_N_FACTOR elements this thread is responsible for
    float total_acc_fp32[UNROLL_N_FACTOR];
    for (int i = 0; i < UNROLL_N_FACTOR; ++i) {
        total_acc_fp32[i] = 0.0f;
    }

    // Shared memory
    // sA: For matrix A, stores a tile of [BLOCK_TILE_M x K_SHARED_STEP]
    // sB: For matrix B, stores a tile of [K_SHARED_STEP x BLOCK_TILE_N] (K rows, N cols for B.T)
    __shared__ __hip_fp8_e4m3_fnuz sA[BLOCK_TILE_M][K_SHARED_STEP];
    __shared__ __hip_fp8_e4m3_fnuz sB[K_SHARED_STEP][BLOCK_TILE_N];

    // Pre-calculate number of N-tiles in b_scale for indexing
    int b_scale_num_n_tiles = (N + SCALE_BLOCK_K - 1) / SCALE_BLOCK_K;

    // Outer K loop: Iterates over K dimension in chunks of SCALE_BLOCK_K (for scaling)
    for (int k_block_base = 0; k_block_base < K; k_block_base += SCALE_BLOCK_K) {
        // Load scales for this K-block
        float scale_a_val = 1.0f;
        float scale_b_vals[UNROLL_N_FACTOR]; // Scales for B for each N-element this thread handles

        if (global_c_row < M) { // Boundary check for A's scale
            scale_a_val = a_scale_ptr[global_c_row + (k_block_base / SCALE_BLOCK_K) * M];
        }

        for (int n_idx = 0; n_idx < UNROLL_N_FACTOR; ++n_idx) {
            int current_c_col = blockIdx.y * BLOCK_TILE_N + ty_logical * UNROLL_N_FACTOR + n_idx;
            if (global_c_row < M && current_c_col < N) { // Boundary check for B's scale
                 scale_b_vals[n_idx] = b_scale_ptr[(current_c_col / SCALE_BLOCK_K) + (k_block_base / SCALE_BLOCK_K) * b_scale_num_n_tiles];
            } else {
                 scale_b_vals[n_idx] = 0.0f; // Elements out of bounds should not contribute
            }
        }

        float current_k_block_sum_fp32[UNROLL_N_FACTOR];
        for (int i = 0; i < UNROLL_N_FACTOR; ++i) {
            current_k_block_sum_fp32[i] = 0.0f;
        }

        // Inner K loop: Iterates through the current SCALE_BLOCK_K segment using 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 sA tile (BLOCK_TILE_M x K_SHARED_STEP)
            // Each thread (tx, ty_logical) loads K_SHARED_STEP / THREADS_N_LOGICAL elements into sA.
            // tx defines the row in sA. ty_logical helps distribute loading across the K_SHARED_STEP columns.
            for(int k_load_iter = 0; k_load_iter < K_SHARED_STEP / THREADS_N_LOGICAL; ++k_load_iter) {
                int k_idx_in_sA = ty_logical * (K_SHARED_STEP / THREADS_N_LOGICAL) + k_load_iter;
                if (global_c_row < M && (current_k_tile_base + k_idx_in_sA) < K) {
                    sA[tx][k_idx_in_sA] = a_ptr[global_c_row + (current_k_tile_base + k_idx_in_sA) * M];
                } else {
                    sA[tx][k_idx_in_sA] = __hip_fp8_e4m3_fnuz(0.0f);
                }
            }

            // Load sB tile (K_SHARED_STEP x BLOCK_TILE_N)
            // Each thread (tx, ty_logical) loads K_SHARED_STEP / THREADS_M elements for its UNROLL_N_FACTOR columns.
            // tx helps distribute loading K_SHARED_STEP rows of sB.
            // ty_logical determines which group of N_UNROLL_FACTOR columns to load.
            for (int n_load_idx = 0; n_load_idx < UNROLL_N_FACTOR; ++n_load_idx) {
                int sB_col_idx = ty_logical * UNROLL_N_FACTOR + n_load_idx; // Column in sB (0 to BLOCK_TILE_N-1)
                int global_c_col_for_b = blockIdx.y * BLOCK_TILE_N + sB_col_idx;

                for (int k_load_iter = 0; k_load_iter < K_SHARED_STEP / THREADS_M; ++k_load_iter) {
                    int k_idx_in_sB = tx * (K_SHARED_STEP / THREADS_M) + k_load_iter; // Row in sB
                    if (global_c_col_for_b < N && (current_k_tile_base + k_idx_in_sB) < K) {
                        sB[k_idx_in_sB][sB_col_idx] = b_ptr[global_c_col_for_b + (current_k_tile_base + k_idx_in_sB) * N];
                    } else {
                        sB[k_idx_in_sB][sB_col_idx] = __hip_fp8_e4m3_fnuz(0.0f);
                    }
                }
            }
            __syncthreads(); // Ensure sA and sB are fully loaded

            // Accumulate product from shared memory
            // Unroll K_SHARED_STEP loop by a factor (e.g., 4) or fully if beneficial
            #pragma unroll 4 // Partial unroll for the K_SHARED_STEP=32 loop
            for (int k_s = 0; k_s < K_SHARED_STEP; ++k_s) {
                float val_a = static_cast<float>(sA[tx][k_s]);
                for (int n_idx = 0; n_idx < UNROLL_N_FACTOR; ++n_idx) {
                    int sB_col_idx = ty_logical * UNROLL_N_FACTOR + n_idx;
                    float val_b = static_cast<float>(sB[k_s][sB_col_idx]);
                    current_k_block_sum_fp32[n_idx] += val_a * val_b;
                }
            }
            __syncthreads(); // Ensure all threads finish using shared mem before next load cycle
        } // end k_shared_offset loop (inner K loop)

        // Accumulate scaled results for the K-block
        for (int n_idx = 0; n_idx < UNROLL_N_FACTOR; ++n_idx) {
             if (global_c_row < M) { // Ensure scale_a_val is valid before using
                 int current_c_col_check = blockIdx.y * BLOCK_TILE_N + ty_logical * UNROLL_N_FACTOR + n_idx;
                 if(current_c_col_check < N) { // Check if this N-element is within bounds
                    total_acc_fp32[n_idx] += current_k_block_sum_fp32[n_idx] * scale_a_val * scale_b_vals[n_idx];
                 }
             }
        }
    } // end k_block_base loop (outer K loop)

    // Write results to global memory C
    for (int n_idx = 0; n_idx < UNROLL_N_FACTOR; ++n_idx) {
        int global_c_col = blockIdx.y * BLOCK_TILE_N + ty_logical * UNROLL_N_FACTOR + n_idx;
        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[n_idx]);
        }
    }
}

// C++ Wrapper function
void fp8_mm_hip_optimized_v3_wrapper(torch::Tensor a_tensor, torch::Tensor b_tensor,
                                   torch::Tensor a_scale_tensor, torch::Tensor b_scale_tensor,
                                   torch::Tensor c_tensor) {
    int M = a_tensor.size(0);
    int K = a_tensor.size(1);
    int N_b_rows = b_tensor.size(0); // N from B (rows of B, which is columns of C)
    // int K_b = b_tensor.size(1); // Should match K from A
    // C is M x N_b_rows

    // Launch configuration: THREADS_M x THREADS_N_LOGICAL threads per block
    dim3 threadsPerBlock(THREADS_M, THREADS_N_LOGICAL); // e.g., 16x8 = 128 threads
    // Grid dimensions to cover the M x N output matrix, considering BLOCK_TILE sizes
    dim3 numBlocks((M + BLOCK_TILE_M - 1) / BLOCK_TILE_M,
                   (N_b_rows + BLOCK_TILE_N - 1) / BLOCK_TILE_N);

    hipLaunchKernelGGL(custom_hip_kernel_optimized_v3, numBlocks, threadsPerBlock, 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, N_b_rows, K
    );
    // hipError_t err = hipGetLastError(); // Optional: error checking
    // if (err != hipSuccess) {
    //    printf("HIP Kernel Launch Error in v3 (%d): %s\\n", err, hipGetErrorString(err));
    // }
}
"""

import os
os.environ["CXX"] = "hipcc" # Use hipcc as it might be more direct for HIP

ROCM_VERSION_PATH = os.environ.get("ROCM_VERSION_PATH", "/opt/rocm-6.3.1") # From your previous log
ROCM_INCLUDE_PATH = f"{ROCM_VERSION_PATH}/include"


module = load_inline(
    name='fp8_mm_hip_optimized_v3', # New unique name for the module
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[HIP_SRC],
    functions=['fp8_mm_hip_optimized_v3_wrapper'], # Exposed C++ function
    verbose=True,
    extra_cuda_cflags=[
        "--offload-arch=gfx942",
        "-std=c++17", # Using C++17
        f"-I{ROCM_INCLUDE_PATH}", # Ensure ROCm headers can be found
        "-O3", # Optimization level
        "-Wall", # Enable all warnings (good for development)
        # "-save-temps" # Uncomment to save intermediate compilation files for inspection
    ]
)

def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    # It's good practice to ensure input tensors are contiguous,
    # especially if the kernel relies on specific memory layouts for performance.
    # However, the kernel loads element by element based on calculated global indices,
    # so contiguity might not be strictly necessary for correctness here,
    # but can impact performance of data_ptr() if views are involved.
    # For now, let's assume inputs are already appropriately laid out or contiguous.
    module.fp8_mm_hip_optimized_v3_wrapper(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_v3/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v3/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v3/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v3/hip.hip [ok]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 0
[1/3] /opt/rocm-6.3.1/bin/hipcc  -DWITH_HIP -DTORCH_EXTENSION_NAME=fp8_mm_hip_optimized_v3 -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++17 -I/opt/rocm-6.3.1/include -O3 -Wall -fno-gpu-rdc -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_hip_optimized_v3/hip.hip -o hip.cuda.o 
[2/3] hipcc -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=fp8_mm_hip_optimized_v3 -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_v3/main.cpp -o main.o -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC
[3/3] hipcc 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_v3.so

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

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