install-faiss-with-mkl.sh

#!/usr/bin/env bash

# The following assumes you have a pyenv virtualenv named faiss. Tested on
# Ubuntu 16.04.

set -eox pipefail

# shellcheck disable=SC1090
source "${HOME}/.pyenv.sh"

FAISS_BRANCH=${FAISS_BRANCH:-"master"}
USE_INTEL_MKL=1

NPROC=$(nproc)
CODE_PATH="${HOME}/code"
mkdir -p "${CODE_PATH}"

# sleep 30 # make sure Ubuntu is fully initialized

function use_intel_mkl() {
  [[ "${USE_INTEL_MKL}" = "1" ]]
}

function venv_site_packages() {
  pyenv shell faiss
  python -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())"
}

function is_cuda_compatible() {
  lspci | grep -i nvidia
}

function install_intel_mkl() {
  pushd /tmp
  wget http://registrationcenter-download.intel.com/akdlm/irc_nas/tec/11544/l_mkl_2017.3.196.tgz
  tar xzf l_mkl_2017.3.196.tgz
  pushd l_mkl_2017.3.196/
  # pre-generated with "sudo ./install.sh -d install.config"
  cat << EOF > ./install.config
ACCEPT_EULA=accept
CONTINUE_WITH_OPTIONAL_ERROR=yes
PSET_INSTALL_DIR=/opt/intel
CONTINUE_WITH_INSTALLDIR_OVERWRITE=yes
PSET_MODE=install
ARCH_SELECTED=ALL
COMPONENTS=;intel-comp-l-all-vars__noarch;intel-openmp-l-all__i486;intel-openmp-l-ps-libs-jp__i486;intel-openmp-l-all__x86_64;intel-openmp-l-ps-libs__x86_64;intel-openmp-l-ps-libs-jp__x86_64;intel-tbb-libs__noarch;intel-mkl-common__noarch;intel-mkl-sta-common__noarch;intel-mkl__i486;intel-mkl__x86_64;intel-mkl-rt__i486;intel-mkl-ps-rt-jp__i486;intel-mkl-rt__x86_64;intel-mkl-ps-rt-jp__x86_64;intel-mkl-doc__noarch;intel-mkl-ps-doc__noarch;intel-mkl-ps-doc-jp__noarch;intel-mkl-gnu__i486;intel-mkl-gnu__x86_64;intel-mkl-gnu-rt__i486;intel-mkl-gnu-rt__x86_64;intel-mkl-ps-common__noarch;intel-mkl-ps-common-jp__noarch;intel-mkl-ps-common__i486;intel-mkl-ps-common-64bit__x86_64;intel-mkl-common-c__noarch;intel-mkl-common-c__i486;intel-mkl-common-c-64bit__x86_64;intel-mkl-ps-common-c__noarch;intel-mkl-doc-c__noarch;intel-mkl-ps-doc-c-jp__noarch;intel-mkl-ps-ss-tbb__i486;intel-mkl-ps-ss-tbb__x86_64;intel-mkl-ps-ss-tbb-rt__i486;intel-mkl-ps-ss-tbb-rt__x86_64;intel-mkl-gnu-c__i486;intel-mkl-gnu-c__x86_64;intel-mkl-ps-common-f__noarch;intel-mkl-ps-common-f__i486;intel-mkl-ps-common-f-64bit__x86_64;intel-mkl-ps-doc-f__noarch;intel-mkl-ps-doc-f-jp__noarch;intel-mkl-ps-gnu-f-rt__i486;intel-mkl-ps-gnu-f-rt__x86_64;intel-mkl-ps-gnu-f__x86_64;intel-mkl-ps-gnu-f__i486;intel-mkl-ps-f95-common__noarch;intel-mkl-ps-f__i486;intel-mkl-ps-f__x86_64;intel-mkl-psxe__noarch;intel-psxe-common__noarch;intel-psxe-common-doc__noarch;intel-compxe-pset
EOF
  sudo ./install.sh -s install.config
  cat << EOF > ~/.intel-mkl.sh
export MKLROOT=/opt/intel/compilers_and_libraries/linux/mkl/
export LD_LIBRARY_PATH=\$MKLROOT/lib/intel64
export LD_PRELOAD=\$MKLROOT/lib/intel64/libmkl_core.so:\$MKLROOT/lib/intel64/libmkl_sequential.so
EOF
  echo '. ~/.intel-mkl.sh' >> "${HOME}/.bashrc"
  popd
  popd
}

function install_faiss() {
  pyenv shell faiss
  pip install numpy
  sudo apt-get install -y \
    build-essential \
    libopenblas-dev \
    gfortran \
    swig

  if use_intel_mkl; then
    install_intel_mkl
  fi

  PREFIX_MAIN=$(pyenv virtualenv-prefix)
  SITE_PACKAGES=$(venv_site_packages)
  PYTHON_NUMPY_INCLUDE_DIR="${SITE_PACKAGES}/numpy/core/include/"
  PYTHON_INCLUDE_DIR="$(echo "${PREFIX_MAIN}/include/python"*"m")/"

  # Download Facebook Faiss library source code
  cd "${CODE_PATH}"
  git clone --depth 1 https://github.com/facebookresearch/faiss.git -b "${FAISS_BRANCH}"
  cd "${CODE_PATH}/faiss"
  # Setup makefile
  cp example_makefiles/makefile.inc.Linux makefile.inc

  if use_intel_mkl; then
    cat << EOF >> makefile.inc
MKLROOT=/opt/intel/compilers_and_libraries/linux/mkl/

BLASLDFLAGS=-Wl,--no-as-needed -L\$(MKLROOT)/lib/intel64 \\
-lmkl_intel_ilp64 -lmkl_core -lmkl_gnu_thread -ldl -lpthread

BLASCFLAGS=-DFINTEGER=long

BLASLDFLAGSNVCC=-Xlinker -L\$(MKLROOT)/lib/intel64  -Xlinker -lmkl_intel_ilp64 \\
-Xlinker -lmkl_core -Xlinker -lmkl_gnu_thread -Xlinker -ldl -Xlinker -lpthread

BLASLDFLAGSSONVCC=\$(BLASLDFLAGSNVCC)
EOF
  else
    export BLASLDFLAGS="/usr/lib/libopenblas.so.0"
  fi

  echo "PYTHONCFLAGS=-I${PYTHON_INCLUDE_DIR} -I${PYTHON_NUMPY_INCLUDE_DIR}" \
    >> makefile.inc

  # Compile faiss
  make -j "${NPROC}"
  # Create python bindings
  make py

  if is_cuda_compatible; then
    # Compile GPU version
    cd "${CODE_PATH}/faiss/gpu"
    make -j "${NPROC}"
    make py
    cd "${CODE_PATH}/faiss"
  else
    echo "No NVIDIA GPU detected, skipping faiss gpu install"
  fi

  # Copy python modules to python system path
  cp -- *.py *.so "$(venv_site_packages)/"
}

install_faiss