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
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; // Single wave per block, responsible for the entire BLOCK_M x BLOCK_N tile
constexpr uint32_t TOTAL_THREADS_PER_BLOCK = TBLOCK_X_DIM * TBLOCK_Y_DIM;

// Number of elements per thread for different fragment types
// constexpr int NUM_FP8_PER_THREAD_A_B = 8;   // For 32x32x16 fp8 MFMA, A/B fragments (rocWMMA detail)
constexpr int NUM_FLOAT_PER_THREAD_ACC = 16; // For 32x32x16 fp8 MFMA, float accumulator (rocWMMA detail)

// rocWMMA Fragment Types
using namespace rocwmma;

using GlobalLayoutA = col_major; // A is M x K, col-major
using GlobalLayoutB = col_major; // B is N x K, col-major
using GlobalLayoutC = row_major; // Output C is M N row-major

// Input fragments for one MFMA op (A: MFMA_M_TILE_M x MFMA_K_TILE_K, B: MFMA_K_TILE_K x MFMA_N_TILE_N)
// LDS for A is K_OUTER rows of BLOCK_M elements. load_matrix_sync needs pitch BLOCK_M.
// MatrixA_MfmaFrag is col_major, consistent with global A.
using MatrixA_MfmaFrag = fragment<matrix_a, MFMA_M_TILE_M, MFMA_N_TILE_N, MFMA_K_TILE_K, fp8_t, col_major>;
// LDS for B is K_OUTER rows of BLOCK_N elements. This is KxN row-major storage.
// MatrixB_MfmaFrag is row_major, for B.T operation using this KxN LDS tile.
using MatrixB_MfmaFrag = fragment<matrix_b, MFMA_M_TILE_M, MFMA_N_TILE_N, MFMA_K_TILE_K, fp8_t, row_major>;
// Accumulator fragment for one MFMA op output (MFMA_M_TILE_M x MFMA_N_TILE_N floats)
// The mapping of p_idx to (m,n) coords is based on MFMA's native distribution,
// usually corresponding to a col_major view of the accumulator tile.
// However, original code used GlobalLayoutC (row_major). Keep this to see effect, or change to col_major if mapping requires it.
// The problem's text implies the mapping is for the hardware layout, which rocWMMA abstracts via frag.x[p].
// Sticking with GlobalLayoutC for now.
using Accumulator_MfmaFrag = fragment<accumulator, MFMA_M_TILE_M, MFMA_N_TILE_N, MFMA_K_TILE_K, float, GlobalLayoutC>;

// --- Helper Device Function: Load a tile from Global to LDS (Vectorized) ---
// LDS layout: A is K_OUTER rows of BLOCK_M elements. B is K_OUTER rows of BLOCK_N elements.
__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, // 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)
{
    constexpr uint32_t VECTOR_SIZE_FP8 = 4; // Load 4 fp8s (one int) at a time

    // 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 (col-major within row sense, or row-major for K rows of M elements).
    // Target LDS storage: lds_a_target[k_in_tile * BLOCK_M + m_in_tile]
    constexpr uint32_t VECTORS_IN_M_DIM_A = BLOCK_M / VECTOR_SIZE_FP8;      // Num vectors along M-dim of one K-slice in LDS (e.g., 64/4 = 16)
    constexpr uint32_t THREAD_GROUP_SIZE_A = VECTORS_IN_M_DIM_A;           // Threads per group for loading M-dim vectors (16 for BLOCK_M=64)
    constexpr uint32_t NUM_THREAD_GROUPS_A = TOTAL_THREADS_PER_BLOCK / THREAD_GROUP_SIZE_A; // (e.g., 64/16 = 4 groups)

    uint32_t group_id_a = wave_lane_id / THREAD_GROUP_SIZE_A;      // This thread's group ID (0 to NUM_THREAD_GROUPS_A-1)
    uint32_t lane_in_group_a = wave_lane_id % THREAD_GROUP_SIZE_A; // This thread's ID within its group (0 to THREAD_GROUP_SIZE_A-1)

    // Each group of threads handles a distinct set of k_idx_in_tile values initially.
    // Iterate if BLOCK_K_OUTER > NUM_THREAD_GROUPS_A (e.g. 128 k-slices / 4 groups = 32 iterations per thread group)
    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;

        // lane_in_group_a determines the vector index along M for this k_idx_in_tile
        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; // M-offset for this vector within the block
        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) { // Ensure gmem access is within bounds
            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;
        }
    }

    // Load B: Global B (N x K, col-major) to LDS B (K_OUTER x BLOCK_N)
    // LDS B stores K_OUTER rows, each of BLOCK_N elements.
    // Target LDS storage: lds_b_target[k_in_tile * BLOCK_N + n_in_tile]
    constexpr uint32_t VECTORS_IN_N_DIM_B = BLOCK_N / VECTOR_SIZE_FP8;      // (16 for BLOCK_N=64)
    constexpr uint32_t THREAD_GROUP_SIZE_B = VECTORS_IN_N_DIM_B;           // (16 threads)
    constexpr uint32_t NUM_THREAD_GROUPS_B = TOTAL_THREADS_PER_BLOCK / THREAD_GROUP_SIZE_B; // (4 groups)

    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; // N-offset for this vector
        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) { // Ensure gmem access is within bounds
            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;
        }
    }
}


// --- Helper Device Function: Compute MMAs on data in LDS ---
__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]) // Output param
{
    // Initialize current_k_block_acc_frags to 0.0f for this K-block computation
    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);
        }
    }
    
    // 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]; // Stores 32x16 tiles of A for the current k_inner_offset
        MatrixB_MfmaFrag b_mfma_sub_frags[BLOCKS_PER_TILE_N]; // Stores 16x32 tiles of B for the current k_inner_offset

        // Load A fragments for all M-tiles of the block
        for (uint32_t m_tile_idx = 0; m_tile_idx < BLOCKS_PER_TILE_M; ++m_tile_idx) {
            // LDS A is (K_OUTER rows, BLOCK_M cols). Pitch for rocWMMA load is BLOCK_M.
            // Pointer to A_LDS[k_inner_offset, m_tile_idx * MFMA_M_TILE_M]
            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);
        }

        // Load B fragments for all N-tiles of the block
        for (uint32_t n_tile_idx = 0; n_tile_idx < BLOCKS_PER_TILE_N; ++n_tile_idx) {
            // LDS B is (K_OUTER rows, BLOCK_N cols). Pitch for rocWMMA load is BLOCK_N.
            // Pointer to B_LDS[k_inner_offset, n_tile_idx * MFMA_N_TILE_N]
            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);
        }
        
        // Perform MFMA for each 32x32 output sub-tile within the 64x64 block
        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],    // A is M_frag x K_frag
                         b_mfma_sub_frags[n_frag_idx],    // B is K_frag x N_frag
                         current_k_block_acc_frags[m_frag_idx][n_frag_idx]);
            }
        }
    } // End k_inner_offset loop
}


__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 ---
    __shared__ fp8_t lds_a_ping[BLOCK_M * BLOCK_K_OUTER]; // 64x128 = 8192 bytes
    __shared__ fp8_t lds_a_pong[BLOCK_M * BLOCK_K_OUTER]; // 64x128 = 8192 bytes
    __shared__ fp8_t lds_b_ping[BLOCK_N * BLOCK_K_OUTER]; // 64x128 = 8192 bytes
    __shared__ fp8_t lds_b_pong[BLOCK_N * BLOCK_K_OUTER]; // 64x128 = 8192 bytes
    // Total LDS: 4 * 8192 = 32KB

    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; // 0 for ping, 1 for pong

    // Global C tile base coordinates for this thread block
    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; // Linear thread ID within the block/wave (0-63)

    // Final accumulators for the entire BLOCK_M x BLOCK_N output C tile.
    // Stored in registers.
    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; } // No computation if K is zero.

    // --- Pre-load first tile into current buffer_selector (0 = ping) ---
    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, /* k_current_block_base */
        M_param, N_out_param, K_param,
        block_c_start_row, block_c_start_col,
        wave_lane_id);

    // K_param is divisible by BLOCK_K_OUTER from problem statement
    int num_k_outer_iterations = K_param / BLOCK_K_OUTER;

    // --- Outer loop over K dimension with pipelining ---
    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;

        // Set compute pointers to the currently loaded buffer (ping or pong)
        fp8_t* lds_a_for_compute = lds_a_double_buffers[buffer_selector];
        fp8_t* lds_b_for_compute = lds_b_double_buffers[buffer_selector];

        // --- Start loading next tile (if one exists) into the *other* buffer ---
        int next_k_block_base = current_k_block_base + BLOCK_K_OUTER;
        if (next_k_block_base < K_param) { // True for all but the last iteration's "next load"
            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, block_c_start_col,
                wave_lane_id);
        }

        // Wait for the current compute buffer to be fully loaded.
        // (This was loaded in the *previous* iteration's "next load" step)
        synchronize_workgroup();

        // Accumulators for the sum over the current BLOCK_K_OUTER (register based)
        Accumulator_MfmaFrag current_k_block_accum_frags[BLOCKS_PER_TILE_M][BLOCKS_PER_TILE_N];
        
        // --- Compute on current tile (data in lds_a_for_compute, lds_b_for_compute) ---
        compute_lds_tile_mma(lds_a_for_compute, lds_b_for_compute, current_k_block_accum_frags);
        
        // --- Apply scaling factors for this K_OUTER_BLOCK and add to final_accum_frags ---
        // This part happens after compute on current tile is finished (sequential).
        // It can overlap with the load of the *next* tile that was initiated above.
        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;
                }
            }
        }
        
        // Wait for the load of the *next* tile (if started) to complete before swapping buffers.
        // The compute and scale application for the current tile are already finished.
        if (next_k_block_base < K_param) { // Only sync if a load was actually dispatched
             synchronize_workgroup();
        }

        // --- Swap buffer roles for the next iteration ---
        buffer_selector = 1 - buffer_selector;
    } // End k_iter_idx loop
    
    // --- Store final_accum_frags to global C (M x N_out_param, row-major) ---
    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> 

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; // Use default stream
  // hipStream_t stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA(); // If using PyTorch's stream

  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) { // TORCH_CHECK can't easily take hipGetErrorString
      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++ 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

python eval.py benchmark /tmp/tmpel_mxe9y
[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
[1/1] 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
