install_optimized.sh

#!/usr/bin/env bash
set -euo pipefail

# Optimized (non-portable) build helper.
# - Linux x86_64: native tuning enabled by default.
# - macOS: portable by default; opt into native tuning with --native-macos
#   or ACTIONET_MACOS_NATIVE=1.
# OpenMP runtime is always required. Set ACTIONET_OPENMP_RUNTIME to
# GNU, INTEL, or LLVM to force a specific runtime; AUTO (default) lets
# CMake detect the compiler's preferred runtime.
#
# GPU support (Linux x86_64 + NVIDIA only):
#   Pass --gpu or set LIBACTIONET_ENABLE_NVIDIA_GPU=ON to opt in.
#   Optionally set CUDAToolkit_ROOT to a CUDA >= 12.2 install directory.
#   Optionally set CMAKE_CUDA_ARCHITECTURES to override the default (80;86;89;90).
#
# Usage:
#   ./install_optimized.sh [--native-macos] [--gpu] [--verbose] [pip args...]
#   ACTIONET_MACOS_NATIVE=1 ./install_optimized.sh [pip args...]
#   LIBACTIONET_ENABLE_NVIDIA_GPU=ON CUDAToolkit_ROOT=/usr/local/cuda ./install_optimized.sh [pip args...]
# Pass --verbose (or -v) to forward the verbose flag to pip.

here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$here"

os="$(uname -s)"
arch="$(uname -m)"
runtime="${ACTIONET_OPENMP_RUNTIME:-AUTO}"  # set to INTEL/GNU/LLVM to force, otherwise auto

enable_macos_native="${ACTIONET_MACOS_NATIVE:-0}"
enable_gpu="${LIBACTIONET_ENABLE_NVIDIA_GPU:-OFF}"
declare -a pip_args=()
for arg in "$@"; do
    if [[ "$arg" == "--native-macos" ]]; then
        enable_macos_native=1
        continue
    fi
    if [[ "$arg" == "--gpu" ]]; then
        enable_gpu="ON"
        continue
    fi
    pip_args+=("$arg")
done

run_pip_install() {
    python -m pip install . "${pip_args[@]}"
}

# Append GPU cmake flags to a cmake_args string (passed by name-ref).
append_gpu_args() {
    local -n _ref=$1
    if [[ "$enable_gpu" == "ON" ]]; then
        _ref+=" -DLIBACTIONET_ENABLE_NVIDIA_GPU=ON"
        if [[ -n "${CUDAToolkit_ROOT:-}" ]]; then
            _ref+=" -DCUDAToolkit_ROOT=${CUDAToolkit_ROOT}"
        fi
        if [[ -n "${CMAKE_CUDA_ARCHITECTURES:-}" ]]; then
            _ref+=" -DCMAKE_CUDA_ARCHITECTURES=${CMAKE_CUDA_ARCHITECTURES}"
        fi
    fi
}

if [[ "$os" == "Darwin" ]]; then
    if [[ "$enable_macos_native" != "1" ]]; then
        echo "[INFO] macOS detected; running portable build."
        echo "[INFO] Use --native-macos (or ACTIONET_MACOS_NATIVE=1) to enable non-portable native CPU tuning."
        cmake_args=""
        if [[ "$runtime" != "AUTO" ]]; then
            cmake_args+=" -DLIBACTIONET_OPENMP_RUNTIME=${runtime}"
        fi
        append_gpu_args cmake_args
        if [[ -n "${cmake_args// }" ]]; then
            export CMAKE_ARGS="${CMAKE_ARGS:-}${cmake_args}"
            echo "[INFO] Using CMAKE_ARGS=${CMAKE_ARGS}"
        fi
        run_pip_install
        exit 0
    fi

    macos_common_flags="-O3 -ffp-contract=fast -fno-strict-aliasing"
    if [[ "$arch" == "arm64" || "$arch" == "aarch64" ]]; then
        native_cxxflags="-mcpu=native ${macos_common_flags}"
    elif [[ "$arch" == "x86_64" || "$arch" == "amd64" ]]; then
        native_cxxflags="-march=native -mtune=native ${macos_common_flags}"
    else
        echo "[WARN] Unrecognized macOS arch ($arch); using generic optimization flags without native CPU target."
        native_cxxflags="${macos_common_flags}"
    fi

    cmake_args="-DCMAKE_CXX_FLAGS='${native_cxxflags}' -DCMAKE_INTERPROCEDURAL_OPTIMIZATION=ON"
    if [[ "$runtime" != "AUTO" ]]; then
        cmake_args+=" -DLIBACTIONET_OPENMP_RUNTIME=${runtime}"
    fi
    append_gpu_args cmake_args

    export CMAKE_ARGS="${CMAKE_ARGS:-} ${cmake_args}"
    echo "[INFO] Using CMAKE_ARGS=${CMAKE_ARGS}"
    echo "[INFO] Invoking pip install with macOS native optimized flags..."
    run_pip_install
    exit 0
fi

if [[ "$arch" != "x86_64" && "$arch" != "amd64" ]]; then
    echo "[WARN] Non-x86_64 architecture ($arch) detected; skipping native flags. Running portable build."
    cmake_args=""
    if [[ "$runtime" != "AUTO" ]]; then
        cmake_args+=" -DLIBACTIONET_OPENMP_RUNTIME=${runtime}"
    fi
    append_gpu_args cmake_args
    if [[ -n "${cmake_args// }" ]]; then
        export CMAKE_ARGS="${CMAKE_ARGS:-}${cmake_args}"
        echo "[INFO] Using CMAKE_ARGS=${CMAKE_ARGS}"
    fi
    run_pip_install
    exit 0
fi

native_cxxflags="-march=native -mtune=native -O3 -ffp-contract=fast -funroll-loops -fomit-frame-pointer -fno-strict-aliasing"

cmake_args="-DCMAKE_CXX_FLAGS='${native_cxxflags}' -DCMAKE_INTERPROCEDURAL_OPTIMIZATION=ON"
if [[ "$runtime" != "AUTO" ]]; then
    cmake_args+=" -DLIBACTIONET_OPENMP_RUNTIME=${runtime}"
fi
append_gpu_args cmake_args

export CMAKE_ARGS="${CMAKE_ARGS:-} ${cmake_args}"

echo "[INFO] Using CMAKE_ARGS=${CMAKE_ARGS}"
echo "[INFO] Invoking pip install with optimized flags..."
run_pip_install