60 lines
1.8 KiB
Bash
Executable File
60 lines
1.8 KiB
Bash
Executable File
#!/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" |