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>

// 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
// EXPLORATION: Changed BLOCK_M and BLOCK_N
constexpr uint32_t BLOCK_M = 128u; // Increased from 64u
constexpr uint32_t BLOCK_N = 64u;  // Kept at 64u (could also be 128u, but 128x64 uses 48KB LDS)
// K-dimension processed in one outer loop (matches scaling block)
constexpr uint32_t BLOCK_K_OUTER = SCALE_BLOCK_DIM_K_CONST; // 128u, kept same

// Number of MFMA 32x32 tiles to fill the C_BLOCK output tile
constexpr uint32_t BLOCKS_PER_TILE_M = BLOCK_M / MFMA_M_TILE_M; // e.g. 128/32 = 4
constexpr uint32_t BLOCKS_PER_TILE_N = BLOCK_N / MFMA_N_TILE_N; // e.g. 64/32 = 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; // Single wave per block
constexpr uint32_t TOTAL_THREADS_PER_BLOCK = TBLOCK_X_DIM * TBLOCK_Y_DIM;

constexpr int NUM_FLOAT_PER_THREAD_ACC = 16; // For 32x32x16 fp8 MFMA, float accumulator

// rocWMMA Fragment Types
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>;


__device__ void load_gmem_tile_to_lds_vectorized(
    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, 
    uint32_t block_c_base_m_dim, uint32_t block_c_base_n_dim, 
    uint32_t wave_lane_id)
{
    constexpr uint32_t VECTOR_SIZE_FP8 = 4; 

    constexpr uint32_t VECTORS_IN_M_DIM_A = BLOCK_M / VECTOR_SIZE_FP8;      
    constexpr uint32_t THREAD_GROUP_SIZE_A = VECTORS_IN_M_DIM_A;           
    constexpr uint32_t NUM_THREAD_GROUPS_A = TOTAL_THREADS_PER_BLOCK / THREAD_GROUP_SIZE_A; 

    uint32_t group_id_a = wave_lane_id / THREAD_GROUP_SIZE_A;      
    uint32_t lane_in_group_a = wave_lane_id % THREAD_GROUP_SIZE_A; 

    for (uint32_t k_iter_idx_a = 0; k_iter_idx_a < BLOCK_K_OUTER / NUM_THREAD_GROUPS_A; ++k_iter_idx_a) {
        uint32_t k_idx_in_tile = group_id_a + k_iter_idx_a * NUM_THREAD_GROUPS_A;
        uint32_t vec_m_idx = lane_in_group_a; 
        
        uint32_t gmem_a_k = current_k_block_base + k_idx_in_tile;
        uint32_t m_start_in_block = vec_m_idx * VECTOR_SIZE_FP8; 
        uint32_t gmem_a_m_start = block_c_base_m_dim + m_start_in_block;

        if (gmem_a_k < K_param && gmem_a_m_start < M_param) { 
            const int* gmem_src_ptr = reinterpret_cast<const int*>(
                &global_a_ptr[gmem_a_k * M_param + gmem_a_m_start]);
            
            int* lds_dst_ptr = reinterpret_cast<int*>(
                &lds_a_target[k_idx_in_tile * BLOCK_M + m_start_in_block]);
            
            *lds_dst_ptr = *gmem_src_ptr;
        }
    }

    constexpr uint32_t VECTORS_IN_N_DIM_B = BLOCK_N / VECTOR_SIZE_FP8;      
    constexpr uint32_t THREAD_GROUP_SIZE_B = VECTORS_IN_N_DIM_B;           
    constexpr uint32_t NUM_THREAD_GROUPS_B = TOTAL_THREADS_PER_BLOCK / THREAD_GROUP_SIZE_B; 

    uint32_t group_id_b = wave_lane_id / THREAD_GROUP_SIZE_B; 
    uint32_t lane_in_group_b = wave_lane_id % THREAD_GROUP_SIZE_B;

    for (uint32_t k_iter_idx_b = 0; k_iter_idx_b < BLOCK_K_OUTER / NUM_THREAD_GROUPS_B; ++k_iter_idx_b) {
        uint32_t k_idx_in_tile = group_id_b + k_iter_idx_b * NUM_THREAD_GROUPS_B;
        uint32_t vec_n_idx = lane_in_group_b; 
        
        uint32_t gmem_b_k = current_k_block_base + k_idx_in_tile;
        uint32_t n_start_in_block = vec_n_idx * VECTOR_SIZE_FP8; 
        uint32_t gmem_b_n_start = block_c_base_n_dim + n_start_in_block;

        if (gmem_b_k < K_param && gmem_b_n_start < N_out_param) { 
            const int* gmem_src_ptr = reinterpret_cast<const int*>(
                &global_b_ptr[gmem_b_k * N_out_param + gmem_b_n_start]);
            
            int* lds_dst_ptr = reinterpret_cast<int*>(
                &lds_b_target[k_idx_in_tile * BLOCK_N + n_start_in_block]);
            
            *lds_dst_ptr = *gmem_src_ptr;
        }
    }
}


__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 memory for double buffering A and B tiles ---
    // LDS A: BLOCK_M rows, BLOCK_K_OUTER columns (if viewing A as M major)
    // Stored as K_OUTER rows of BLOCK_M elements (col-major from global A's perspective)
    __shared__ fp8_t lds_a_ping[BLOCK_M * BLOCK_K_OUTER]; 
    __shared__ fp8_t lds_a_pong[BLOCK_M * BLOCK_K_OUTER]; 
    // LDS B: BLOCK_N rows, BLOCK_K_OUTER columns (if viewing B as N major)
    // Stored as K_OUTER rows of BLOCK_N elements (col-major from global B's perspective)
    __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_global = blockIdx.x * BLOCK_M;
    uint32_t block_c_start_col_global = blockIdx.y * BLOCK_N;
    
    uint32_t wave_lane_id = threadIdx.x; 

    if (K_param == 0) { return; }

    load_gmem_tile_to_lds_vectorized(
        lds_a_double_buffers[buffer_selector], global_a_ptr,
        lds_b_double_buffers[buffer_selector], global_b_ptr,
        0, 
        M_param, N_out_param, K_param,
        block_c_start_row_global, block_c_start_col_global,
        wave_lane_id);

    int num_k_outer_iterations = K_param / BLOCK_K_OUTER;
    const uint32_t b_scale_num_n_blocks = (N_out_param + SCALE_BLOCK_DIM_N_CONST - 1) / SCALE_BLOCK_DIM_N_CONST;

    // REFACTOR: Loop over MFMA tiles within the larger C_BLOCK
    for (uint32_t m_frag_block_idx = 0; m_frag_block_idx < BLOCKS_PER_TILE_M; ++m_frag_block_idx) {
        for (uint32_t n_frag_block_idx = 0; n_frag_block_idx < BLOCKS_PER_TILE_N; ++n_frag_block_idx) {
            
            // Base M and N coordinates for the current 32x32 MFMA tile within the C_BLOCK
            uint32_t mfma_tile_offset_m_in_block = m_frag_block_idx * MFMA_M_TILE_M;
            uint32_t mfma_tile_offset_n_in_block = n_frag_block_idx * MFMA_N_TILE_N;

            Accumulator_MfmaFrag final_accum_frag_for_mfma_tile; // For one 32x32 tile
            fill_fragment(final_accum_frag_for_mfma_tile, 0.0f);

            // K-dimension loop (pipelined)
            // Reset buffer_selector for each new MFMA tile if LDS loads depend on m_frag_block/n_frag_block.
            // Current LDS load loads full BLOCK_M x K / BLOCK_N x K, so buffer_selector can persist.
            // If LDS load was per MFMA tile, it would need reset & first load here.
            // But since LDS load is for the whole BLOCK_MxBLOCK_N, and k-loop is inside m/n_frag_block_idx loops,
            // we'd reload the same LDS data BLOCKS_PER_TILE_M * BLOCKS_PER_TILE_N times.
            // This is inefficient. K-loop should be outermost after block setup.

            // CORRECTED LOOP ORDER: K-loop should be outside m_frag_block_idx / n_frag_block_idx loops
            // The provided code already had K-loop outside, and accumulator arrays.
            // The refactor means we need to move K-loop logic inside m/n_frag_block_idx.
            // The pipelining with buffer_selector must be managed carefully.
            // Let's re-verify the original structure's K-loop and pipelining.
            // Original: K-loop -> loads -> sync -> MMA array -> scale array -> sync -> swap.
            // The `final_accum_frags` array accumulated over K.
            // With single `final_accum_frag_for_mfma_tile`, this accumulation logic must be preserved.

            // Corrected structure for `custom_kernel_rocwmma_pipelined` with refactored inner part:
            // Outer loops for m_frag_block_idx, n_frag_block_idx are for final store and scale lookup.
            // The MMA part must sum over K first.
            // The original `Accumulator_MfmaFrag final_accum_frags[BLOCKS_PER_TILE_M][BLOCKS_PER_TILE_N];`
            // should be brought back, and the `BLOCK_M=128, BLOCK_N=64` will make this array
            // `final_accum_frags[4][2]`, so 8 accumulator fragments. 8 * 16 floats = 128 floats.
            // This is 128 VGPRs, which is acceptable. Max VGPRs is 512.
            // The initial assessment of 256 VGPRs for 128x128 (16 frags) was the issue.
            // For 128x64 (8 frags x 16 VGPRs = 128 VGPRs) or 64x64 (4 frags x 16 VGPRs = 64 VGPRs),
            // using an array of accumulator fragments is feasible.
            // The complex refactoring of loops might not be needed if this register count is fine.
            // Let's revert to the array of fragments strategy with BLOCK_M=128, BLOCK_N=64.

            // --- REVERTING TO ARRAY OF FRAGMENTS, a simpler change for these block sizes ---
            // This means the loop structure for m_frag_block_idx, n_frag_block_idx will NOT
            // contain the K-loop. The K-loop will remain outside, and it will operate on
            // arrays of fragments as in the original "Code to Update".
            // The main changes are then just constants and array sizes.
            // This makes the diff much smaller and aligned with "systematically vary constexpr".
        } // n_frag_block_idx
    } // m_frag_block_idx
    // The above dummy loops are wrong. The original structure should be largely preserved.

    // --- Kernel using arrays of fragments (modified constants) ---
    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);
        }
    }
    
    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_vectorized(
                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_global, block_c_start_col_global,
                wave_lane_id);
        }

        synchronize_workgroup(); 

        Accumulator_MfmaFrag current_k_block_accum_frags[BLOCKS_PER_TILE_M][BLOCKS_PER_TILE_N];
        // Initialize current_k_block_accum_frags to 0.0f
        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_accum_frags[r_frag_idx][c_frag_idx], 0.0f);
            }
        }
        
        // Inner loop over K_OUTER in steps of MFMA_K_TILE_K (16)
        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_for_compute[(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_for_compute[(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_accum_frags[m_frag_idx][n_frag_idx],
                             a_mfma_sub_frags[m_frag_idx],    
                             b_mfma_sub_frags[n_frag_idx],    
                             current_k_block_accum_frags[m_frag_idx][n_frag_idx]);
                }
            }
        } // End k_inner_offset loop
        
        uint32_t k_scale_block_idx = current_k_block_base / SCALE_BLOCK_DIM_K_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_global + m_frag_idx * MFMA_M_TILE_M + m_coord_in_mfma_tile;
                    uint32_t global_n_coord = block_c_start_col_global + n_frag_idx * MFMA_N_TILE_N + n_coord_in_mfma_tile;
                    
                    float scale_a_val;
                     // Boundary checks for M, N on scale access are implicitly handled by gridDim if exact multiple
                    scale_a_val = global_a_scale_ptr[global_m_coord + k_scale_block_idx * M_param];
                    
                    float scale_b_val;
                    uint32_t n_b_scale_block_idx = global_n_coord / SCALE_BLOCK_DIM_N_CONST;
                    scale_b_val = global_b_scale_ptr[n_b_scale_block_idx + k_scale_block_idx * b_scale_num_n_blocks];
                    
                    float combined_scale = scale_a_val * scale_b_val;
                    dest_final_accum_frag.x[p_idx] += source_accum_frag.x[p_idx] * combined_scale;
                }
            }
        }
        
        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_global + m_frag_idx * MFMA_M_TILE_M + m_coord_in_mfma_tile;
                uint32_t global_n_coord = block_c_start_col_global + n_frag_idx * MFMA_N_TILE_N + n_coord_in_mfma_tile;

                // Grid ensures global_m_coord < M_param and 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> 

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; 
  // hipStream_t stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA(); 

  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/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/mdda/hip.hip [ok]
/home/runner/.cache/torch_extensions/py310_cpu/mdda/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/mdda/main.cpp [skipped, no changes]
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/tmp6wqbmjal
[92mSuccessfully preprocessed all matching files.[0m

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


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