Files
rustytorch/scripts/run_pinn_comparison.sh
T
2026-03-04 00:08:42 +00:00

413 lines
14 KiB
Bash
Executable File

#!/bin/bash
# =============================================================================
# Unified PyTorch vs RustyTorch++ PINN Benchmark
# =============================================================================
#
# This script runs comprehensive PINN benchmarks comparing PyTorch and RustyTorch++
# on both CPU and GPU, generating a unified HTML report with performance charts.
#
# Usage:
# ./scripts/run_pinn_comparison.sh [OPTIONS]
#
# Options:
# --quick Run quick benchmarks (fewer iterations/point sizes)
# --cpu-only Run only CPU benchmarks (skip GPU)
# --gpu-only Run only GPU benchmarks (skip CPU)
# --no-venv Use system Python instead of virtual environment
# -h, --help Show this help message
#
# Output:
# benchmark_reports/<hostname>-rust-py-pinn-benchmark-<MM-DD-YYYY>/
# - system_info.json
# - pytorch_cpu_results.json
# - pytorch_gpu_results.json
# - rust_cpu_output.txt
# - rust_gpu_output.txt
# - comparison_report.html
# - comparison_report.md
#
# =============================================================================
set -e
# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m' # No Color
# =============================================================================
# Progress Indicator Functions
# =============================================================================
# Track start time
START_SECONDS=$SECONDS
TOTAL_STEPS=5
CURRENT_STEP=0
# Format time as Xm Ys
format_time() {
local secs=$1
if [[ $secs -lt 60 ]]; then
printf "%ds" $secs
else
printf "%dm %ds" $((secs/60)) $((secs%60))
fi
}
# Display progress bar with current step and elapsed time
show_progress() {
local step=$1
local step_name=$2
local width=40
local percent=$((step * 100 / TOTAL_STEPS))
local filled=$((width * step / TOTAL_STEPS))
local empty=$((width - filled))
local elapsed=$((SECONDS - START_SECONDS))
# Build progress bar
local bar=""
for ((i=0; i<filled; i++)); do bar+="█"; done
for ((i=0; i<empty; i++)); do bar+="░"; done
# Clear previous progress display (3 lines + header box)
echo -ne "\033[6A\033[J"
# Draw progress box
echo -e "${CYAN}╔══════════════════════════════════════════════════════════════════════╗${NC}"
echo -e "${CYAN}${BOLD}PyTorch vs RustyTorch++ PINN Benchmark${NC}${CYAN}${NC}"
echo -e "${CYAN}╠══════════════════════════════════════════════════════════════════════╣${NC}"
printf "${CYAN}${NC} Progress: [${GREEN}%s${NC}] %3d%% (%d/%d steps) ${CYAN}${NC}\n" "$bar" "$percent" "$step" "$TOTAL_STEPS"
printf "${CYAN}${NC} Current: %-54s ${CYAN}${NC}\n" "$step_name"
printf "${CYAN}${NC} Elapsed: %-54s ${CYAN}${NC}\n" "$(format_time $elapsed)"
echo -e "${CYAN}╚══════════════════════════════════════════════════════════════════════╝${NC}"
}
# Initial progress display (placeholder before first update)
init_progress() {
echo ""
echo ""
echo ""
echo ""
echo ""
echo ""
echo ""
}
# Script location and project root
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
PINN_DIR="$PROJECT_ROOT/examples/pinn_mre_helmholtz"
# Use python3 by default, allow override via PYTHON env var
PYTHON="${PYTHON:-python3}"
# Configuration defaults
QUICK_MODE=false
CPU_ONLY=false
GPU_ONLY=false
VENV_MODE=true # Always use venv by default
# Auto-detect platform
if [[ "$(uname)" == "Darwin" ]]; then
GPU_BACKEND="metal"
PLATFORM="macOS"
else
GPU_BACKEND="cuda"
PLATFORM="Linux"
fi
# Hostname and date for output directory
HOSTNAME_SHORT=$(hostname | cut -d'.' -f1)
DATE_STR=$(date +"%m-%d-%Y")
OUTPUT_DIR="$PROJECT_ROOT/benchmark_reports/${HOSTNAME_SHORT}-rust-py-pinn-benchmark-${DATE_STR}"
# =============================================================================
# Parse command line arguments
# =============================================================================
show_help() {
echo "Unified PyTorch vs RustyTorch++ PINN Benchmark"
echo ""
echo "Usage: $0 [OPTIONS]"
echo ""
echo "Options:"
echo " --quick Run quick benchmarks (fewer iterations)"
echo " --cpu-only Run only CPU benchmarks (skip GPU)"
echo " --gpu-only Run only GPU benchmarks (skip CPU)"
echo " --no-venv Use system Python instead of virtual environment"
echo " -h, --help Show this help message"
echo ""
echo "Output:"
echo " benchmark_reports/<hostname>-rust-py-pinn-benchmark-<date>/"
}
while [[ $# -gt 0 ]]; do
case $1 in
--quick)
QUICK_MODE=true
shift
;;
--cpu-only)
CPU_ONLY=true
shift
;;
--gpu-only)
GPU_ONLY=true
shift
;;
--no-venv)
VENV_MODE=false
shift
;;
-h|--help)
show_help
exit 0
;;
*)
echo "Unknown option: $1"
show_help
exit 1
;;
esac
done
# Determine what to run
RUN_CPU=true
RUN_GPU=true
if [[ "$GPU_ONLY" == true ]]; then
RUN_CPU=false
RUN_GPU=true
elif [[ "$CPU_ONLY" == true ]]; then
RUN_CPU=true
RUN_GPU=false
fi
# Set point sizes based on quick mode
if [[ "$QUICK_MODE" == true ]]; then
POINT_SIZES="200,1000"
RUST_BENCH_ARGS="--noplot"
else
POINT_SIZES="200,1000,10000"
RUST_BENCH_ARGS=""
fi
# =============================================================================
# Banner and Initial Progress Display
# =============================================================================
echo -e "Host: ${YELLOW}${HOSTNAME_SHORT}${NC} | Platform: ${YELLOW}${PLATFORM}${NC} | GPU: ${YELLOW}${GPU_BACKEND}${NC}"
echo -e "Output: ${CYAN}${OUTPUT_DIR}${NC}"
echo ""
# Initialize progress display (creates space for progress box)
init_progress
show_progress 0 "Initializing..."
# =============================================================================
# Create output directory
# =============================================================================
mkdir -p "$OUTPUT_DIR"
# =============================================================================
# Phase 1: Collect System Information
# =============================================================================
show_progress 1 "Collecting system information..."
SYSTEM_INFO_FILE="$OUTPUT_DIR/system_info.json"
# Collect OS info based on platform
if [[ "$PLATFORM" == "macOS" ]]; then
OS_NAME="macOS"
OS_VERSION=$(sw_vers -productVersion 2>/dev/null || echo "Unknown")
KERNEL=$(uname -r)
CPU_MODEL=$(sysctl -n machdep.cpu.brand_string 2>/dev/null || echo "Apple Silicon")
CPU_CORES=$(sysctl -n hw.physicalcpu 2>/dev/null || echo "Unknown")
CPU_THREADS=$(sysctl -n hw.logicalcpu 2>/dev/null || echo "Unknown")
TOTAL_MEM_BYTES=$(sysctl -n hw.memsize 2>/dev/null || echo "0")
TOTAL_MEM="$((TOTAL_MEM_BYTES / 1024 / 1024 / 1024)) GB"
GPU_NAME="Apple Silicon GPU"
GPU_MEMORY="Unified Memory"
GPU_DRIVER="Metal"
GPU_COMPUTE="N/A"
else
OS_NAME=$(cat /etc/os-release 2>/dev/null | grep "^NAME=" | cut -d'"' -f2 || echo "Unknown")
OS_VERSION=$(cat /etc/os-release 2>/dev/null | grep "^VERSION=" | cut -d'"' -f2 || echo "Unknown")
KERNEL=$(uname -r)
CPU_MODEL=$(lscpu 2>/dev/null | grep "Model name" | cut -d':' -f2 | xargs || echo "Unknown")
CPU_CORES=$(lscpu 2>/dev/null | grep "^CPU(s):" | cut -d':' -f2 | xargs || echo "Unknown")
CPU_THREADS=$(lscpu 2>/dev/null | grep "Thread(s) per core" | cut -d':' -f2 | xargs || echo "1")
TOTAL_MEM=$(free -h 2>/dev/null | grep "Mem:" | awk '{print $2}' || echo "Unknown")
# GPU info
if command -v nvidia-smi &> /dev/null; then
GPU_INFO=$(nvidia-smi --query-gpu=name,memory.total,driver_version,compute_cap --format=csv,noheader 2>/dev/null || echo "")
if [[ -n "$GPU_INFO" ]]; then
GPU_NAME=$(echo "$GPU_INFO" | cut -d',' -f1 | xargs)
GPU_MEMORY=$(echo "$GPU_INFO" | cut -d',' -f2 | xargs)
GPU_DRIVER=$(echo "$GPU_INFO" | cut -d',' -f3 | xargs)
GPU_COMPUTE=$(echo "$GPU_INFO" | cut -d',' -f4 | xargs)
else
GPU_NAME="Not detected"
GPU_MEMORY="N/A"
GPU_DRIVER="N/A"
GPU_COMPUTE="N/A"
fi
else
GPU_NAME="nvidia-smi not available"
GPU_MEMORY="N/A"
GPU_DRIVER="N/A"
GPU_COMPUTE="N/A"
fi
fi
# Get git info
GIT_COMMIT=$(cd "$PROJECT_ROOT" && git rev-parse --short HEAD 2>/dev/null || echo "unknown")
GIT_BRANCH=$(cd "$PROJECT_ROOT" && git branch --show-current 2>/dev/null || echo "unknown")
# Write system info JSON
cat > "$SYSTEM_INFO_FILE" << EOF
{
"timestamp": "$(date -Iseconds)",
"hostname": "$HOSTNAME_SHORT",
"commit": "$GIT_COMMIT",
"branch": "$GIT_BRANCH",
"os": {
"name": "$OS_NAME",
"version": "$OS_VERSION",
"kernel": "$KERNEL",
"arch": "$(uname -m)"
},
"cpu": {
"model": "$CPU_MODEL",
"cores": "$CPU_CORES",
"threads_per_core": "$CPU_THREADS"
},
"gpu": {
"name": "$GPU_NAME",
"memory": "$GPU_MEMORY",
"driver": "$GPU_DRIVER",
"compute_capability": "$GPU_COMPUTE"
},
"memory": {
"total": "$TOTAL_MEM"
}
}
EOF
# =============================================================================
# Phase 2: Setup Python Environment
# =============================================================================
show_progress 2 "Setting up Python environment..."
cd "$PINN_DIR"
if [[ "$VENV_MODE" == true ]]; then
if [[ ! -d ".venv" ]]; then
$PYTHON -m venv .venv
fi
source .venv/bin/activate
PYTHON="python"
# Install dependencies if needed
if [[ -f "python/requirements.txt" ]]; then
pip install -q -r python/requirements.txt 2>/dev/null || true
fi
fi
# Check if Python and torch are available
if ! $PYTHON -c "import torch" 2>/dev/null; then
SKIP_PYTHON=true
else
SKIP_PYTHON=false
fi
# =============================================================================
# Phase 3: Run PyTorch Benchmarks
# =============================================================================
show_progress 3 "Running PyTorch benchmarks..."
if [[ "$SKIP_PYTHON" != true ]]; then
if [[ "$RUN_CPU" == true ]]; then
show_progress 3 "Running PyTorch CPU benchmarks..."
$PYTHON python/benchmark_runner.py \
--device cpu \
--point-sizes "$POINT_SIZES" \
--output "$OUTPUT_DIR/pytorch_cpu_results.json" >/dev/null 2>&1 || true
fi
if [[ "$RUN_GPU" == true ]]; then
if $PYTHON -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then
show_progress 3 "Running PyTorch GPU benchmarks..."
$PYTHON python/benchmark_runner.py \
--device cuda \
--point-sizes "$POINT_SIZES" \
--output "$OUTPUT_DIR/pytorch_gpu_results.json" >/dev/null 2>&1 || true
fi
fi
fi
# =============================================================================
# Phase 4: Run Rust Benchmarks
# =============================================================================
show_progress 4 "Running RustyTorch++ benchmarks..."
cd "$PINN_DIR"
if [[ "$RUN_CPU" == true ]]; then
show_progress 4 "Running RustyTorch++ CPU benchmarks..."
cargo bench --bench pinn_benchmark $RUST_BENCH_ARGS 2>&1 > "$OUTPUT_DIR/rust_cpu_output.txt" || true
fi
if [[ "$RUN_GPU" == true ]]; then
if cargo check --features cuda 2>/dev/null; then
show_progress 4 "Running RustyTorch++ GPU benchmarks..."
cargo bench --bench pinn_benchmark --features cuda $RUST_BENCH_ARGS 2>&1 > "$OUTPUT_DIR/rust_gpu_output.txt" || true
fi
fi
# =============================================================================
# Phase 5: Generate Report
# =============================================================================
show_progress 5 "Generating comparison report..."
cd "$PROJECT_ROOT"
# Build report generator arguments
REPORT_ARGS="--output-dir $OUTPUT_DIR"
REPORT_ARGS="$REPORT_ARGS --system-info $OUTPUT_DIR/system_info.json"
[[ -f "$OUTPUT_DIR/pytorch_cpu_results.json" ]] && REPORT_ARGS="$REPORT_ARGS --pytorch-cpu $OUTPUT_DIR/pytorch_cpu_results.json"
[[ -f "$OUTPUT_DIR/pytorch_gpu_results.json" ]] && REPORT_ARGS="$REPORT_ARGS --pytorch-gpu $OUTPUT_DIR/pytorch_gpu_results.json"
[[ -f "$OUTPUT_DIR/rust_cpu_output.txt" ]] && REPORT_ARGS="$REPORT_ARGS --rust-cpu $OUTPUT_DIR/rust_cpu_output.txt"
[[ -f "$OUTPUT_DIR/rust_gpu_output.txt" ]] && REPORT_ARGS="$REPORT_ARGS --rust-gpu $OUTPUT_DIR/rust_gpu_output.txt"
$PYTHON "$SCRIPT_DIR/generate_comparison_report.py" $REPORT_ARGS >/dev/null 2>&1
# =============================================================================
# Completion - Clear progress and show final summary
# =============================================================================
TOTAL_ELAPSED=$((SECONDS - START_SECONDS))
# Clear the progress display
echo -ne "\033[7A\033[J"
# Final completion banner
echo -e "${GREEN}╔════════════════════════════════════════════════════════════════════╗${NC}"
echo -e "${GREEN}${BOLD}PyTorch vs RustyTorch++ PINN Benchmark Complete!${NC}${GREEN}${NC}"
echo -e "${GREEN}╚════════════════════════════════════════════════════════════════════╝${NC}"
echo ""
echo -e "Output: ${CYAN}$OUTPUT_DIR${NC}"
echo ""
echo -e "Total time: ${YELLOW}$(format_time $TOTAL_ELAPSED)${NC}"
echo ""
echo "Files generated:"
ls -1 "$OUTPUT_DIR" | while read f; do
echo -e " - ${CYAN}$f${NC}"
done
echo ""
echo -e "View report: ${YELLOW}firefox $OUTPUT_DIR/comparison_report.html${NC}"
echo ""