Support AVX2 extension (#65)

Support swigfaiss_avx2 extension in x86_64 arch
kyamagu · Nov 17, 2022 · 7b57ac6 · 7b57ac6
1 parent 4bc4cd6
commit 7b57ac6
Show file tree

Hide file tree

Showing 7 changed files with 83 additions and 91 deletions.
diff --git 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,50 +37,60 @@ 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
 
       - name: Build wheels
         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
@@ -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,60 +63,26 @@ 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.
 
 ```bash
 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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',