413 lines
14 KiB
Bash
Executable File
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 ""
|