CODES-compile-instructions.sh

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

# Switches
swm_enable=0
union_enable=0
torch_enable=0

# Uncomment below for MPICH
#export PATH=/usr/local/mpich-4.1.2/bin/:"$PATH"
# Note: remember to compile MPICH with nemesis not with UCX support

################## Actual scripts starts from here ##################

# SWM has to be enabled for UNION to work
if [ $union_enable = 1 ]; then
    swm_enable=1
fi

# What to compile
CUR_DIR="$PWD"



##### Downloading everything #####

if [ ! -d codes/.git ]; then
    git clone https://github.com/codes-org/codes --depth=100 --branch=v1.5.0
else
    echo "Using existing codes checkout: $(realpath codes)"
fi

if [ ! -d ross/.git ]; then
    git clone https://github.com/ross-org/ross --depth=100 --branch=v8.1.0
else
    echo "Using existing ross checkout: $(realpath ross)"
fi

if [ $swm_enable = 1 ]; then
    if [ ! -d argobots/.git ]; then
    git clone https://github.com/pmodels/argobots --depth=1
else
    echo "Using existing argobots checkout: $(realpath argobots)"
fi
    if [ ! -d swm-workloads/.git ]; then
    git clone https://github.com/codes-org/swm-workloads --branch=v1.2
else
    echo "Using existing swm-workloads checkout: $(realpath swm-workloads)"
fi
fi

if [ $union_enable = 1 ]; then
    # Downloading conceptual
    curl -L https://sourceforge.net/projects/conceptual/files/conceptual/1.5.1b/conceptual-1.5.1b.tar.gz -o conceptual-1.5.1b.tar.gz
    tar xvf conceptual-1.5.1b.tar.gz
    # Downloading union
    if [ ! -d Union/.git ]; then
        git clone https://github.com/SPEAR-UIC/Union
    else
        echo "Using existing Union checkout: $(realpath Union)"
    fi
    pushd Union && git checkout 99b3df3 && popd
fi

##### COMPILING #####

mkdir -p ross/build
pushd ross/build

# ---- CODES CUDA arch autodetection from CUDA_HOME ----
# CUDA is opt-in: set CUDA_HOME to enable GPU/Torch CUDA compilation.
# The script derives CMAKE_CUDA_ARCHITECTURES, CUDAARCHS, and
# TORCH_CUDA_ARCH_LIST automatically so users do not have to export them.
if [ -n "${CUDA_HOME:-}" ]; then
    if [ ! -x "$CUDA_HOME/bin/nvcc" ]; then
        echo "ERROR: CUDA_HOME is set to '$CUDA_HOME', but '$CUDA_HOME/bin/nvcc' is not executable."
        exit 1
    fi

    cuda_compute_cap=""

    if command -v nvidia-smi >/dev/null 2>&1; then
        cuda_compute_cap="$(
            nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null \
                | head -n 1 \
                | tr -d '[:space:]' \
                || true
        )"
    fi

    if [ -z "$cuda_compute_cap" ]; then
        cuda_compute_cap="$(
            python3 - <<'PY_DETECT_CUDA_ARCH'
try:
    import torch
    if torch.cuda.is_available():
        major, minor = torch.cuda.get_device_capability(0)
        print(f"{major}.{minor}")
except Exception:
    pass
PY_DETECT_CUDA_ARCH
        )"
    fi

    if [ -z "$cuda_compute_cap" ]; then
        echo "ERROR: CUDA_HOME is set, but GPU compute capability could not be detected."
        echo "       Check nvidia-smi and torch.cuda.is_available()."
        echo "       Refusing to continue because an empty CUDA arch causes nvcc to receive compute_."
        exit 1
    fi

    cuda_arch="${cuda_compute_cap/./}"

    if [ -z "$cuda_arch" ]; then
        echo "ERROR: detected empty CUDA architecture from compute capability '$cuda_compute_cap'."
        exit 1
    fi

    export CMAKE_CUDA_ARCHITECTURES="$cuda_arch"
    export CUDAARCHS="$cuda_arch"
    export TORCH_CUDA_ARCH_LIST="$cuda_compute_cap"

    echo "CUDA enabled via CUDA_HOME=$CUDA_HOME"
    echo "Detected CUDA compute capability: $cuda_compute_cap"
    echo "Using CMAKE_CUDA_ARCHITECTURES=$CMAKE_CUDA_ARCHITECTURES"
    echo "Using TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST"
else
    unset CMAKE_CUDA_ARCHITECTURES
    unset CUDAARCHS
    unset TORCH_CUDA_ARCH_LIST
    echo "CUDA_HOME is not set; building without CUDA arch overrides."
fi
# ---- end CODES CUDA arch autodetection ----

cmake .. -DROSS_BUILD_MODELS=ON -DCMAKE_INSTALL_PREFIX="$(realpath ./bin)" \
  -DCMAKE_BUILD_TYPE=Debug -DCMAKE_C_FLAGS="-g -Wall"
#make VERBOSE=1
make install -j4
err=$?
[[ $err -ne 0 ]] && exit $err
popd

if [ $swm_enable = 1 ]; then
    pushd swm-workloads/swm
    ./prepare.sh
    mkdir -p build
    pushd build
    ../configure --disable-shared --prefix="$(realpath ./bin)" CC=mpicc CXX=mpicxx CFLAGS=-g CXXFLAGS=-g
    #make V=1 && make install
    make -j4 && make install
    err=$?
    [[ $err -ne 0 ]] && exit $err
    popd && popd

    pushd argobots
    ./autogen.sh
    mkdir -p build
    pushd build
    #../configure --enable-debug=all --disable-fast --disable-shared --prefix="$(realpath ./bin)" CC=mpicc CXX=mpicxx CFLAGS=-g CXXFLAGS=-g
    ../configure --disable-shared --prefix="$(realpath ./bin)" CC=mpicc CXX=mpicxx CFLAGS=-g CXXFLAGS=-g
    #make V=1 && make install
    make -j4 && make install
    err=$?
    [[ $err -ne 0 ]] && exit $err
    popd && popd
fi

if [ $union_enable = 1 ]; then
    pushd conceptual-1.5.1b
    PYTHON=python2 ./configure --prefix="$(realpath ./install)" LIBS=-lm
    make -j4 && make install
    err=$?
    [[ $err -ne 0 ]] && exit $err
    popd

    pushd Union
    # Python 2 override. Union expects Python 2 ONLY
    mkdir -p python-override
    if [ -e python-override/python ] || [ -L python-override/python ]; then
    rm -f python-override/python
fi
ln -s /usr/bin/python2 python-override/python
    # compiling
    ./prepare.sh
    PYTHON=python2 ./configure --disable-shared --with-conceptual="$(realpath ../conceptual-1.5.1b/install)" --with-conceptual-src="$(realpath ../conceptual-1.5.1b)" --prefix="$(realpath ./install)" CC=mpicc CXX=mpicxx
    PATH="$PWD/python-override:$PATH" make -j4 && make install
    err=$?
    [[ $err -ne 0 ]] && exit $err
    popd
fi



# Make system pkg-config metadata visible even when Conda's pkg-config is active.
# This is needed for libzmq.pc on systems where ZeroMQ is installed through the OS
# but the active Conda environment's pkg-config only searches Conda pkgconfig dirs.
if ! pkg-config --exists libzmq 2>/dev/null; then
    for pcdir in \
        /usr/lib/x86_64-linux-gnu/pkgconfig \
        /usr/lib64/pkgconfig \
        /usr/lib/pkgconfig \
        /usr/local/lib/pkgconfig \
        /usr/local/lib64/pkgconfig \
        /opt/homebrew/lib/pkgconfig \
        /usr/share/pkgconfig
    do
        if [ -d "$pcdir" ]; then
            export PKG_CONFIG_PATH="$pcdir:${PKG_CONFIG_PATH:-}"
        fi
    done
fi

if ! pkg-config --exists libzmq 2>/dev/null; then
    echo "WARNING: pkg-config still cannot find libzmq.pc." >&2
    echo "         If ZMQML fails to build, install the ZeroMQ development package" >&2
    echo "         or set PKG_CONFIG_PATH to the directory containing libzmq.pc." >&2
fi

# The zmqml requester is built by CODES' own CMake (src/surrogate/zmqml) when
# CODES_USE_ZEROMQ resolves on — no separate make step. The pkg-config setup
# above is what lets CMake find libzmq; rapidjson is picked up from the system
# include path.

mkdir -p codes/build
pushd codes/build

torch_cmake_prefix=""
torch_dir=""

if [ "$torch_enable" = 1 ]; then
    torch_cmake_prefix="$(python3 - <<'INNERPY'
import torch
print(torch.utils.cmake_prefix_path)
INNERPY
)"
    torch_dir="${torch_cmake_prefix}/Torch"

    if [ ! -f "${torch_dir}/TorchConfig.cmake" ]; then
        echo "ERROR: TorchConfig.cmake not found at: ${torch_dir}/TorchConfig.cmake" >&2
        echo "       torch.utils.cmake_prefix_path returned: ${torch_cmake_prefix}" >&2
        exit 1
    fi

    echo "Using Torch CMake prefix: ${torch_cmake_prefix}"
    echo "Using Torch_DIR: ${torch_dir}"

    # CUDA is intentionally opt-in.
    # Default to CPU-only Torch-JIT compilation unless CUDA_HOME is explicitly set.
    #
    # To enable CUDA, run for example:
    #   export CUDA_HOME=/usr/local/cuda-12.4
    #   ./CODES-compile-instructions.sh
    torch_cuda_version="$(python3 - <<'INNERPY'
import torch
print(torch.version.cuda or "")
INNERPY
)"

    cuda_arch=""
    if [ -z "${CUDA_HOME:-}" ] && [ -n "${torch_cuda_version}" ]; then
        echo "ERROR: CUDA_HOME is not set, so this script is defaulting to CPU-only Torch-JIT compilation." >&2
        echo "       However, the active Python environment has a CUDA-enabled PyTorch build:" >&2
        echo "       torch.version.cuda=${torch_cuda_version}" >&2
        echo "" >&2
        echo "       CMake cannot use a CUDA-enabled PyTorch package as a CPU-only LibTorch package." >&2
        echo "       Choose one of the following:" >&2
        echo "         1. For CPU-only compilation, install a CPU-only PyTorch build in this environment." >&2
        echo "         2. For CUDA compilation, export CUDA_HOME to your CUDA toolkit root." >&2
        echo "" >&2
        echo "       Example CUDA build:" >&2
        echo "         export CUDA_HOME=/usr/local/cuda-12.4" >&2
        echo "         bash CODES-compile-instructions.sh" >&2
        exit 1
    fi

    if [ -n "${CUDA_HOME:-}" ]; then
        if [ ! -f "${CUDA_HOME}/include/cuda_runtime_api.h" ]; then
            echo "ERROR: CUDA_HOME is set, but missing CUDA header: ${CUDA_HOME}/include/cuda_runtime_api.h" >&2
            exit 1
        fi

        if [ ! -f "${CUDA_HOME}/lib64/libcudart.so" ] && [ ! -f "${CUDA_HOME}/lib/libcudart.so" ]; then
            echo "ERROR: CUDA_HOME is set, but missing CUDA runtime library under ${CUDA_HOME}/lib64 or ${CUDA_HOME}/lib" >&2
            exit 1
        fi

        if [ ! -x "${CUDA_HOME}/bin/nvcc" ]; then
            echo "ERROR: CUDA_HOME is set, but missing CUDA compiler: ${CUDA_HOME}/bin/nvcc" >&2
            exit 1
        fi

        if [ ! -d "${CUDA_HOME}/nvvm/libdevice" ]; then
            echo "ERROR: CUDA_HOME is set, but missing CUDA libdevice directory: ${CUDA_HOME}/nvvm/libdevice" >&2
            exit 1
        fi

        if command -v nvidia-smi >/dev/null 2>&1; then
            cuda_arch="$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -n1 | tr -d '.[:space:]' || true)"
        fi

        if [ -z "${cuda_arch}" ]; then
            echo "WARNING: Could not auto-detect GPU compute capability with nvidia-smi." >&2
            echo "         Falling back to CMAKE_CUDA_ARCHITECTURES=80." >&2
            cuda_arch="80"
        fi

        export CUDA_HOME
        export CUDA_PATH="${CUDA_HOME}"
        export CUDA_ROOT="${CUDA_HOME}"
        export CUDA_TOOLKIT_ROOT_DIR="${CUDA_HOME}"
        export CUDAToolkit_ROOT="${CUDA_HOME}"
        export CUDACXX="${CUDA_HOME}/bin/nvcc"
        export PATH="${CUDA_HOME}/bin:${PATH}"
        export LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${CUDA_HOME}/lib:${LD_LIBRARY_PATH:-}"

        echo "CUDA_HOME is set; enabling CUDA Torch-JIT compilation."
        echo "Using CUDA_HOME: ${CUDA_HOME}"
        echo "Using CUDACXX: ${CUDACXX}"
        echo "Using CMAKE_CUDA_ARCHITECTURES=${cuda_arch}"
    else
        echo "CUDA_HOME is not set; forcing CPU-only Torch-JIT compilation."

        # Prevent accidental CUDA discovery from /usr/local/cuda, nvcc on PATH,
        # inherited CMake cache variables, or CUDA-enabled PyTorch metadata.
        unset CUDA_HOME
        unset CUDA_PATH
        unset CUDA_ROOT
        unset CUDA_TOOLKIT_ROOT_DIR
        unset CUDAToolkit_ROOT
        unset CUDACXX
        unset CMAKE_CUDA_COMPILER
    fi
fi

ross_config="$(find "$CUR_DIR/ross/build" \( -name ROSSConfig.cmake -o -name ross-config.cmake \) | head -n 1)"
if [ -z "$ross_config" ]; then
    echo "ERROR: Could not find built ROSSConfig.cmake under $CUR_DIR/ross/build." >&2
    echo "       Try rebuilding ROSS or check whether make install completed." >&2
    exit 1
fi

ross_dir="$(dirname "$ross_config")"
echo "Using ROSS_DIR=${ross_dir}"

cmake_prefix_path="${ross_dir}"
if [ "$torch_enable" = 1 ]; then
    cmake_prefix_path="${cmake_prefix_path};${torch_cmake_prefix}"
fi

make_args_codes=(
    -DROSS_DIR="${ross_dir}"
    -DCMAKE_PREFIX_PATH="${cmake_prefix_path}"
    -DCMAKE_C_FLAGS="-g -Wall"
    -DCMAKE_CXX_FLAGS="-g -Wall"
    -DTHREADS_PREFER_PTHREAD_FLAG=ON
    -DCMAKE_THREAD_LIBS_INIT="-pthread"
    -DCMAKE_HAVE_THREADS_LIBRARY=1
    -DCMAKE_USE_PTHREADS_INIT=1
    -DCMAKE_USE_WIN32_THREADS_INIT=0
    -DCMAKE_BUILD_TYPE=Debug -DBUILD_TESTING=ON
    -DCMAKE_INSTALL_PREFIX="$(realpath bin)"
)
# SWM/argobots/union are located via pkg-config. Their .pc files sit in the
# in-tree autotools build dirs (e.g. swm/build/maint), which are NOT under a
# standard <prefix>/lib/pkgconfig, so CMAKE_PREFIX_PATH can't reach them —
# expose them via PKG_CONFIG_PATH (the correct mechanism for that case, and it
# avoids the deprecated *_PKG_CONFIG_PATH cache vars). Mirrors the zmq handling
# above.
if [ $swm_enable = 1 ]; then
    export PKG_CONFIG_PATH="$(realpath "$CUR_DIR/swm-workloads/swm/build/maint"):$(realpath "$CUR_DIR/argobots/build/maint"):${PKG_CONFIG_PATH:-}"
fi
if [ $union_enable = 1 ]; then
    export PKG_CONFIG_PATH="$(realpath "$CUR_DIR/Union/install/lib/pkgconfig"):${PKG_CONFIG_PATH:-}"
fi
if [ "$torch_enable" = 1 ]; then
    make_args_codes=(
        "${make_args_codes[@]}"
        -DCODES_USE_TORCH=ON
        -DTorch_DIR="${torch_dir}"
    )

    if [ -n "${CUDA_HOME:-}" ]; then
        make_args_codes=(
            "${make_args_codes[@]}"
            -DCUDA_TOOLKIT_ROOT_DIR="${CUDA_HOME}"
            -DCUDAToolkit_ROOT="${CUDA_HOME}"
            -DCUDA_PATH="${CUDA_HOME}"
            -DCUDA_ROOT="${CUDA_HOME}"
            -DCMAKE_CUDA_COMPILER="${CUDA_HOME}/bin/nvcc"
            -DCMAKE_CUDA_ARCHITECTURES="${cuda_arch}"
            -DCUDA_INCLUDE_DIRS="${CUDA_HOME}/include"
            -DCUDA_CUDART_LIBRARY="${CUDA_HOME}/lib64/libcudart.so"
        )
    else
        make_args_codes=(
            "${make_args_codes[@]}"
            -DCMAKE_DISABLE_FIND_PACKAGE_CUDA=ON
            -DCMAKE_DISABLE_FIND_PACKAGE_CUDAToolkit=ON
        )
    fi
else
    make_args_codes=("${make_args_codes[@]}" -DCODES_USE_TORCH=OFF)
fi

cmake .. "${make_args_codes[@]}"
#make VERBOSE=1
make -j4
err=$?
[[ $err -ne 0 ]] && exit $err

popd