# 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>

constexpr const int BLOCK = 128;

// 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 with tiling and shared memory
__global__ void fp8_gemm_tiled_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
) {
    const int TILE_SIZE = 32;
    const int tid = threadIdx.x + threadIdx.y * blockDim.x;
    
    // Shared memory for tiles
    extern __shared__ float smem[];
    float* tile_a = smem;
    float* tile_b = smem + TILE_SIZE * TILE_SIZE;
    
    // Thread block computes TILE_SIZE x TILE_SIZE output
    const int tile_row = blockIdx.y;
    const int tile_col = blockIdx.x;
    
    // Each thread computes one element
    const int row = tile_row * TILE_SIZE + threadIdx.y;
    const int col = tile_col * TILE_SIZE + threadIdx.x;
    
    float sum = 0.0f;
    
    // Main loop over K dimension
    for (int k_start = 0; k_start < k; k_start += BLOCK) {
        // Process 128x128 blocks for scaling
        for (int k_tile = k_start; k_tile < min(k_start + BLOCK, k); k_tile += TILE_SIZE) {
            // Load A tile
            __syncthreads();
            if (threadIdx.y < TILE_SIZE && threadIdx.x < TILE_SIZE) {
                int a_row = tile_row * TILE_SIZE + threadIdx.y;
                int a_col = k_tile + threadIdx.x;
                if (a_row < m && a_col < k) {
                    tile_a[threadIdx.y * TILE_SIZE + threadIdx.x] = 
                        fp8_to_float(a[a_row + a_col * m]);
                } else {
                    tile_a[threadIdx.y * TILE_SIZE + threadIdx.x] = 0.0f;
                }
            }
            
            // Load B tile
            if (threadIdx.y < TILE_SIZE && threadIdx.x < TILE_SIZE) {
                int b_row = tile_col * TILE_SIZE + threadIdx.x;
                int b_col = k_tile + threadIdx.y;
                if (b_row < n && b_col < k) {
                    tile_b[threadIdx.y * TILE_SIZE + threadIdx.x] = 
                        fp8_to_float(b[b_row + b_col * n]);
                } else {
                    tile_b[threadIdx.y * TILE_SIZE + threadIdx.x] = 0.0f;
                }
            }
            __syncthreads();
            
            // Compute partial dot product
            if (row < m && col < n) {
                float partial_sum = 0.0f;
                #pragma unroll
                for (int i = 0; i < TILE_SIZE; ++i) {
                    partial_sum += tile_a[threadIdx.y * TILE_SIZE + i] * 
                                   tile_b[i * TILE_SIZE + threadIdx.x];
                }
                
                // Apply scaling factors for this 128-block
                int scale_k_idx = k_start / BLOCK;
                float a_scale = (row < m) ? as[row + scale_k_idx * m] : 1.0f;
                float b_scale = (col < n) ? bs[col / BLOCK + scale_k_idx * ((n + BLOCK - 1) / BLOCK)] : 1.0f;
                
                sum += partial_sum * a_scale * b_scale;
            }
        }
    }
    
    // Write result
    if (row < m && col < n) {
        c[row * n + col] = __float2bfloat16(sum);
    }
}

// Simple kernel for all cases (more reliable)
__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);
    
    // For now, use simple kernel for all sizes to ensure correctness
    // Can optimize later once correctness is verified
    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
    ],
)

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/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 [ok]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 0
[1/3] amdclang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=fp8_mm_optimized -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/fp8_mm_optimized/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=fp8_mm_optimized -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 -O3 -ffast-math -ffinite-math-only -c /home/runner/.cache/torch_extensions/py310_cpu/fp8_mm_optimized/hip.hip -o hip.cuda.o 
[3/3] 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

python eval.py test /tmp/tmpdi6vx9oh
[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
ninja: no work to do.
