# 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 for peak performance
constexpr int TILE_M = 64;   // Reduced from 128
constexpr int TILE_N = 64;   // Reduced from 128
constexpr int TILE_K = 32;   // Reduced from 64
constexpr int WARP_M = 16;   // Adjusted
constexpr int WARP_N = 16;   // Adjusted
constexpr int WARP_K = 8;    // Adjusted
constexpr int WARP_SIZE = 64;
constexpr int WARPS_M = TILE_M / WARP_M;  // 4
constexpr int WARPS_N = TILE_N / WARP_N;  // 4
constexpr int THREADS = WARPS_M * WARPS_N * WARP_SIZE;  // 1024

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

// High performance kernel with aggressive optimizations
template<int BM, int BN, int BK>
__global__ void __launch_bounds__(1024, 1)
fp8_gemm_peak_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
) {
    // Double buffering for shared memory
    __shared__ float smem_a[2][BM][BK + 8];  // +8 padding for bank conflict avoidance
    __shared__ float smem_b[2][BK][BN + 8];
    
    // Thread and warp indices
    const int tid = threadIdx.x;
    const int warp_id = tid / WARP_SIZE;
    const int lane_id = tid % WARP_SIZE;
    const int warp_row = warp_id / WARPS_N;
    const int warp_col = warp_id % WARPS_N;
    
    // Block position
    const int block_row = blockIdx.y * BM;
    const int block_col = blockIdx.x * BN;
    
    // Thread assignment for global memory loads (vectorized)
    const int load_a_row = tid / (BK / 4);
    const int load_a_col = (tid % (BK / 4)) * 4;
    const int load_b_row = tid / (BN / 4);
    const int load_b_col = (tid % (BN / 4)) * 4;
    
    // Register storage for accumulation
    float acc[WARP_M / 4][WARP_N / 4][4][4] = {0.0f};  // Now 4x4x4x4
    
    // Prefetch scale factors
    const int scale_n = (n + BLOCK - 1) / BLOCK;
    
    // Pipeline state
    int write_stage = 0;
    int read_stage = 0;
    
    // Prefetch first tile
    if (load_a_row < BM && load_a_col < BK && block_row + load_a_row < m && load_a_col < k) {
        // Load 4 elements individually
        #pragma unroll
        for (int i = 0; i < 4; i++) {
            if (load_a_col + i < k) {
                __hip_fp8_e4m3_fnuz val = a[block_row + load_a_row + (load_a_col + i) * m];
                float fval = fp8_to_float(val);
                
                // Apply scaling
                int scale_idx = (load_a_col + i) / BLOCK;
                fval *= as[block_row + load_a_row + scale_idx * m];
                smem_a[write_stage][load_a_row][load_a_col + i] = fval;
            }
        }
    }
    
    if (load_b_row < BK && load_b_col < BN && block_col + load_b_col < n && load_b_row < k) {
        // Load 4 elements individually
        #pragma unroll
        for (int i = 0; i < 4; i++) {
            if (load_b_col + i < n) {
                __hip_fp8_e4m3_fnuz val = b[block_col + load_b_col + i + load_b_row * n];
                float fval = fp8_to_float(val);
                
                // Apply scaling
                int scale_row = (block_col + load_b_col + i) / BLOCK;
                int scale_col = load_b_row / BLOCK;
                fval *= bs[scale_row + scale_col * scale_n];
                smem_b[write_stage][load_b_row][load_b_col + i] = fval;
            }
        }
    }
    
    __syncthreads();
    write_stage ^= 1;
    
    // Main computation loop
    for (int k_tile = BK; k_tile < k; k_tile += BK) {
        // Prefetch next tile while computing current
        if (load_a_row < BM && load_a_col < BK && 
            block_row + load_a_row < m && k_tile + load_a_col < k) {
            #pragma unroll
            for (int i = 0; i < 4; i++) {
                if (k_tile + load_a_col + i < k) {
                    __hip_fp8_e4m3_fnuz val = a[block_row + load_a_row + (k_tile + load_a_col + i) * m];
                    float fval = fp8_to_float(val);
                    
                    int scale_idx = (k_tile + load_a_col + i) / BLOCK;
                    fval *= as[block_row + load_a_row + scale_idx * m];
                    smem_a[write_stage][load_a_row][load_a_col + i] = fval;
                }
            }
        }
        
        if (load_b_row < BK && load_b_col < BN && 
            block_col + load_b_col < n && k_tile + load_b_row < k) {
            #pragma unroll
            for (int i = 0; i < 4; i++) {
                if (load_b_col + i < n) {
                    __hip_fp8_e4m3_fnuz val = b[block_col + load_b_col + i + (k_tile + load_b_row) * n];
                    float fval = fp8_to_float(val);
                    
                    int scale_row = (block_col + load_b_col + i) / BLOCK;
                    int scale_col = (k_tile + load_b_row) / BLOCK;
                    fval *= bs[scale_row + scale_col * scale_n];
                    smem_b[write_stage][load_b_row][load_b_col + i] = fval;
                }
            }
        }
        
        // Compute on current tile
        #pragma unroll
        for (int k = 0; k < BK; k += 4) {
            // Prefetch from shared to registers
            float a_frag[WARP_M / 4][4];
            float b_frag[4][WARP_N / 4];
            
            #pragma unroll
            for (int i = 0; i < WARP_M / 4; i++) {
                #pragma unroll
                for (int kk = 0; kk < 4; kk++) {
                    a_frag[i][kk] = smem_a[read_stage][warp_row * WARP_M + i * 4 + lane_id % 4][k + kk];
                }
            }
            
            #pragma unroll
            for (int kk = 0; kk < 4; kk++) {
                #pragma unroll
                for (int j = 0; j < WARP_N / 4; j++) {
                    b_frag[kk][j] = smem_b[read_stage][k + kk][warp_col * WARP_N + j * 4 + lane_id / 4];
                }
            }
            
            // Compute outer product
            #pragma unroll
            for (int i = 0; i < WARP_M / 4; i++) {
                #pragma unroll
                for (int j = 0; j < WARP_N / 4; j++) {
                    #pragma unroll
                    for (int ii = 0; ii < 4; ii++) {
                        #pragma unroll
                        for (int jj = 0; jj < 4; jj++) {
                            #pragma unroll
                            for (int kk = 0; kk < 4; kk++) {
                                acc[i][j][ii][jj] += a_frag[i][kk] * b_frag[kk][j];
                            }
                        }
                    }
                }
            }
        }
        
        __syncthreads();
        read_stage ^= 1;
        write_stage ^= 1;
    }
    
    // Process last tile
    #pragma unroll
    for (int k = 0; k < BK; k += 4) {
        float a_frag[WARP_M / 4][4];
        float b_frag[4][WARP_N / 4];
        
        #pragma unroll
        for (int i = 0; i < WARP_M / 4; i++) {
            #pragma unroll
            for (int kk = 0; kk < 4; kk++) {
                a_frag[i][kk] = smem_a[read_stage][warp_row * WARP_M + i * 4 + lane_id % 4][k + kk];
            }
        }
        
        #pragma unroll
        for (int kk = 0; kk < 4; kk++) {
            #pragma unroll
            for (int j = 0; j < WARP_N / 4; j++) {
                b_frag[kk][j] = smem_b[read_stage][k + kk][warp_col * WARP_N + j * 4 + lane_id / 4];
            }
        }
        
        #pragma unroll
        for (int i = 0; i < WARP_M / 4; i++) {
            #pragma unroll
            for (int j = 0; j < WARP_N / 4; j++) {
                #pragma unroll
                for (int ii = 0; ii < 4; ii++) {
                    #pragma unroll
                    for (int jj = 0; jj < 4; jj++) {
                        #pragma unroll
                        for (int kk = 0; kk < 4; kk++) {
                            acc[i][j][ii][jj] += a_frag[i][kk] * b_frag[kk][j];
                        }
                    }
                }
            }
        }
    }
    
    // Write results
    const int thread_row = warp_row * WARP_M + (lane_id % 4) * 4;
    const int thread_col = warp_col * WARP_N + (lane_id / 4) * 4;
    
    #pragma unroll
    for (int i = 0; i < WARP_M / 4; i++) {
        #pragma unroll
        for (int j = 0; j < WARP_N / 4; j++) {
            #pragma unroll
            for (int ii = 0; ii < 4; ii++) {
                #pragma unroll
                for (int jj = 0; jj < 4; jj++) {
                    int row = block_row + thread_row + i * 16 + ii;
                    int col = block_col + thread_col + j * 16 + jj;
                    if (row < m && col < n) {
                        c[row * n + col] = __float2bfloat16(acc[i][j][ii][jj]);
                    }
                }
            }
        }
    }
}

// Large tile kernel for better performance on big matrices
template<int TILE_M, int TILE_N, int TILE_K>
__global__ void __launch_bounds__(256, 2)
fp8_gemm_large_tile_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__ float smem_a[TILE_M][TILE_K + 1];  // +1 for bank conflict avoidance
    __shared__ float smem_b[TILE_K][TILE_N + 1];
    
    const int tid = threadIdx.x + threadIdx.y * blockDim.x;
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    
    const int block_row = blockIdx.y * TILE_M;
    const int block_col = blockIdx.x * TILE_N;
    
    // Accumulator
    float acc[4][4] = {0.0f};
    
    const int thread_row = ty * 4;
    const int thread_col = tx * 4;
    
    const int sn = (n + BLOCK - 1) / BLOCK;
    
    // Main loop
    for (int k_start = 0; k_start < k; k_start += TILE_K) {
        // Cooperative load of A tile
        __syncthreads();
        #pragma unroll
        for (int i = tid; i < TILE_M * TILE_K; i += blockDim.x * blockDim.y) {
            int row = i / TILE_K;
            int col = i % TILE_K;
            int global_row = block_row + row;
            int global_col = k_start + col;
            
            if (global_row < m && global_col < k) {
                __hip_fp8_e4m3_fnuz val = a[global_row + global_col * m];
                float fval = fp8_to_float(val);
                
                // Apply scaling
                int scale_idx = global_col / BLOCK;
                fval *= as[global_row + scale_idx * m];
                
                smem_a[row][col] = fval;
            } else {
                smem_a[row][col] = 0.0f;
            }
        }
        
        // Cooperative load of B tile
        #pragma unroll
        for (int i = tid; i < TILE_K * TILE_N; i += blockDim.x * blockDim.y) {
            int row = i / TILE_N;
            int col = i % TILE_N;
            int global_row = k_start + row;
            int global_col = block_col + col;
            
            if (global_row < k && global_col < n) {
                __hip_fp8_e4m3_fnuz val = b[global_col + global_row * n];
                float fval = fp8_to_float(val);
                
                // Apply scaling
                int scale_row = global_col / BLOCK;
                int scale_col = global_row / BLOCK;
                fval *= bs[scale_row + scale_col * sn];
                
                smem_b[row][col] = fval;
            } else {
                smem_b[row][col] = 0.0f;
            }
        }
        
        __syncthreads();
        
        // Compute 4x4 output per thread
        #pragma unroll
        for (int k = 0; k < TILE_K; k++) {
            #pragma unroll
            for (int i = 0; i < 4; i++) {
                #pragma unroll
                for (int j = 0; j < 4; j++) {
                    if (thread_row + i < TILE_M && thread_col + j < TILE_N) {
                        acc[i][j] += smem_a[thread_row + i][k] * smem_b[k][thread_col + j];
                    }
                }
            }
        }
    }
    
    // Write results
    #pragma unroll
    for (int i = 0; i < 4; i++) {
        #pragma unroll
        for (int j = 0; j < 4; j++) {
            int global_row = block_row + thread_row + i;
            int global_col = block_col + thread_col + j;
            if (global_row < m && global_col < n) {
                c[global_row * n + global_col] = __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;
    
    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;
    
    for (int k_start = 0; k_start < k; k_start += TILE_SIZE) {
        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;
        }
        
        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();
        
        #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];
                
                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();
    }
    
    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;
    
    for (int i = 0; i < k; i += BLOCK) {
        float block_sum = 0.0f;
        
        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;
        }
        
        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 peak performance kernel for benchmark sizes - temporarily disabled for debugging
    if (false && m >= 512 && n >= 512 && k >= 256) {
        dim3 block(THREADS);
        dim3 grid((n + TILE_N - 1) / TILE_N, (m + TILE_M - 1) / TILE_M);
        
        hipLaunchKernelGGL(
            HIP_KERNEL_NAME(fp8_gemm_peak_kernel<TILE_M, TILE_N, TILE_K>),
            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
        );
    }
    // Use large tile kernel for big matrices
    else if (m >= 256 && n >= 256 && k >= 128) {
        const int TILE_M = 64;
        const int TILE_N = 64;
        const int TILE_K = 32;
        dim3 block(16, 16);  // 256 threads
        dim3 grid((n + TILE_N - 1) / TILE_N, (m + TILE_M - 1) / TILE_M);
        
        hipLaunchKernelGGL(
            HIP_KERNEL_NAME(fp8_gemm_large_tile_kernel<TILE_M, TILE_N, TILE_K>),
            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
        );
    }
    // Use tiled kernel for medium to large sizes
    else 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
        "-mllvm", "-amdgpu-function-calls=false",
        "-mllvm", "-enable-post-misched=0",  # Disable post-RA scheduling
        "-mllvm", "-amdgpu-coerce-illegal-types=1",
    ],
)

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/tmpn1fyz_d0
[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++ 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 fp8_mm_optimized.so
