#!/bin/bash
# Compile GPU kernels for RustyTorch++

set -e

# Source shell config for CUDA paths
source ~/.zshrc 2>/dev/null || source ~/.bashrc 2>/dev/null

# Configuration
KERNEL_CACHE_DIR="target/kernel_cache"
GPU_ARCH="${GPU_ARCH:-sm_120}"
RUSTG_PATH="../rust/rustg"

echo "=== RustyTorch++ Kernel Compilation ==="
echo "GPU Architecture: $GPU_ARCH"
echo "CUDA Version: $(nvcc --version | grep release | awk '{print $5}' | sed 's/,//')"

# Create cache directories
mkdir -p "$KERNEL_CACHE_DIR/$GPU_ARCH"

# Function to compile a CUDA kernel
compile_cuda_kernel() {
    local kernel_name=$1
    local source_file=$2
    local output_file="$KERNEL_CACHE_DIR/$GPU_ARCH/${kernel_name}.ptx"
    
    echo "Compiling kernel: $kernel_name"
    nvcc -ptx \
        --gpu-architecture=$GPU_ARCH \
        -O3 \
        --use_fast_math \
        --ptxas-options=-v \
        -o "$output_file" \
        "$source_file"
    
    if [ $? -eq 0 ]; then
        echo "✓ Compiled $kernel_name successfully"
        echo "  Output: $output_file"
    else
        echo "✗ Failed to compile $kernel_name"
        exit 1
    fi
}

# Example: Compile vector_add kernel if CUDA source exists
if [ -f "crates/rtx-kernel/src/kernels/vector_add.cu" ]; then
    compile_cuda_kernel "vector_add" "crates/rtx-kernel/src/kernels/vector_add.cu"
fi

# Use cargo-g if available for Rust GPU compilation
if [ -x "$RUSTG_PATH/cargo-g/target/release/cargo-g" ]; then
    echo "Using cargo-g for Rust GPU compilation..."
    "$RUSTG_PATH/cargo-g/target/release/cargo-g" build --release
elif [ -x "$RUSTG_PATH/mock-cargo-g.sh" ]; then
    echo "Using mock cargo-g for testing..."
    "$RUSTG_PATH/mock-cargo-g.sh" build --release
fi

echo "=== Kernel compilation complete ==="
ls -la "$KERNEL_CACHE_DIR/$GPU_ARCH/" 2>/dev/null || echo "No kernels compiled yet"