Dockerfile

# Dockerfile for QSoftmax with ONNX Runtime
# Optimized for ARM64 (Apple Silicon / Neoverse / Graviton)
#
# Build:
#   docker build -t qsoftmax-arm .
#
# Run:
#   docker run --rm qsoftmax-arm
#
# Run with volume mount for results:
#   docker run --rm -v $(pwd)/results:/app/results qsoftmax-arm

# Multi-platform base image for ARM64
FROM --platform=linux/arm64 arm64v8/ubuntu:22.04

LABEL maintainer="QSoftmax Project"
LABEL description="QSoftmax: Quantum-Inspired Attention with ARM NEON and ONNX Runtime"
LABEL architecture="arm64"

# Prevent interactive prompts during package installation
ENV DEBIAN_FRONTEND=noninteractive

# Set locale
ENV LANG=C.UTF-8
ENV LC_ALL=C.UTF-8

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    gcc-12 \
    g++-12 \
    build-essential \
    cmake \
    python3 \
    python3-pip \
    python3-dev \
    python3-venv \
    libomp-dev \
    libopenblas-dev \
    && rm -rf /var/lib/apt/lists/* \
    && apt-get clean

# Set gcc-12 as default compiler
RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 100 && \
    update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-12 100

# Create app directory
WORKDIR /app

# Copy requirements first for better caching
COPY requirements-arm.txt .

# Install Python dependencies optimized for ARM
RUN pip3 install --no-cache-dir --upgrade pip && \
    pip3 install --no-cache-dir -r requirements-arm.txt

# Copy source files
COPY qsoftmax/ ./qsoftmax/
COPY qsoftmax_onnx.py .

# Build C/NEON benchmark with ARM optimizations
WORKDIR /app/qsoftmax

# Detect ARM architecture and build with appropriate flags
RUN ARCH=$(uname -m) && \
    echo "Building for architecture: $ARCH" && \
    if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then \
        # Try Neoverse-optimized build first (AWS Graviton, Ampere, etc.)
        gcc -O3 -mcpu=neoverse-n1 -ffast-math -fopenmp \
            -DARM_NEON -ftree-vectorize \
            -o qsoftmax_bench benchmark_c.c -lm -fopenmp 2>/dev/null || \
        # Fall back to generic ARM64 build
        gcc -O3 -march=armv8-a -ffast-math -fopenmp \
            -DARM_NEON -ftree-vectorize \
            -o qsoftmax_bench benchmark_c.c -lm -fopenmp 2>/dev/null || \
        # Last resort: basic optimized build
        gcc -O3 -ffast-math -o qsoftmax_bench benchmark_c.c -lm; \
    else \
        echo "Warning: Non-ARM architecture detected"; \
        gcc -O3 -ffast-math -o qsoftmax_bench benchmark_c.c -lm; \
    fi

WORKDIR /app

# Create output directories
RUN mkdir -p results/benchmark results/onnx results/llama2

# Create entrypoint script
RUN cat > /app/run_all.sh << 'SCRIPT'
#!/bin/bash
set -e

echo "=============================================="
echo "QSoftmax Benchmark Suite for ARM"
echo "=============================================="
echo ""
echo "Platform: $(uname -m)"
echo "Kernel: $(uname -r)"
echo "CPU: $(cat /proc/cpuinfo | grep 'model name' | head -1 | cut -d: -f2 || echo 'ARM Processor')"
echo ""

# Run C/NEON benchmark
echo "=== C/NEON Benchmark ==="
cd /app/qsoftmax
./qsoftmax_bench | tee /app/results/benchmark/c_neon_results.txt
cd /app
echo ""

# Run Python benchmark
echo "=== Python Benchmark ==="
python3 qsoftmax/benchmark.py --output-dir /app/results/benchmark
echo ""

# Export and benchmark ONNX model
echo "=== ONNX Export and Benchmark ==="
python3 qsoftmax_onnx.py --info
echo ""
python3 qsoftmax_onnx.py --export /app/results/onnx/qsoftmax.onnx --seq-len 512 --head-dim 64
echo ""
python3 qsoftmax_onnx.py --benchmark /app/results/onnx/qsoftmax.onnx --seq-len 512 --head-dim 64 --iterations 100 | tee /app/results/onnx/benchmark_results.txt
echo ""

# Run Llama-2 benchmark if available
if [ -f "qsoftmax/llama2_benchmark.py" ]; then
    echo "=== Llama-2 Pattern Benchmark ==="
    python3 qsoftmax/llama2_benchmark.py --output-dir /app/results/llama2 --int4 2>/dev/null || \
        echo "Llama-2 benchmark skipped (missing dependencies)"
    echo ""
fi

echo "=============================================="
echo "All benchmarks complete!"
echo "Results saved to /app/results/"
echo "=============================================="
SCRIPT

RUN chmod +x /app/run_all.sh

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
    CMD python3 -c "import numpy; import onnxruntime; print('OK')" || exit 1

# Default command
CMD ["/app/run_all.sh"]