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

151 lines
4.2 KiB
Bash
Executable File

#!/bin/bash
# Metal Kernel Benchmark: CubeCL vs Hand-crafted on Apple Silicon
#
# Usage:
# ./scripts/run_metal_kernel_bench.sh
# ./scripts/run_metal_kernel_bench.sh --sizes 128,256,512
# ./scripts/run_metal_kernel_bench.sh --iterations 200 --output-format json
set -e
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
echo ""
echo "=============================================="
echo " RustyTorch Metal Kernel Benchmark"
echo " CubeCL vs Hand-crafted MSL Kernels"
echo "=============================================="
echo ""
# Check we're on macOS
if [[ "$(uname)" != "Darwin" ]]; then
echo "ERROR: This script must be run on macOS (requires Metal)"
exit 1
fi
# Default parameters
OPERATIONS="${OPERATIONS:-add,mul,exp,relu,gelu,silu}"
SIZES="${SIZES:-64,128,256}"
ITERATIONS="${ITERATIONS:-100}"
WARMUP="${WARMUP:-10}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-table}"
OUTPUT_FILE=""
VERIFY="${VERIFY:-false}"
# Parse command line arguments
while [[ $# -gt 0 ]]; do
case $1 in
--operations)
OPERATIONS="$2"
shift 2
;;
--sizes)
SIZES="$2"
shift 2
;;
--iterations)
ITERATIONS="$2"
shift 2
;;
--warmup)
WARMUP="$2"
shift 2
;;
--output-format)
OUTPUT_FORMAT="$2"
shift 2
;;
--output|-o)
OUTPUT_FILE="$2"
shift 2
;;
--verify)
VERIFY="true"
shift
;;
--help|-h)
echo "Usage: $0 [OPTIONS]"
echo ""
echo "Options:"
echo " --operations OPS Comma-separated operations (default: add,mul,exp,relu,gelu,silu)"
echo " --sizes SIZES Comma-separated tensor sizes (default: 64,128,256)"
echo " --iterations N Number of benchmark iterations (default: 100)"
echo " --warmup N Number of warmup iterations (default: 10)"
echo " --output-format FMT Output format: table, markdown, json, csv (default: table)"
echo " --output, -o FILE Output file path"
echo " --verify Verify outputs match between implementations"
echo " --help, -h Show this help message"
exit 0
;;
*)
echo "Unknown option: $1"
exit 1
;;
esac
done
cd "$PROJECT_ROOT"
echo "System Info:"
echo " macOS Version: $(sw_vers -productVersion)"
echo " Chip: $(sysctl -n machdep.cpu.brand_string 2>/dev/null || system_profiler SPHardwareDataType | grep 'Chip' | awk -F': ' '{print $2}')"
echo ""
echo "Benchmark Configuration:"
echo " Operations: $OPERATIONS"
echo " Sizes: $SIZES"
echo " Iterations: $ITERATIONS (warmup: $WARMUP)"
echo " Output format: $OUTPUT_FORMAT"
echo ""
# Build the benchmark binary with Metal support
echo "Building benchmark binary..."
cargo build --release -p rtx-kernel-bench --bin compare-kernels 2>&1 | tail -5
if [[ ! -f "$PROJECT_ROOT/target/release/compare-kernels" ]]; then
echo "ERROR: Failed to build compare-kernels binary"
exit 1
fi
echo ""
echo "Running benchmark..."
echo ""
# Build command arguments
CMD_ARGS=(
--backend metal
--operations "$OPERATIONS"
--sizes "$SIZES"
--iterations "$ITERATIONS"
--warmup "$WARMUP"
--output-format "$OUTPUT_FORMAT"
)
if [[ -n "$OUTPUT_FILE" ]]; then
# Create output directory if needed
mkdir -p "$(dirname "$OUTPUT_FILE")"
CMD_ARGS+=(--output "$OUTPUT_FILE")
fi
if [[ "$VERIFY" == "true" ]]; then
CMD_ARGS+=(--verify)
fi
# Run the benchmark
"$PROJECT_ROOT/target/release/compare-kernels" "${CMD_ARGS[@]}"
# Save results with timestamp if no output file specified
if [[ -z "$OUTPUT_FILE" && "$OUTPUT_FORMAT" == "markdown" ]]; then
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
RESULT_FILE="$PROJECT_ROOT/results/metal_benchmark_${TIMESTAMP}.md"
echo ""
echo "Saving results to: $RESULT_FILE"
"$PROJECT_ROOT/target/release/compare-kernels" "${CMD_ARGS[@]}" > "$RESULT_FILE"
fi
echo ""
echo "=============================================="
echo " Benchmark Complete"
echo "=============================================="