configure

#!/bin/bash
set -e

cd $(dirname $0)
. ./scripts/utils/measure.sh

if [ -z "${BUILD_GTEST}" ]; then
    BUILD_GTEST=0 # download and build gtest from source
fi
BUILD_TENSORFLOW_OPS=0
BUILD_TESTS=0 # build unit tests
BUILD_TOOLS=0
ENABLE_MPI=0
ENABLE_NCCL=0
ENABLE_TRACE=0
VERBOSE=0

# defaults
NCCL_HOME=$HOME/local/nccl
MPI_HOME=$HOME/local/openmpi

parse_args() {
    for i in "$@"; do
        case $i in
        --prefix=*)
            PREFIX="${i#*=}"
            ;;
        --build-tensorflow-ops)
            BUILD_TENSORFLOW_OPS=1
            ;;
        --build-tools)
            BUILD_TOOLS=1
            ;;
        --build-gtest)
            BUILD_GTEST=1
            ;;
        --build-tests)
            BUILD_TESTS=1
            ;;
        --no-tests)
            BUILD_TESTS=0
            ;;
        --disable-cxx11-abi)
            DISABLE_CXX11_ABI=1
            ;;
        --enable-trace)
            ENABLE_TRACE=1
            ;;
        --enable-mpi)
            ENABLE_MPI=1
            ;;
        --enable-nccl)
            ENABLE_NCCL=1
            ;;
        --with-mpi=*)
            MPI_HOME="${i#*=}"
            ENABLE_MPI=1
            ;;
        --with-nccl=*)
            NCCL_HOME="${i#*=}"
            ENABLE_NCCL=1
            ;;
        --verbose)
            VERBOSE=1
            ;;
        *)
            echo "unknown argument $i"
            exit 1
            ;;
        esac
    done
}

CMAKE_FLAGS=

add_cmake_flag() {
    echo "using $1=$2"
    CMAKE_FLAGS="$CMAKE_FLAGS -D$1=$2"
}

add_cmake_flags() {
    if [ ${BUILD_GTEST} = 1 ] && [ ! -z ${GTEST_GIT_URL} ]; then
        add_cmake_flag GTEST_GIT_URL ${GTEST_GIT_URL}
    fi

    if [ ! -z ${ENABLE_TRACE} ]; then
        add_cmake_flag KUNGFU_ENABLE_TRACE ${ENABLE_TRACE}
        if [ ! -z ${STDTRACER_GIT_URL} ]; then
            add_cmake_flag STDTRACER_GIT_URL ${STDTRACER_GIT_URL}
        fi
    fi

    if [ ${BUILD_TENSORFLOW_OPS} = 1 ]; then
        if [ -z "${PYTHON}" ]; then
            PYTHON=$(which python3)
        fi
        add_cmake_flag PYTHON $PYTHON
    fi

    add_cmake_flag CMAKE_PREFIX_PATH "$HOME/local;$HOME/local/gtest"
    add_cmake_flag CMAKE_BUILD_TYPE Release
    add_cmake_flag LIBRARY_OUTPUT_PATH $(pwd)/lib
    add_cmake_flag CMAKE_RUNTIME_OUTPUT_DIRECTORY $(pwd)/bin

    add_cmake_flag CMAKE_VERBOSE_MAKEFILE ${VERBOSE}
    add_cmake_flag CMAKE_EXPORT_COMPILE_COMMANDS ${VERBOSE}

    add_cmake_flag KUNGFU_BUILD_GTEST ${BUILD_GTEST}
    add_cmake_flag KUNGFU_BUILD_TESTS ${BUILD_TESTS}
    add_cmake_flag KUNGFU_BUILD_TF_OPS ${BUILD_TENSORFLOW_OPS}
    add_cmake_flag KUNGFU_BUILD_TOOLS ${BUILD_TOOLS}

    add_cmake_flag DISABLE_CXX11_ABI $DISABLE_CXX11_ABI

    if [ ! -z ${PREFIX} ]; then
        add_cmake_flag CMAKE_INSTALL_PREFIX ${PREFIX}
    fi
    if [ ${ENABLE_MPI} -gt 0 ]; then
        add_cmake_flag MPI_HOME ${MPI_HOME}
    fi
    if [ ${ENABLE_NCCL} -gt 0 ]; then
        add_cmake_flag NCCL_HOME ${NCCL_HOME}
        add_cmake_flag KUNGFU_ENABLE_NCCL 1
    fi

    if [ $(uname) = "Linux" ]; then 
        add_cmake_flag KUNGFU_ENABLE_AFFINITY 1
    fi
}

cleanup_cmake_cache() {
    if [ -f CMakeCache.txt ]; then
        rm CMakeCache.txt
    fi
}

main() {
    parse_args $@
    add_cmake_flags
    cleanup_cmake_cache
    cmake . $CMAKE_FLAGS
}

measure main $@