import os
os.environ['PYTORCH_ROCM_ARCH']='gfx942'  # Maybe 'gfx942:xnack-' later?
os.environ["CXX"] = "clang++"

# This script provides a template for using load_inline to run a HIP kernel for MI300
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 a_scale, torch::Tensor b_scale, torch::Tensor c);
"""

HIP_SRC = r"""
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>
#include <rocwmma/rocwmma.hpp>
#include <cstring> // For memcpy

// Helper for checking HIP errors (optional, good for standalone)
#define HIP_CHECK(cmd)                                                         \\
    do {                                                                       \\
        hipError_t e = cmd;                                                    \\
        if (e != hipSuccess) {                                                 \\
            printf("HIP error %s:%d '%s' (%d)\\n", __FILE__, __LINE__,          \\
                   hipGetErrorString(e), e);                                   \\
            abort();                                                           \\
        }                                                                      \\
    } while (0)

// Typedefs for HIP datatypes
typedef __hip_fp8_e4m3_fnuz fp8_t;
typedef __hip_bfloat16 bf16_t;

// Constants from problem description
constexpr int SCALE_BLOCK_DIM_K_CONST = 128;
constexpr int SCALE_BLOCK_DIM_N_CONST = 128; // For b_scale indexing

// rocWMMA Parameters
constexpr uint32_t MFMA_M_TILE_M = 32u; // MFMA output tile M dimension
constexpr uint32_t MFMA_N_TILE_N = 32u; // MFMA output tile N dimension
constexpr uint32_t MFMA_K_TILE_K = 16u; // K-slice per MFMA operation

// Output C tile computed by one thread block
constexpr uint32_t BLOCK_M = 64u;
constexpr uint32_t BLOCK_N = 64u;
// K-dimension processed in one outer loop (matches scaling block)
constexpr uint32_t BLOCK_K_OUTER = SCALE_BLOCK_DIM_K_CONST; // 128u

// Number of MFMA 32x32 tiles to fill the C_BLOCK (64x64) output tile
constexpr uint32_t BLOCKS_PER_TILE_M = BLOCK_M / MFMA_M_TILE_M; // 2
constexpr uint32_t BLOCKS_PER_TILE_N = BLOCK_N / MFMA_N_TILE_N; // 2

// Thread block configuration (one wave = 64 threads for gfx942)
constexpr uint32_t TBLOCK_X_DIM = rocwmma::Constants::AMDGCN_WAVE_SIZE_64;
constexpr uint32_t TBLOCK_Y_DIM = 1u; 
constexpr uint32_t TOTAL_THREADS_PER_BLOCK = TBLOCK_X_DIM * TBLOCK_Y_DIM;

constexpr int NUM_FLOAT_PER_THREAD_ACC = 16; 

using namespace rocwmma;

using GlobalLayoutA = col_major; 
using GlobalLayoutB = col_major; 
using GlobalLayoutC = row_major; 

using MatrixA_MfmaFrag = fragment<matrix_a, MFMA_M_TILE_M, MFMA_N_TILE_N, MFMA_K_TILE_K, fp8_t, col_major>;
using MatrixB_MfmaFrag = fragment<matrix_b, MFMA_M_TILE_M, MFMA_N_TILE_N, MFMA_K_TILE_K, fp8_t, row_major>;
using Accumulator_MfmaFrag = fragment<accumulator, MFMA_M_TILE_M, MFMA_N_TILE_N, MFMA_K_TILE_K, float, GlobalLayoutC>;


// --- Modified Helper Device Function: Load a tile from Global to LDS ---
// This version simulates launching block-level asynchronous copies (e.g., TMA).
// One thread (wave_lane_id == 0) issues these conceptual copies.
// `memcpy` is used as a placeholder for actual MI300 TMA intrinsics/calls.
// In a real TMA scenario, these `memcpy` calls would be non-blocking dispatches.
__device__ void load_gmem_tile_to_lds_async_placeholder(
    fp8_t* lds_a_target, const fp8_t* __restrict__ global_a_ptr,
    fp8_t* lds_b_target, const fp8_t* __restrict__ global_b_ptr,
    int current_k_block_base,
    int M_param, int N_out_param, int K_param, // Global dimensions
    uint32_t block_c_base_m_dim, uint32_t block_c_base_n_dim, // Block's base in C
    uint32_t wave_lane_id)
{
    // Only one thread in the wave issues the (conceptual) async block copies
    if (wave_lane_id == 0) {
        // Load A: Global A (M x K, col-major) to LDS A (K_OUTER x BLOCK_M)
        // LDS A stores K_OUTER rows, each of BLOCK_M elements.
        // This requires BLOCK_K_OUTER separate 1D copies.
        for (uint32_t k_local_idx = 0; k_local_idx < BLOCK_K_OUTER; ++k_local_idx) {
            uint32_t gmem_k_absolute = current_k_block_base + k_local_idx;

            if (gmem_k_absolute < K_param) { // Check if this k-slice is within bounds
                // Check if the M-dimension for this block is within global M bounds
                // block_c_base_m_dim is the start M-coord for this block.
                // We copy BLOCK_M elements. So, block_c_base_m_dim + 0 to block_c_base_m_dim + BLOCK_M-1
                if (block_c_base_m_dim < M_param) {
                     const fp8_t* gmem_src_a = &global_a_ptr[gmem_k_absolute * M_param + block_c_base_m_dim];
                     fp8_t* lds_dst_a = &lds_a_target[k_local_idx * BLOCK_M];
                     
                     // Placeholder for an asynchronous hardware copy primitive (e.g., TMA call)
                     // This memcpy is synchronous, but in a real TMA scenario, it would be a non-blocking dispatch.
                     // The size is BLOCK_M bytes (since fp8_t is 1 byte).
                     // Ensure we don't read past M_param if block_c_base_m_dim + BLOCK_M > M_param.
                     // However, problem constraints state M is divisible by BLOCK_M, so this full copy is safe.
                     memcpy(lds_dst_a, gmem_src_a, BLOCK_M * sizeof(fp8_t));
                }

                // Load B: Global B (N x K, col-major) to LDS B (K_OUTER x BLOCK_N)
                // Similar logic for B
                if (block_c_base_n_dim < N_out_param) {
                    const fp8_t* gmem_src_b = &global_b_ptr[gmem_k_absolute * N_out_param + block_c_base_n_dim];
                    fp8_t* lds_dst_b = &lds_b_target[k_local_idx * BLOCK_N];

                    // Placeholder for an asynchronous hardware copy primitive
                    // Ensure we don't read past N_out_param. N_out is divisible by BLOCK_N.
                    memcpy(lds_dst_b, gmem_src_b, BLOCK_N * sizeof(fp8_t));
                }
            }
        }
    }
    // This function returns, and the main kernel loop will use synchronize_workgroup()
    // to wait for these conceptual asynchronous copies to complete and be visible in LDS.
}


// --- Helper Device Function: Compute MMAs on data in LDS (Unchanged) ---
__device__ void compute_lds_tile_mma(
    const fp8_t* lds_a_current, 
    const fp8_t* lds_b_current,
    Accumulator_MfmaFrag current_k_block_acc_frags[BLOCKS_PER_TILE_M][BLOCKS_PER_TILE_N])
{
    for (uint32_t r_frag_idx = 0; r_frag_idx < BLOCKS_PER_TILE_M; ++r_frag_idx) {
        for (uint32_t c_frag_idx = 0; c_frag_idx < BLOCKS_PER_TILE_N; ++c_frag_idx) {
            fill_fragment(current_k_block_acc_frags[r_frag_idx][c_frag_idx], 0.0f);
        }
    }
    
    for (int k_inner_offset = 0; k_inner_offset < BLOCK_K_OUTER; k_inner_offset += MFMA_K_TILE_K) {
        MatrixA_MfmaFrag a_mfma_sub_frags[BLOCKS_PER_TILE_M]; 
        MatrixB_MfmaFrag b_mfma_sub_frags[BLOCKS_PER_TILE_N]; 

        for (uint32_t m_tile_idx = 0; m_tile_idx < BLOCKS_PER_TILE_M; ++m_tile_idx) {
            const fp8_t* lds_a_sub_tile_ptr = &lds_a_current[(k_inner_offset * BLOCK_M) + (m_tile_idx * MFMA_M_TILE_M)];
            load_matrix_sync(a_mfma_sub_frags[m_tile_idx], lds_a_sub_tile_ptr, BLOCK_M);
        }

        for (uint32_t n_tile_idx = 0; n_tile_idx < BLOCKS_PER_TILE_N; ++n_tile_idx) {
            const fp8_t* lds_b_sub_tile_ptr = &lds_b_current[(k_inner_offset * BLOCK_N) + (n_tile_idx * MFMA_N_TILE_N)];
            load_matrix_sync(b_mfma_sub_frags[n_tile_idx], lds_b_sub_tile_ptr, BLOCK_N);
        }
        
        for (uint32_t m_frag_idx = 0; m_frag_idx < BLOCKS_PER_TILE_M; ++m_frag_idx) {
            for (uint32_t n_frag_idx = 0; n_frag_idx < BLOCKS_PER_TILE_N; ++n_frag_idx) {
                mma_sync(current_k_block_acc_frags[m_frag_idx][n_frag_idx],
                         a_mfma_sub_frags[m_frag_idx],    
                         b_mfma_sub_frags[n_frag_idx],    
                         current_k_block_acc_frags[m_frag_idx][n_frag_idx]);
            }
        }
    } 
}


__global__ void __launch_bounds__(TOTAL_THREADS_PER_BLOCK)
    custom_kernel_rocwmma_pipelined(
    const fp8_t* __restrict__ global_a_ptr,
    const fp8_t* __restrict__ global_b_ptr,
    const float* __restrict__ global_a_scale_ptr,
    const float* __restrict__ global_b_scale_ptr,
    bf16_t* __restrict__ global_c_ptr,
    int M_param, int N_out_param, int K_param)
{
    __shared__ fp8_t lds_a_ping[BLOCK_M * BLOCK_K_OUTER]; 
    __shared__ fp8_t lds_a_pong[BLOCK_M * BLOCK_K_OUTER]; 
    __shared__ fp8_t lds_b_ping[BLOCK_N * BLOCK_K_OUTER]; 
    __shared__ fp8_t lds_b_pong[BLOCK_N * BLOCK_K_OUTER]; 
    
    fp8_t* lds_a_double_buffers[2] = {lds_a_ping, lds_a_pong};
    fp8_t* lds_b_double_buffers[2] = {lds_b_ping, lds_b_pong};
    int buffer_selector = 0; 

    uint32_t block_c_start_row = blockIdx.x * BLOCK_M;
    uint32_t block_c_start_col = blockIdx.y * BLOCK_N;
    
    uint32_t wave_lane_id = threadIdx.x; 

    Accumulator_MfmaFrag final_accum_frags[BLOCKS_PER_TILE_M][BLOCKS_PER_TILE_N];
    for (uint32_t i = 0; i < BLOCKS_PER_TILE_M; ++i) {
        for (uint32_t j = 0; j < BLOCKS_PER_TILE_N; ++j) {
            fill_fragment(final_accum_frags[i][j], 0.0f);
        }
    }
    
    if (K_param == 0) { return; }

    // Pre-load first tile using the new async placeholder loader
    load_gmem_tile_to_lds_async_placeholder(
        lds_a_double_buffers[buffer_selector], global_a_ptr,
        lds_b_double_buffers[buffer_selector], global_b_ptr,
        0, /* k_current_block_base */
        M_param, N_out_param, K_param,
        block_c_start_row, block_c_start_col,
        wave_lane_id);

    int num_k_outer_iterations = K_param / BLOCK_K_OUTER;

    for (int k_iter_idx = 0; k_iter_idx < num_k_outer_iterations; ++k_iter_idx) {
        int current_k_block_base = k_iter_idx * BLOCK_K_OUTER;

        fp8_t* lds_a_for_compute = lds_a_double_buffers[buffer_selector];
        fp8_t* lds_b_for_compute = lds_b_double_buffers[buffer_selector];

        int next_k_block_base = current_k_block_base + BLOCK_K_OUTER;
        if (next_k_block_base < K_param) { 
            load_gmem_tile_to_lds_async_placeholder(
                lds_a_double_buffers[1 - buffer_selector], global_a_ptr,
                lds_b_double_buffers[1 - buffer_selector], global_b_ptr,
                next_k_block_base,
                M_param, N_out_param, K_param,
                block_c_start_row, block_c_start_col,
                wave_lane_id);
        }

        // SYNC-A: Wait for the current compute buffer (loaded in previous iteration or pre-load)
        // to be fully loaded and visible in LDS.
        synchronize_workgroup();

        Accumulator_MfmaFrag current_k_block_accum_frags[BLOCKS_PER_TILE_M][BLOCKS_PER_TILE_N];
        compute_lds_tile_mma(lds_a_for_compute, lds_b_for_compute, current_k_block_accum_frags);
        
        uint32_t k_scale_block_idx = current_k_block_base / SCALE_BLOCK_DIM_K_CONST;
        uint32_t b_scale_num_n_blocks = (N_out_param + SCALE_BLOCK_DIM_N_CONST - 1) / SCALE_BLOCK_DIM_N_CONST;

        for (uint32_t m_frag_idx = 0; m_frag_idx < BLOCKS_PER_TILE_M; ++m_frag_idx) {
            for (uint32_t n_frag_idx = 0; n_frag_idx < BLOCKS_PER_TILE_N; ++n_frag_idx) {
                auto& source_accum_frag = current_k_block_accum_frags[m_frag_idx][n_frag_idx];
                auto& dest_final_accum_frag = final_accum_frags[m_frag_idx][n_frag_idx];

                for (int p_idx = 0; p_idx < NUM_FLOAT_PER_THREAD_ACC; ++p_idx) {
                    uint32_t p_div_4 = p_idx / 4;
                    uint32_t p_mod_4 = p_idx % 4;
                    uint32_t lane_div_32 = wave_lane_id / 32; 
                    uint32_t lane_mod_32 = wave_lane_id % 32; 
                    
                    uint32_t m_coord_in_mfma_tile = (8 * p_div_4) + (4 * lane_div_32) + p_mod_4;
                    uint32_t n_coord_in_mfma_tile = lane_mod_32;

                    uint32_t global_m_coord = block_c_start_row + m_frag_idx * MFMA_M_TILE_M + m_coord_in_mfma_tile;
                    uint32_t global_n_coord = block_c_start_col + n_frag_idx * MFMA_N_TILE_N + n_coord_in_mfma_tile;

                    float scale_a_val;
                    if (global_m_coord < M_param) {
                         scale_a_val = global_a_scale_ptr[global_m_coord + k_scale_block_idx * M_param];
                    } else { continue; }

                    float scale_b_val;
                    uint32_t n_b_scale_block_idx = global_n_coord / SCALE_BLOCK_DIM_N_CONST;
                    if (global_n_coord < N_out_param && n_b_scale_block_idx < b_scale_num_n_blocks) {
                        scale_b_val = global_b_scale_ptr[n_b_scale_block_idx + k_scale_block_idx * b_scale_num_n_blocks];
                    } else { continue; }
                    
                    float combined_scale = scale_a_val * scale_b_val;
                    dest_final_accum_frag.x[p_idx] += source_accum_frag.x[p_idx] * combined_scale;
                }
            }
        }
        
        // SYNC-B: Wait for the load of the *next* tile (if one was dispatched) to complete
        // This ensures the data is in LDS before swapping buffers for the next iteration.
        if (next_k_block_base < K_param) { 
             synchronize_workgroup();
        }

        buffer_selector = 1 - buffer_selector;
    } 
    
    for (uint32_t m_frag_idx = 0; m_frag_idx < BLOCKS_PER_TILE_M; ++m_frag_idx) {
        for (uint32_t n_frag_idx = 0; n_frag_idx < BLOCKS_PER_TILE_N; ++n_frag_idx) {
            auto& result_frag_to_store = final_accum_frags[m_frag_idx][n_frag_idx];
            for (int p_idx = 0; p_idx < NUM_FLOAT_PER_THREAD_ACC; ++p_idx) {
                uint32_t p_div_4 = p_idx / 4;
                uint32_t p_mod_4 = p_idx % 4;
                uint32_t lane_div_32 = wave_lane_id / 32;
                uint32_t lane_mod_32 = wave_lane_id % 32;
                uint32_t m_coord_in_mfma_tile = (8 * p_div_4) + (4 * lane_div_32) + p_mod_4;
                uint32_t n_coord_in_mfma_tile = lane_mod_32;

                uint32_t global_m_coord = block_c_start_row + m_frag_idx * MFMA_M_TILE_M + m_coord_in_mfma_tile;
                uint32_t global_n_coord = block_c_start_col + n_frag_idx * MFMA_N_TILE_N + n_coord_in_mfma_tile;

                if (global_m_coord < M_param && global_n_coord < N_out_param) {
                    global_c_ptr[global_m_coord * N_out_param + global_n_coord] = static_cast<bf16_t>(result_frag_to_store.x[p_idx]);
                }
            }
        }
    }
}

// Entry-point from PyTorch (this function signature is fixed, and unchangeable)
#include <torch/types.h> 
// Required for TORCH_CHECK and other PyTorch types if not already included by hip_runtime.h context
#if __HIP_PLATFORM_HCC__ // Ensure this is HIP context for PyTorch headers
#include <c10/hip/HIPStream.h> 
#endif


void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor a_scale, torch::Tensor b_scale, torch::Tensor c) {
  int M = a.size(0);
  int K = a.size(1); 
  int N_out = b.size(0); 


  
  
  uint32_t b_scale_expected_dim0 = (N_out + SCALE_BLOCK_DIM_N_CONST -1) / SCALE_BLOCK_DIM_N_CONST;



  dim3 grid_dim(M / BLOCK_M, N_out / BLOCK_N);
  dim3 block_dim(TBLOCK_X_DIM, TBLOCK_Y_DIM); 
  
  hipStream_t stream = nullptr; 
  #if __HIP_PLATFORM_HCC__ && TORCH_VERSION_MAJOR >= 1 // Basic check for PyTorch context
    // It's good practice to use the stream from PyTorch if in that environment
    // stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA(); // Might require c10::hip::impl...
    // For simplicity, using nullptr (default stream) if not straightforwardly available.
    // If this code is compiled standalone or specific PyTorch linkage is not set up,
    // relying on default stream is safer.
    // The problem implies this is part of a PyTorch extension, so stream integration is relevant.
    // However, the current template uses nullptr, so sticking to that for now.
  #endif


  custom_kernel_rocwmma_pipelined<<<grid_dim, block_dim, 0, stream>>>(
    reinterpret_cast<const fp8_t*>(a.data_ptr()), 
    reinterpret_cast<const fp8_t*>(b.data_ptr()), 
    a_scale.data_ptr<float>(), 
    b_scale.data_ptr<float>(), 
    reinterpret_cast<bf16_t*>(c.data_ptr()), 
    M, N_out, K);
    
  hipError_t last_error = hipGetLastError();
  if (last_error != hipSuccess) { 
      std::string err_msg = "HIP kernel launch failed: ";
      err_msg += hipGetErrorString(last_error);
  }
}
"""


module = load_inline(
    name='mdda',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[HIP_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/mdda/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/mdda/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/mdda/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/mdda/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=mdda -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/mdda/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=mdda -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/mdda/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 mdda.so

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

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


Total number of replaced kernel launches: 1
[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 mdda.so
