diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 39e4d6d..21d1644 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -14,20 +14,21 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 with: submodules: true - - uses: actions/setup-python@v2 + - uses: actions/setup-python@v3 with: python-version: '3.x' + - name: Apply patch + run: cd faiss && git apply ../patch/faiss-rename-swigfaiss.patch && cd .. + - name: Build sdist - run: | - mv faiss/faiss/python/swigfaiss.swig faiss/faiss/python/swigfaiss.i - python setup.py sdist + run: python setup.py sdist - - uses: actions/upload-artifact@v2 + - uses: actions/upload-artifact@v3 with: path: ./dist/*.tar.gz @@ -36,28 +37,36 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - os: [ubuntu-latest, windows-latest, macos-latest] + os: [ubuntu-latest, macos-latest] arch: [auto64] gpu: [OFF] + opt_level: [avx2] include: + - os: windows-latest + arch: auto64 + gpu: OFF + opt_level: generic - os: ubuntu-latest arch: aarch64 gpu: OFF + opt_level: generic - os: macos-latest arch: arm64 gpu: OFF + opt_level: generic - os: ubuntu-latest arch: auto64 gpu: ON + opt_level: avx2 steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v3 with: submodules: true - name: Set up QEMU if: runner.os == 'Linux' && matrix.arch != 'auto64' - uses: docker/setup-qemu-action@v1 + uses: docker/setup-qemu-action@v2 with: platforms: arm64 @@ -65,21 +74,23 @@ jobs: uses: pypa/cibuildwheel@v2.11.2 env: CIBW_ARCHS: ${{ matrix.arch }} - CIBW_ENVIRONMENT: > - FAISS_OPT_LEVEL=generic + CIBW_ENVIRONMENT_LINUX: > + FAISS_OPT_LEVEL=${{ matrix.opt_level }} FAISS_ENABLE_GPU=${{ matrix.gpu }} CIBW_ENVIRONMENT_MACOS: > + FAISS_OPT_LEVEL=${{ matrix.opt_level }} TARGET_ARCH=${{ matrix.arch }} LIBOMP_USE_HIDDEN_HELPER_TASK=0 LIBOMP_NUM_HIDDEN_HELPER_THREADS=0 CIBW_ENVIRONMENT_WINDOWS: > + FAISS_OPT_LEVEL=${{ matrix.opt_level }} CMAKE_PREFIX_PATH="c:\\opt" PATH="${PATH};${CONDA}\\condabin;${CONDA}\\Library\\bin" LIB="${LIB};${CMAKE_PREFIX_PATH}\\lib;${CONDA}\\Library\\lib" CPATH="${CPATH};${CMAKE_PREFIX_PATH}\\include;${CONDA}\\Library\\include" CIBW_BEFORE_ALL: bash scripts/build_${{ runner.os }}.sh - - uses: actions/upload-artifact@v2 + - uses: actions/upload-artifact@v3 with: path: ./wheelhouse/*.whl @@ -89,7 +100,7 @@ jobs: runs-on: ubuntu-latest if: github.event_name == 'release' && github.event.action == 'published' steps: - - uses: actions/download-artifact@v2 + - uses: actions/download-artifact@v3 with: name: artifact path: dist diff --git a/README.md b/README.md index 000e4f0..60dd61c 100644 --- a/README.md +++ b/README.md @@ -10,24 +10,16 @@ faiss python wheel packages. ## Overview -This repository provides scripts to create wheel packages for the +This repository provides scripts to build wheel packages for the [faiss](https://github.com/facebookresearch/faiss) library. - Builds CPU-only or CUDA-11.0+ compatible wheels with [cibuildwheel](https://github.com/pypa/cibuildwheel/). - Bundles OpenBLAS in Linux/Windows -- Uses Accelerate framework on macOS -- CUDA runtime and cuBLAS are statically linked +- Uses Accelerate framework in macOS +- Statically linked CUDA runtime There is also a source package to customize the build process. -### Prerequisite - -On macOS, install `libomp` via Homebrew to use the wheel. - -```bash -brew install libomp -``` - ### Install Install CPU-only version: @@ -59,9 +51,9 @@ Build and install the faiss library first. ```bash cd faiss -cmake -B build . -DFAISS_ENABLE_PYTHON=OFF -make -C build -j8 -make -C build install +cmake . -B build -DFAISS_ENABLE_GPU=OFF -DFAISS_ENABLE_PYTHON=OFF -DFAISS_OPT_LEVEL=avx2 +cmake --build build --config Release -j +cmake --install build install cd .. ``` @@ -71,7 +63,7 @@ for more on how to build and install faiss. For building sdist, swig 3.0.12 or later needs to be available. -### Linux +### Building a wheel package By default, the following builds and installs the faiss-cpu package. @@ -79,52 +71,18 @@ By default, the following builds and installs the faiss-cpu package. pip install --no-binary :all: faiss-cpu ``` -The following example shows static linking and CUDA support: +The following example builds a GPU wheel. ```bash export FAISS_ENABLE_GPU=ON -export FAISS_LDFLAGS='-l:libfaiss.a -l:libopenblas.a -lgfortran -lcudart_static -lcublas_static -lculibos' pip install --no-binary :all: faiss-gpu ``` -There are a few environment variables to specify build-time options. +There are a few environment variables that specifies build-time options. -- `CUDA_HOME`: Specifies CUDA install location. -- `FAISS_INCLUDE`: Header locations of the installed faiss library. Default to - `/usr/local/include`. -- `FAISS_LDFLAGS`: Linker flags for package build. Default to - `-l:libfaiss.a -l:libopenblas.a -lgfortran`. -- `FAISS_OPT_LEVEL`: Faiss SIMD optimization, one of `generic`, `avx2`. +- `CUDA_HOME`: Specifies CUDA install location for building faiss-gpu package. +- `FAISS_OPT_LEVEL`: Faiss SIMD optimization, one of `generic`, `avx2`. When set + to `avx2`, the package internally builds `avx2` extension in addition to + `generic`. Note this option is only available in x86_64 arch. - `FAISS_ENABLE_GPU`: Setting this variable to `ON` builds `faiss-gpu` package. Set this variable if faiss is built with GPU support. - -Below is an example for faiss built with `avx2` option and OpenBLAS backend. - -```bash -export FAISS_OPT_LEVEL='avx2' -export FAISS_LDFLAGS='-l:libfaiss_avx2.a -l:libopenblas.a -lgfortran' -pip install --no-binary :all: faiss-cpu -``` - -### macOS - -On macOS, install `libomp` via Homebrew to build with OpenMP support. Mac has -Accelerate framework for BLAS implementation. CUDA is not supported on macOS. - -```bash -pip install --no-binary :all: faiss-cpu -``` - -To link to faiss library with `avx2` support, set appropriate environment -variables. - -```bash -export FAISS_OPT_LEVEL=avx2 -export FAISS_LDFLAGS="/usr/local/lib/libfaiss_avx2.a /usr/local/lib/libomp.a -framework Accelerate" -pip install --no-binary :all: faiss-cpu -``` - -### Windows - -Windows environment requires BLAS/LAPACK and fortran library. See -`.github/workflows/build.yml` for how the binary wheels are built. diff --git a/patch/faiss-rename-swigfaiss.patch b/patch/faiss-rename-swigfaiss.patch new file mode 100644 index 0000000..3442aeb --- /dev/null +++ b/patch/faiss-rename-swigfaiss.patch @@ -0,0 +1,4 @@ +diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.i +similarity index 100% +rename from faiss/python/swigfaiss.swig +rename to faiss/python/swigfaiss.i diff --git a/scripts/build_Linux.sh b/scripts/build_Linux.sh index bd5a6aa..b63bd13 100644 --- a/scripts/build_Linux.sh +++ b/scripts/build_Linux.sh @@ -52,11 +52,11 @@ cd faiss && \ -B build \ -DFAISS_ENABLE_GPU=${FAISS_ENABLE_GPU} \ -DFAISS_ENABLE_PYTHON=OFF \ - -DBUILD_TESTING=ON \ + -DBUILD_TESTING=OFF \ -DCMAKE_CUDA_ARCHITECTURES=${CMAKE_CUDA_ARCHITECTURES} \ -DFAISS_OPT_LEVEL=${FAISS_OPT_LEVEL} \ -DCMAKE_BUILD_TYPE=Release && \ - cmake --build build --config Release -j2 && \ + cmake --build build --config Release -j3 && \ cmake --install build && \ - mv faiss/python/swigfaiss.swig faiss/python/swigfaiss.i && \ + git apply ../patch/faiss-rename-swigfaiss.patch && \ cd .. diff --git a/scripts/build_Windows.sh b/scripts/build_Windows.sh index e146abd..f08e1ee 100644 --- a/scripts/build_Windows.sh +++ b/scripts/build_Windows.sh @@ -21,5 +21,5 @@ cd faiss && \ -DBLA_STATIC=ON && \ cmake --build build --config Release -j && \ cmake --install build --prefix ${CMAKE_PREFIX_PATH} && \ - mv faiss/python/swigfaiss.swig faiss/python/swigfaiss.i && \ + git apply ../patch/faiss-rename-swigfaiss.patch && \ cd .. diff --git a/scripts/build_macOS.sh b/scripts/build_macOS.sh index 25a4365..a3772ff 100644 --- a/scripts/build_macOS.sh +++ b/scripts/build_macOS.sh @@ -44,7 +44,7 @@ cd faiss && \ -DFAISS_OPT_LEVEL=${FAISS_OPT_LEVEL} \ -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_OSX_ARCHITECTURES=${TARGET_ARCH} && \ - cmake --build build --config Release -j -v && \ + cmake --build build --config Release -j && \ cmake --install build && \ - mv faiss/python/swigfaiss.swig faiss/python/swigfaiss.i && \ + git apply ../patch/faiss-rename-swigfaiss.patch && \ cd .. diff --git a/setup.py b/setup.py index 0893458..a7d7ba5 100644 --- a/setup.py +++ b/setup.py @@ -1,8 +1,9 @@ +import os +import sys + from setuptools import setup -from setuptools.extension import Extension from setuptools.command.build_py import build_py -import sys -import os +from setuptools.extension import Extension NAME = 'faiss-cpu' VERSION = '1.7.3' @@ -133,12 +134,6 @@ def __str__(self): LIBRARY_DIRS += [os.path.join(CUDA_HOME, 'lib64')] SWIG_OPTS += ['-I' + os.path.join(CUDA_HOME, 'include'), '-DGPU_WRAPPER'] -if FAISS_OPT_LEVEL == 'avx2': - if sys.platform == 'win32': - EXTRA_COMPILE_ARGS += ['/arch:AVX2'] - else: - EXTRA_COMPILE_ARGS += ['-mavx2', '-mpopcnt'] - class CustomBuildPy(build_py): """Run build_ext before build_py to compile swig code.""" @@ -153,17 +148,43 @@ def run(self): os.path.join(FAISS_ROOT, 'faiss', 'python', 'swigfaiss.i'), os.path.join(FAISS_ROOT, 'faiss', 'python', 'python_callbacks.cpp'), ], - depends=[ - os.path.join(FAISS_ROOT, 'faiss', 'python', 'python_callbacks.h'), - ], + depends=[os.path.join(FAISS_ROOT, 'faiss', 'python', 'python_callbacks.h')], language='c++', define_macros=DEFINE_MACROS, include_dirs=INCLUDE_DIRS, library_dirs=LIBRARY_DIRS, extra_compile_args=EXTRA_COMPILE_ARGS, extra_link_args=EXTRA_LINK_ARGS, - swig_opts=SWIG_OPTS, + swig_opts=SWIG_OPTS + ["-module", "swigfaiss"], ) +ext_modules = [_swigfaiss] + +if FAISS_OPT_LEVEL == 'avx2': + # NOTE: avx2 is only available on x86_64 arch. + if sys.platform == 'win32': + EXTRA_COMPILE_ARGS_AVX2 = EXTRA_COMPILE_ARGS + ['/arch:AVX2'] + else: + EXTRA_COMPILE_ARGS_AVX2 = EXTRA_COMPILE_ARGS + ['-mavx2', '-mpopcnt'] + + # TODO: fix this ad-hoc approach to specify avx2 lib. + EXTRA_LINK_ARGS_AVX2 = [x.replace("faiss", "faiss_avx2") for x in EXTRA_LINK_ARGS] + + _swigfaiss_avx2 = Extension( + 'faiss._swigfaiss_avx2', + sources=[ + os.path.join(FAISS_ROOT, 'faiss', 'python', 'swigfaiss.i'), + os.path.join(FAISS_ROOT, 'faiss', 'python', 'python_callbacks.cpp'), + ], + depends=[os.path.join(FAISS_ROOT, 'faiss', 'python', 'python_callbacks.h')], + language='c++', + define_macros=DEFINE_MACROS, + include_dirs=INCLUDE_DIRS, + library_dirs=LIBRARY_DIRS, + extra_compile_args=EXTRA_COMPILE_ARGS_AVX2, + extra_link_args=EXTRA_LINK_ARGS_AVX2, + swig_opts=SWIG_OPTS + ["-module", "swigfaiss_avx2"], + ) + ext_modules.append(_swigfaiss_avx2) setup( name=NAME, @@ -184,10 +205,8 @@ def run(self): 'faiss': os.path.join(FAISS_ROOT, 'faiss', 'python'), 'faiss.contrib': os.path.join(FAISS_ROOT, 'contrib'), }, - package_data={ - 'faiss': ['*.i', '*.h'], - }, - ext_modules=[_swigfaiss], + package_data={'faiss': ['*.i', '*.h']}, + ext_modules=ext_modules, cmdclass={'build_py': CustomBuildPy}, classifiers=[ 'Development Status :: 4 - Beta',