# This script provides an optimized FP8 GEMM kernel for AMD MI300X
import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

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 as, torch::Tensor bs, torch::Tensor c);
"""

CUDA_SRC = """
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>
#include <hip/hip_runtime.h>
#include <hip/hip_cooperative_groups.h>

constexpr const int BLOCK = 128;

// MI300X specific configurations
constexpr int TILE_M = 128;
constexpr int TILE_N = 128;
constexpr int TILE_K = 32;
constexpr int WARPS_PER_BLOCK_M = 2;
constexpr int WARPS_PER_BLOCK_N = 2;
constexpr int WARP_SIZE = 64;
constexpr int THREADS = WARPS_PER_BLOCK_M * WARPS_PER_BLOCK_N * WARP_SIZE;

// Convert FP8 to float with optimized path
__device__ __forceinline__ float fp8_to_float(__hip_fp8_e4m3_fnuz val) {
    __half h = __hip_cvt_fp8_to_halfraw(val.__x, __HIP_E4M3_FNUZ);
    return __half2float(h);
}

// Optimized kernel using shared memory and better tiling
template<int BLOCK_M, int BLOCK_N, int BLOCK_K>
__global__ void __launch_bounds__(256, 2)
fp8_gemm_optimized_kernel(
    const __hip_fp8_e4m3_fnuz* __restrict__ a,
    const __hip_fp8_e4m3_fnuz* __restrict__ b,
    const float* __restrict__ as,
    const float* __restrict__ bs,
    __hip_bfloat16* __restrict__ c,
    int m, int n, int k
) {
    // Shared memory for A and B tiles
    __shared__ float smem_a[BLOCK_M][BLOCK_K + 1];  // +1 to avoid bank conflicts
    __shared__ float smem_b[BLOCK_K][BLOCK_N + 1];  // +1 to avoid bank conflicts
    
    // Thread and warp IDs
    const int tid = threadIdx.x;
    const int warp_id = tid / WARP_SIZE;
    const int lane_id = tid % WARP_SIZE;
    
    // Block indices
    const int block_m = blockIdx.x;
    const int block_n = blockIdx.y;
    
    // Warp assignment within block
    const int warp_m = warp_id / WARPS_PER_BLOCK_N;
    const int warp_n = warp_id % WARPS_PER_BLOCK_N;
    
    // Thread assignment for loading
    const int load_a_row = tid / (BLOCK_K / 4);
    const int load_a_col = (tid % (BLOCK_K / 4)) * 4;
    const int load_b_row = tid / (BLOCK_N / 4);
    const int load_b_col = (tid % (BLOCK_N / 4)) * 4;
    
    // Accumulator registers
    float acc[8][8] = {0.0f};
    
    // Scale factor precomputation
    const int scale_n = (n + BLOCK - 1) / BLOCK;
    const int base_m = block_m * BLOCK_M;
    const int base_n = block_n * BLOCK_N;
    
    // Main K loop
    for (int k_block = 0; k_block < k; k_block += BLOCK_K) {
        // Load A tile (coalesced access)
        __syncthreads();
        
        if (load_a_row < BLOCK_M && load_a_col < BLOCK_K) {
            int global_m = base_m + load_a_row;
            int global_k = k_block + load_a_col;
            
            #pragma unroll 4
            for (int i = 0; i < 4; i++) {
                if (global_m < m && global_k + i < k) {
                    __hip_fp8_e4m3_fnuz val = a[global_m + (global_k + i) * m];
                    smem_a[load_a_row][load_a_col + i] = fp8_to_float(val);
                    
                    // Apply A scaling
                    int scale_idx = (global_k + i) / BLOCK;
                    smem_a[load_a_row][load_a_col + i] *= as[global_m + scale_idx * m];
                } else {
                    smem_a[load_a_row][load_a_col + i] = 0.0f;
                }
            }
        }
        
        // Load B tile (coalesced access)
        if (load_b_row < BLOCK_K && load_b_col < BLOCK_N) {
            int global_k = k_block + load_b_row;
            int global_n = base_n + load_b_col;
            
            #pragma unroll 4
            for (int i = 0; i < 4; i++) {
                if (global_k < k && global_n + i < n) {
                    __hip_fp8_e4m3_fnuz val = b[global_n + i + global_k * n];
                    smem_b[load_b_row][load_b_col + i] = fp8_to_float(val);
                    
                    // Apply B scaling
                    int scale_row = (global_n + i) / BLOCK;
                    int scale_col = global_k / BLOCK;
                    smem_b[load_b_row][load_b_col + i] *= bs[scale_row + scale_col * scale_n];
                } else {
                    smem_b[load_b_row][load_b_col + i] = 0.0f;
                }
            }
        }
        
        __syncthreads();
        
        // Compute matrix multiplication
        const int thread_m = warp_m * 64 + (lane_id / 8) * 8;
        const int thread_n = warp_n * 64 + (lane_id % 8) * 8;
        
        #pragma unroll
        for (int k = 0; k < BLOCK_K; k++) {
            #pragma unroll
            for (int i = 0; i < 8; i++) {
                #pragma unroll
                for (int j = 0; j < 8; j++) {
                    if (thread_m + i < BLOCK_M && thread_n + j < BLOCK_N) {
                        acc[i][j] += smem_a[thread_m + i][k] * smem_b[k][thread_n + j];
                    }
                }
            }
        }
    }
    
    // Write results to global memory
    const int global_m = base_m + warp_m * 64 + (lane_id / 8) * 8;
    const int global_n = base_n + warp_n * 64 + (lane_id % 8) * 8;
    
    #pragma unroll
    for (int i = 0; i < 8; i++) {
        #pragma unroll
        for (int j = 0; j < 8; j++) {
            if (global_m + i < m && global_n + j < n) {
                c[(global_m + i) * n + global_n + j] = __float2bfloat16(acc[i][j]);
            }
        }
    }
}

// Medium optimization kernel - stable and faster
template<int TILE_SIZE>
__global__ void fp8_gemm_tiled_stable(
    const __hip_fp8_e4m3_fnuz* __restrict__ a,
    const __hip_fp8_e4m3_fnuz* __restrict__ b,
    const float* __restrict__ as,
    const float* __restrict__ bs,
    __hip_bfloat16* __restrict__ c,
    int m, int n, int k
) {
    extern __shared__ float smem[];
    float* tile_a = smem;
    float* tile_b = smem + TILE_SIZE * TILE_SIZE;
    
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int bx = blockIdx.x;
    const int by = blockIdx.y;
    
    // Global position
    const int row = by * TILE_SIZE + ty;
    const int col = bx * TILE_SIZE + tx;
    
    float sum = 0.0f;
    const int sn = (n + BLOCK - 1) / BLOCK;
    
    // Process K dimension in tiles
    for (int k_start = 0; k_start < k; k_start += TILE_SIZE) {
        // Load tile from A
        if (row < m && k_start + tx < k) {
            __hip_fp8_e4m3_fnuz fp8_val = a[row + (k_start + tx) * m];
            tile_a[ty * TILE_SIZE + tx] = fp8_to_float(fp8_val);
        } else {
            tile_a[ty * TILE_SIZE + tx] = 0.0f;
        }
        
        // Load tile from B
        if (col < n && k_start + ty < k) {
            __hip_fp8_e4m3_fnuz fp8_val = b[col + (k_start + ty) * n];
            tile_b[ty * TILE_SIZE + tx] = fp8_to_float(fp8_val);
        } else {
            tile_b[ty * TILE_SIZE + tx] = 0.0f;
        }
        
        __syncthreads();
        
        // Compute partial product
        #pragma unroll
        for (int i = 0; i < TILE_SIZE; ++i) {
            if (k_start + i < k) {
                float a_val = tile_a[ty * TILE_SIZE + i];
                float b_val = tile_b[i * TILE_SIZE + tx];
                
                // Apply scaling
                int k_idx = k_start + i;
                int scale_k = k_idx / BLOCK;
                
                if (row < m) {
                    a_val *= as[row + scale_k * m];
                }
                if (col < n) {
                    b_val *= bs[col / BLOCK + scale_k * sn];
                }
                
                sum += a_val * b_val;
            }
        }
        
        __syncthreads();
    }
    
    // Write result
    if (row < m && col < n) {
        c[row * n + col] = __float2bfloat16(sum);
    }
}

// Fallback kernel for small sizes
__global__ void fp8_gemm_simple_kernel(
    const __hip_fp8_e4m3_fnuz* __restrict__ a,
    const __hip_fp8_e4m3_fnuz* __restrict__ b,
    const float* __restrict__ as,
    const float* __restrict__ bs,
    __hip_bfloat16* __restrict__ c,
    int m, int n, int k
) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    int col = blockIdx.y * blockDim.y + threadIdx.y;
    
    if (row >= m || col >= n) return;
    
    float sum = 0.0f;
    const int sn = (n + BLOCK - 1) / BLOCK;
    
    // Process in 128-element chunks for scaling
    for (int i = 0; i < k; i += BLOCK) {
        float block_sum = 0.0f;
        
        // Compute dot product for this block
        int block_end = min(i + BLOCK, k);
        for (int j = i; j < block_end; ++j) {
            float a_val = fp8_to_float(a[row + j * m]);
            float b_val = fp8_to_float(b[col + j * n]);
            block_sum += a_val * b_val;
        }
        
        // Apply scaling factors for this block
        int scale_k_idx = i / BLOCK;
        float a_scale = as[row + scale_k_idx * m];
        float b_scale = bs[col / BLOCK + scale_k_idx * sn];
        
        sum += block_sum * a_scale * b_scale;
    }
    
    c[row * n + col] = __float2bfloat16(sum);
}

void fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor as, torch::Tensor bs, torch::Tensor c) {
    int m = a.size(0);
    int n = b.size(0);
    int k = a.size(1);
    
    // Use tiled kernel for medium to large sizes
    if (m >= 128 && n >= 128 && k >= 128) {
        const int TILE_SIZE = 32;
        dim3 block(TILE_SIZE, TILE_SIZE);
        dim3 grid((n + TILE_SIZE - 1) / TILE_SIZE, (m + TILE_SIZE - 1) / TILE_SIZE);
        
        size_t smem_size = 2 * TILE_SIZE * TILE_SIZE * sizeof(float);
        
        hipLaunchKernelGGL(
            HIP_KERNEL_NAME(fp8_gemm_tiled_stable<TILE_SIZE>),
            grid, block, smem_size, 0,
            (__hip_fp8_e4m3_fnuz*)a.data_ptr(),
            (__hip_fp8_e4m3_fnuz*)b.data_ptr(),
            as.data_ptr<float>(),
            bs.data_ptr<float>(),
            (__hip_bfloat16*)c.data_ptr(),
            m, n, k
        );
    } else {
        // Use simple kernel for small sizes
        const int BLOCK_SIZE = 16;
        dim3 block(BLOCK_SIZE, BLOCK_SIZE);
        dim3 grid((m + BLOCK_SIZE - 1) / BLOCK_SIZE, (n + BLOCK_SIZE - 1) / BLOCK_SIZE);
        
        hipLaunchKernelGGL(
            fp8_gemm_simple_kernel,
            grid, block, 0, 0,
            (__hip_fp8_e4m3_fnuz*)a.data_ptr(),
            (__hip_fp8_e4m3_fnuz*)b.data_ptr(),
            as.data_ptr<float>(),
            bs.data_ptr<float>(),
            (__hip_bfloat16*)c.data_ptr(),
            m, n, k
        );
    }
    
    // Check for errors
    hipError_t err = hipGetLastError();
    if (err != hipSuccess) {
        printf("Kernel launch failed: %s\\n", hipGetErrorString(err));
    }
}
"""

# Compiler flags optimized for MI300X
os.environ["CXX"] = "amdclang++"

module = load_inline(
    name='fp8_mm_optimized',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['fp8_mm'],
    verbose=True,
    extra_cuda_cflags=[
        "--offload-arch=gfx942",      # MI300X architecture
        "-std=c++20",
        "-O3",                         # Maximum optimization
        "-ffast-math",                 # Fast math operations
        "-ffinite-math-only",         # Assume no NaN/Inf
        "-famd-opt",                  # AMD-specific optimizations
        "-mllvm", "-amdgpu-early-inline-all=true",  # Aggressive inlining
    ],
)

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/tmps95v0uf8
[92mSuccessfully preprocessed all matching files.[0m

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


Total number of replaced kernel launches: 0
[1/1] amdclang++ 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_optimized.so
