run_arm_docker.sh

#!/bin/bash
# run_arm_docker.sh
# Build and run QSoftmax on ARM via Docker
#
# Usage:
#   ./run_arm_docker.sh              # Build and run full benchmark
#   ./run_arm_docker.sh --build      # Force rebuild
#   ./run_arm_docker.sh --onnx       # ONNX benchmark only
#   ./run_arm_docker.sh --shell      # Interactive shell
#   ./run_arm_docker.sh --clean      # Remove container and images

set -e

IMAGE_NAME="qsoftmax-arm"
CONTAINER_NAME="qsoftmax-benchmark"
RESULTS_DIR="$(pwd)/results"

# Detect architecture
ARCH=$(uname -m)
OS=$(uname -s)

echo "=============================================="
echo "QSoftmax ARM Docker Runner"
echo "=============================================="
echo "Architecture: $ARCH"
echo "OS: $OS"
echo ""

# Check if running on ARM
if [[ "$ARCH" != "arm64" && "$ARCH" != "aarch64" ]]; then
    echo "Warning: Not running on native ARM architecture."
    echo "Docker will use emulation (slower performance)."
    echo ""
    read -p "Continue anyway? [y/N] " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        exit 1
    fi
fi

# Set Docker platform for ARM
export DOCKER_DEFAULT_PLATFORM=linux/arm64

# Create results directory
mkdir -p "$RESULTS_DIR/benchmark" "$RESULTS_DIR/onnx" "$RESULTS_DIR/llama2"

# Parse arguments
FORCE_BUILD=false
ONNX_ONLY=false
INTERACTIVE=false
CLEAN=false

for arg in "$@"; do
    case $arg in
        --build)
            FORCE_BUILD=true
            ;;
        --onnx)
            ONNX_ONLY=true
            ;;
        --shell)
            INTERACTIVE=true
            ;;
        --clean)
            CLEAN=true
            ;;
        --help|-h)
            echo "Usage: $0 [options]"
            echo ""
            echo "Options:"
            echo "  --build    Force rebuild of Docker image"
            echo "  --onnx     Run ONNX benchmark only"
            echo "  --shell    Start interactive shell"
            echo "  --clean    Remove container and images"
            echo "  --help     Show this help"
            exit 0
            ;;
    esac
done

# Clean command
if [ "$CLEAN" = true ]; then
    echo "Cleaning up..."
    docker rm -f "$CONTAINER_NAME" 2>/dev/null || true
    docker rmi -f "$IMAGE_NAME" 2>/dev/null || true
    echo "Done."
    exit 0
fi

# Check if image exists or force build
if [ "$FORCE_BUILD" = true ] || ! docker image inspect "$IMAGE_NAME" > /dev/null 2>&1; then
    echo "Building Docker image for ARM64..."
    echo ""

    docker build \
        --platform linux/arm64 \
        --tag "$IMAGE_NAME:latest" \
        --progress=plain \
        .

    echo ""
    echo "Build complete."
    echo ""
fi

# Remove existing container if present
docker rm -f "$CONTAINER_NAME" 2>/dev/null || true

# Run based on mode
if [ "$INTERACTIVE" = true ]; then
    echo "Starting interactive shell..."
    docker run \
        --rm \
        -it \
        --platform linux/arm64 \
        --name "$CONTAINER_NAME" \
        -v "$(pwd):/app/workspace" \
        -v "$RESULTS_DIR:/app/results" \
        -e OMP_NUM_THREADS=4 \
        -w /app/workspace \
        "$IMAGE_NAME" \
        bash

elif [ "$ONNX_ONLY" = true ]; then
    echo "Running ONNX benchmark..."
    docker run \
        --rm \
        --platform linux/arm64 \
        --name "$CONTAINER_NAME" \
        -v "$RESULTS_DIR:/app/results" \
        -e OMP_NUM_THREADS=4 \
        "$IMAGE_NAME" \
        bash -c "
            python3 qsoftmax_onnx.py --info &&
            python3 qsoftmax_onnx.py --export /app/results/onnx/qsoftmax.onnx --seq-len 512 --head-dim 64 &&
            echo '' &&
            echo 'Benchmarking across sequence lengths...' &&
            for seq in 128 256 512 1024 2048; do
                echo \"=== seq_len=\$seq ===\" &&
                python3 qsoftmax_onnx.py --benchmark /app/results/onnx/qsoftmax.onnx --seq-len \$seq --head-dim 64 --iterations 50
            done
        "

else
    echo "Running full benchmark suite..."
    docker run \
        --rm \
        --platform linux/arm64 \
        --name "$CONTAINER_NAME" \
        -v "$RESULTS_DIR:/app/results" \
        -e OMP_NUM_THREADS=4 \
        "$IMAGE_NAME"
fi

echo ""
echo "=============================================="
echo "Results saved to: $RESULTS_DIR"
echo "=============================================="
ls -la "$RESULTS_DIR"