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

# 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; 
constexpr uint32_t MFMA_N_TILE_N = 32u; 
constexpr uint32_t MFMA_K_TILE_K = 16u; 

// Output C tile computed by one thread block - MODIFIED
constexpr uint32_t TILE_M_PER_BLOCK = 32u; // MODIFIED from 64u to reduce VGPR pressure
constexpr uint32_t TILE_N_PER_BLOCK = 64u; // Retained from 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 output tile - MODIFIED
constexpr uint32_t MFMA_TILES_M_PER_BLOCK = TILE_M_PER_BLOCK / MFMA_M_TILE_M; // 32/32 = 1
constexpr uint32_t MFMA_TILES_N_PER_BLOCK = TILE_N_PER_BLOCK / MFMA_N_TILE_N; // 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; 
constexpr uint32_t TOTAL_THREADS_PER_BLOCK = TBLOCK_X_DIM * TBLOCK_Y_DIM;

// 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, GlobalLayoutA>;
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>;


__global__ void __launch_bounds__(TOTAL_THREADS_PER_BLOCK)
    custom_kernel_rocwmma(
    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, int N_out, int K) 
{
    // Shared memory for A and B tiles - MODIFIED SIZES
    __shared__ fp8_t lds_a_tile_ptr[TILE_M_PER_BLOCK * BLOCK_K_OUTER]; // 32x128 = 4KB
    __shared__ fp8_t lds_b_tile_ptr[TILE_N_PER_BLOCK * BLOCK_K_OUTER]; // 64x128 = 8KB

    // Shared memory for scales - NEW
    __shared__ float lds_a_scale_cache[TILE_M_PER_BLOCK]; // 32 floats for a_scale
    __shared__ float lds_b_scale_cache[2]; // Max 2 floats needed for b_scale if N-tile (64) spans N-scale-blocks (128)

    // Global C tile base coordinates for this thread block
    uint32_t block_c_base_m_dim = blockIdx.x * TILE_M_PER_BLOCK;
    uint32_t block_c_base_n_dim = blockIdx.y * TILE_N_PER_BLOCK;

    // Final accumulators for the entire TILE_M_PER_BLOCK x TILE_N_PER_BLOCK output C tile.
    // Array dimensions match MFMA_TILES_M_PER_BLOCK x MFMA_TILES_N_PER_BLOCK.
    Accumulator_MfmaFrag final_accum_frags[MFMA_TILES_M_PER_BLOCK][MFMA_TILES_N_PER_BLOCK]; // [1][2]

    // Initialize final accumulators to 0
    for (uint32_t i = 0; i < MFMA_TILES_M_PER_BLOCK; ++i) { // Loop iterates once for i=0
        for (uint32_t j = 0; j < MFMA_TILES_N_PER_BLOCK; ++j) { // Loop iterates twice for j=0,1
            fill_fragment(final_accum_frags[i][j], 0.0f);
        }
    }
    
    uint32_t wave_lane_id = threadIdx.x; // Linear thread ID within the block/wave (0-63)

    // --- Outer loop over K dimension in steps of BLOCK_K_OUTER (128) ---
    for (int k_outer_base = 0; k_outer_base < K; k_outer_base += BLOCK_K_OUTER) {
        // --- 1. Load A and B tiles from global to LDS ---
        constexpr uint32_t lds_a_elements = TILE_M_PER_BLOCK * BLOCK_K_OUTER;
        constexpr uint32_t lds_b_elements = TILE_N_PER_BLOCK * BLOCK_K_OUTER;
        constexpr uint32_t loads_per_thread_a = (lds_a_elements + TOTAL_THREADS_PER_BLOCK - 1) / TOTAL_THREADS_PER_BLOCK;
        constexpr uint32_t loads_per_thread_b = (lds_b_elements + TOTAL_THREADS_PER_BLOCK - 1) / TOTAL_THREADS_PER_BLOCK;

        // Load A tile (TILE_M_PER_BLOCK x BLOCK_K_OUTER)
        for (uint32_t i = 0; i < loads_per_thread_a; ++i) {
            uint32_t lds_flat_idx = wave_lane_id * loads_per_thread_a + i;
            if (lds_flat_idx < lds_a_elements) {
                uint32_t lds_m = lds_flat_idx % TILE_M_PER_BLOCK; 
                uint32_t lds_k = lds_flat_idx / TILE_M_PER_BLOCK;
                
                uint32_t gmem_a_m = block_c_base_m_dim + lds_m;
                uint32_t gmem_a_k = k_outer_base + lds_k;

                if (gmem_a_m < M && gmem_a_k < K) {
                    lds_a_tile_ptr[lds_k * TILE_M_PER_BLOCK + lds_m] = global_a_ptr[gmem_a_k * M + gmem_a_m];
                } else {
                    lds_a_tile_ptr[lds_k * TILE_M_PER_BLOCK + lds_m] = static_cast<fp8_t>(0.0f);
                }
            }
        }
        // Load B tile (TILE_N_PER_BLOCK x BLOCK_K_OUTER)
        for (uint32_t i = 0; i < loads_per_thread_b; ++i) {
            uint32_t lds_flat_idx = wave_lane_id * loads_per_thread_b + i;
            if (lds_flat_idx < lds_b_elements) {
                uint32_t lds_n = lds_flat_idx % TILE_N_PER_BLOCK; 
                uint32_t lds_k = lds_flat_idx / TILE_N_PER_BLOCK;

                uint32_t gmem_b_n = block_c_base_n_dim + lds_n;
                uint32_t gmem_b_k = k_outer_base + lds_k;
                
                if (gmem_b_n < N_out && gmem_b_k < K) {
                    lds_b_tile_ptr[lds_k * TILE_N_PER_BLOCK + lds_n] = global_b_ptr[gmem_b_k * N_out + gmem_b_n];
                } else {
                    lds_b_tile_ptr[lds_k * TILE_N_PER_BLOCK + lds_n] = static_cast<fp8_t>(0.0f);
                }
            }
        }
        
        // --- 1.1 Load scales for this k_outer_base into LDS --- NEW
        uint32_t k_scale_idx = k_outer_base / SCALE_BLOCK_DIM_K_CONST;
        
        // Load a_scales (TILE_M_PER_BLOCK floats)
        // Each of the first TILE_M_PER_BLOCK threads loads one scale factor.
        if (wave_lane_id < TILE_M_PER_BLOCK) {
            uint32_t cur_m_global = block_c_base_m_dim + wave_lane_id; // wave_lane_id is m_offset here
            if (cur_m_global < M) {
                lds_a_scale_cache[wave_lane_id] = global_a_scale_ptr[cur_m_global + k_scale_idx * M];
            } else {
                lds_a_scale_cache[wave_lane_id] = 0.0f; 
            }
        }

        // Load b_scales (1 or 2 floats)
        uint32_t n_scale_idx_start_abs = block_c_base_n_dim / SCALE_BLOCK_DIM_N_CONST;
        uint32_t n_scale_idx_end_abs   = (block_c_base_n_dim + TILE_N_PER_BLOCK - 1) / SCALE_BLOCK_DIM_N_CONST;
        uint32_t num_b_scales_to_load = n_scale_idx_end_abs - n_scale_idx_start_abs + 1;
        uint32_t b_scale_global_stride_N_blocks = (N_out + SCALE_BLOCK_DIM_N_CONST - 1) / SCALE_BLOCK_DIM_N_CONST;

        if (wave_lane_id == 0) {
            uint32_t n_global_rep0 = block_c_base_n_dim; // Representative N for first scale of this N-tile
            if (n_global_rep0 < N_out) {
                 lds_b_scale_cache[0] = global_b_scale_ptr[n_scale_idx_start_abs + k_scale_idx * b_scale_global_stride_N_blocks];
            } else {
                 lds_b_scale_cache[0] = 0.0f; 
            }
        }
        if (num_b_scales_to_load == 2 && wave_lane_id == 1) {
            // Smallest N global coord that uses the second scale factor relevant to this block
            uint32_t n_global_rep1 = n_scale_idx_end_abs * SCALE_BLOCK_DIM_N_CONST; 
            if (n_global_rep1 < N_out) {
                lds_b_scale_cache[1] = global_b_scale_ptr[n_scale_idx_end_abs + k_scale_idx * b_scale_global_stride_N_blocks];
            } else {
                lds_b_scale_cache[1] = 0.0f;
            }
        }
        synchronize_workgroup(); // Ensure LDS A, B, and scales are populated

        // Accumulators for the sum over the current BLOCK_K_OUTER (128 K-depth)
        Accumulator_MfmaFrag current_k_block_accum_frags[MFMA_TILES_M_PER_BLOCK][MFMA_TILES_N_PER_BLOCK]; // [1][2]
        for (uint32_t r = 0; r < MFMA_TILES_M_PER_BLOCK; ++r) { // r=0
            for (uint32_t c = 0; c < MFMA_TILES_N_PER_BLOCK; ++c) { // c=0,1
                fill_fragment(current_k_block_accum_frags[r][c], 0.0f);
            }
        }
        
        // --- Inner loop over K_OUTER in steps of MFMA_K_TILE_K (16) --- UNROLLED
        #pragma unroll
        for (int k_inner_step = 0; k_inner_step < (BLOCK_K_OUTER / MFMA_K_TILE_K); ++k_inner_step) {
            int k_inner = k_inner_step * MFMA_K_TILE_K;
            
            MatrixA_MfmaFrag a_frag; 
            MatrixB_MfmaFrag b_frags[MFMA_TILES_N_PER_BLOCK]; // For n_idx = 0, 1

            // Load A fragment (32x16 tile from 32x128 LDS A tile)
            // m_idx is 0 since MFMA_TILES_M_PER_BLOCK = 1
            const fp8_t* lds_a_k_slice_ptr = &lds_a_tile_ptr[k_inner * TILE_M_PER_BLOCK];
            load_matrix_sync(a_frag, lds_a_k_slice_ptr, TILE_M_PER_BLOCK); 

            // Load B fragments (16x32 tiles from 64x128 LDS B tile)
            for (uint32_t n_idx = 0; n_idx < MFMA_TILES_N_PER_BLOCK; ++n_idx) { // n_idx = 0, 1
                const fp8_t* lds_b_k_slice_ptr = &lds_b_tile_ptr[k_inner * TILE_N_PER_BLOCK + (n_idx * MFMA_N_TILE_N)];
                load_matrix_sync(b_frags[n_idx], lds_b_k_slice_ptr, TILE_N_PER_BLOCK); 
            }
            
            // Perform MFMA for each 32x32 output sub-tile
            // m_fidx is 0 since MFMA_TILES_M_PER_BLOCK = 1
            for (uint32_t n_fidx = 0; n_fidx < MFMA_TILES_N_PER_BLOCK; ++n_fidx) { // n_fidx = 0, 1
                mma_sync(current_k_block_accum_frags[0][n_fidx],
                         a_frag, 
                         b_frags[n_fidx],
                         current_k_block_accum_frags[0][n_fidx]);
            }
        } // End k_inner loop (accumulated over 128 K-depth)
        
        // --- Apply scaling factors for this K_OUTER_BLOCK and add to final_accum_frags ---
        // m_fidx is 0
        for (uint32_t n_fidx = 0; n_fidx < MFMA_TILES_N_PER_BLOCK; ++n_fidx) { // n_fidx = 0, 1
            auto& src_k_block_acc_frag = current_k_block_accum_frags[0][n_fidx];
            auto& dst_final_acc_frag = final_accum_frags[0][n_fidx];

            // Each thread processes its elements in the accumulator fragment
            for (int p = 0; p < Accumulator_MfmaFrag::num_elements; ++p) {
                // Map (wave_lane_id, p) to logical (m_tile, n_tile) coords within this 32x32 MFMA output tile
                uint32_t p_div_4 = p / 4;
                uint32_t p_mod_4 = p % 4;
                uint32_t lane_div_32 = wave_lane_id / 32;
                uint32_t lane_mod_32 = wave_lane_id % 32;
                uint32_t m_in_mfma_tile = (8 * p_div_4) + (4 * lane_div_32) + p_mod_4; // 0..31 relative to this MFMA tile's M
                uint32_t n_in_mfma_tile = lane_mod_32;                                  // 0..31 relative to this MFMA tile's N

                // Global coordinates for this specific accumulator element
                // m_fidx is 0, so m_global is relative to block_c_base_m_dim
                uint32_t m_global = block_c_base_m_dim + m_in_mfma_tile;
                uint32_t n_global = block_c_base_n_dim + n_fidx * MFMA_N_TILE_N + n_in_mfma_tile;

                float scale_a_val;
                if (m_global < M) {
                     // m_in_mfma_tile is local to the TILE_M_PER_BLOCK dimension.
                     scale_a_val = lds_a_scale_cache[m_in_mfma_tile]; 
                } else { // Element is out of M bounds for C output
                    // Value in src_k_block_acc_frag.x[p] for out-of-bound m_global should be 0 due to LDS padding.
                    // Adding 0 ensures correctness if it wasn't exactly 0.
                    dst_final_acc_frag.x[p] += 0.0f; 
                    continue; 
                }

                float scale_b_val;
                if (n_global < N_out) {
                    uint32_t n_global_b_scale_block_idx_abs = n_global / SCALE_BLOCK_DIM_N_CONST;
                    uint32_t cache_idx_b = n_global_b_scale_block_idx_abs - n_scale_idx_start_abs;
                    scale_b_val = lds_b_scale_cache[cache_idx_b];
                } else { // Element is out of N bounds for C output
                    dst_final_acc_frag.x[p] += 0.0f;
                    continue;
                }
                
                float combined_scale = scale_a_val * scale_b_val;
                dst_final_acc_frag.x[p] += src_k_block_acc_frag.x[p] * combined_scale;
            }
        }
    } // End k_outer_base loop

    // --- Store final_accum_frags to global C (M x N_out, row-major) ---
    // m_fidx is 0
    for (uint32_t n_fidx = 0; n_fidx < MFMA_TILES_N_PER_BLOCK; ++n_fidx) { // n_fidx = 0, 1
        auto& result_frag_to_store = final_accum_frags[0][n_fidx]; 

        // Each thread stores its elements from the accumulator fragment
        for (int p = 0; p < Accumulator_MfmaFrag::num_elements; ++p) {
            uint32_t p_div_4 = p / 4;
            uint32_t p_mod_4 = p % 4;
            uint32_t lane_div_32 = wave_lane_id / 32;
            uint32_t lane_mod_32 = wave_lane_id % 32;
            uint32_t m_in_mfma_tile = (8 * p_div_4) + (4 * lane_div_32) + p_mod_4;
            uint32_t n_in_mfma_tile = lane_mod_32;

            uint32_t m_global = block_c_base_m_dim + m_in_mfma_tile;
            uint32_t n_global = block_c_base_n_dim + n_fidx * MFMA_N_TILE_N + n_in_mfma_tile;

            if (m_global < M && n_global < N_out) {
                global_c_ptr[m_global * N_out + n_global] = static_cast<bf16_t>(result_frag_to_store.x[p]);
            }
        }
    }
}

// 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); 


  // Updated grid dimensions for new TILE_M_PER_BLOCK
  dim3 grid_dim((M + TILE_M_PER_BLOCK - 1) / TILE_M_PER_BLOCK, 
                (N_out + TILE_N_PER_BLOCK - 1) / TILE_N_PER_BLOCK);
  dim3 block_dim(TBLOCK_X_DIM, TBLOCK_Y_DIM); 
  
  hipStream_t stream = nullptr; // Use default stream or get current stream if needed

  custom_kernel_rocwmma<<<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();
  TORCH_CHECK(last_error == hipSuccess, "HIP kernel launch failed: ", hipGetErrorString(last_error));
}
"""

import os
os.environ["CXX"] = "clang++"

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 eval.py leaderboard /tmp/tmp8idcl0ll
[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
