Upgrade faiss to v1.9.0 (#113)

kyamagu · Oct 8, 2024 · 9bcced8 · 9bcced8
1 parent 4fb2285
commit 9bcced8
Show file tree

Hide file tree

Showing 8 changed files with 103 additions and 57 deletions.
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -68,13 +68,20 @@ jobs:
           python-version: "3.x"
 
       - name: Set up QEMU
-        if: runner.os == 'Linux'
+        if: runner.os == 'Linux' && matrix.arch == 'aarch64'
         uses: docker/setup-qemu-action@v3
         with:
           platforms: arm64
 
+      - name: Set up conda
+        if: runner.os == 'Windows'
+        uses: conda-incubator/setup-miniconda@v3
+        with:
+          auto-activate-base: true
+          activate-environment: ""
+
       - name: Install cibuildwheel
-        run: python -m pip install cibuildwheel==2.16.5
+        run: python -m pip install cibuildwheel==2.21.1
 
       - name: Build wheels
         run: python -m cibuildwheel --output-dir wheelhouse
@@ -89,18 +96,17 @@ jobs:
             LIBOMP_USE_HIDDEN_HELPER_TASK=0
             LIBOMP_NUM_HIDDEN_HELPER_THREADS=0
             KMP_DUPLICATE_LIB_OK=TRUE
+            OMP_NUM_THREADS=1
             MACOSX_DEPLOYMENT_TARGET=10.14
           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
           CIBW_BEFORE_TEST_LINUX: pip install torch --index-url https://download.pytorch.org/whl/cpu
           CIBW_TEST_COMMAND: >
-            env -u FAISS_OPT_LEVEL pytest {project}/faiss/tests &&
-            env -u FAISS_OPT_LEVEL pytest -s {project}/faiss/tests/torch_test_contrib.py
+            env -u FAISS_OPT_LEVEL pytest {project}/faiss/tests
 
       - uses: actions/upload-artifact@v4
         with:

diff --git a/faiss b/faiss
diff --git a/patch/faiss-remove-lapack.patch b/patch/faiss-remove-lapack.patch
@@ -1,16 +1,17 @@
 diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt
-index 22dac7cb..e330cf36 100644
+index 2871d974..f823348c 100644
 --- a/faiss/CMakeLists.txt
 +++ b/faiss/CMakeLists.txt
-@@ -307,11 +307,6 @@ else()
-   target_link_libraries(faiss PRIVATE ${BLAS_LIBRARIES})
+@@ -365,12 +365,6 @@ else()
    target_link_libraries(faiss_avx2 PRIVATE ${BLAS_LIBRARIES})
    target_link_libraries(faiss_avx512 PRIVATE ${BLAS_LIBRARIES})
+   target_link_libraries(faiss_sve PRIVATE ${BLAS_LIBRARIES})
 -
 -  find_package(LAPACK REQUIRED)
 -  target_link_libraries(faiss PRIVATE ${LAPACK_LIBRARIES})
 -  target_link_libraries(faiss_avx2 PRIVATE ${LAPACK_LIBRARIES})
 -  target_link_libraries(faiss_avx512 PRIVATE ${LAPACK_LIBRARIES})
+-  target_link_libraries(faiss_sve PRIVATE ${LAPACK_LIBRARIES})
  endif()
 
  install(TARGETS faiss
diff --git a/pyproject.toml b/pyproject.toml
@@ -2,19 +2,19 @@
 requires = [
     "setuptools",
     "wheel",
-    "oldest-supported-numpy",
+    "numpy>=2.0,<3",
 ]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "faiss-cpu"
-version = "1.8.0.post1"
+version = "1.9.0"
 authors = [
     { name = "Kota Yamaguchi", email = "[email protected]" },
 ]
 description = "A library for efficient similarity search and clustering of dense vectors."
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 keywords = ["faiss", "similarity search", "clustering", "machine learning"]
 license = { text = "MIT License" }
 classifiers = [
@@ -25,14 +25,13 @@ classifiers = [
     "Operating System :: MacOS :: MacOS X",
     "Operating System :: Microsoft :: Windows",
     "Operating System :: POSIX",
-    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
-dependencies = ["numpy>=1.0,<2.0", "packaging"]
+dependencies = ["numpy>=1.25.0,<3.0", "packaging"]
 
 [project.urls]
 Repository = "https://github.com/kyamagu/faiss-wheels"
@@ -42,12 +41,16 @@ minversion = "6.0"
 addopts = [
     # Failing tests for numerical issues.
     "--deselect=faiss/tests/test_local_search_quantizer.py::TestProductLocalSearchQuantizer::test_lut",
+    "--deselect=faiss/tests/test_local_search_quantizer.py::TestComponents::test_update_codebooks_with_double",
+    "--deselect=faiss/tests/test_contrib.py::TestPreassigned::test_float",
+    "--deselect=faiss/tests/test_contrib.py::TestClustering::test_ivf_train_2level",
 ]
 testpaths = ["faiss/tests"]
+python_files = "test_*.py"  # torch_test_*.py will break the other tests
 
 [tool.cibuildwheel]
-skip = "pp* *-musllinux*"
-test-skip = "*-macosx_arm64 *-manylinux_aarch64 *-win_amd64 cp312-*"
+skip = "pp* cp38-* cp313-* *-musllinux*"
+test-skip = "*-manylinux_aarch64 *-win_amd64"
 
 test-requires = ["pytest", "scipy", "torch"]
 test-command = "pytest {project}/faiss/tests"

diff --git a/scripts/build_Linux.sh b/scripts/build_Linux.sh
@@ -40,8 +40,15 @@ if [[ ${FAISS_ENABLE_GPU} == "ON" ]]; then
     export CUDAFLAGS="--compiler-options=${CXXFLAGS// /,}"
 fi
 
+# Check if swig is available
+if ! command -v swig &> /dev/null; then
+    echo "swig is not available. Please install swig."
+    exit 1
+fi
+swig -version
+
 # Install system dependencies
-yum install -y openblas-devel openblas-static swig3
+yum install -y openblas-devel openblas-static
 
 # Build and patch faiss
 cd faiss && \

diff --git a/scripts/build_Windows.sh b/scripts/build_Windows.sh
@@ -4,8 +4,11 @@ set -eux
 
 CMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH:-"c:\\opt"}
 
+# Show env
+env
+
 # Install system dependencies
-conda.bat install -c conda-forge openblas
+conda install -y -c conda-forge openblas
 
 # Build and patch faiss
 cd faiss && \

diff --git a/scripts/build_macOS.sh b/scripts/build_macOS.sh
@@ -21,25 +21,29 @@ sudo mkdir -p /usr/local/share && \
 brew install swig
 
 # Build libomp: needed for cross compilation
-echo "Building libomp"
-git clone \
-        --depth 1 \
-        --filter=blob:none \
-        --sparse \
-        --branch ${LLVM_VERSION:-"llvmorg-17.0.6"} \
-        https://github.com/llvm/llvm-project.git \
-        third-party/llvm-project && \
-    cd third-party/llvm-project && \
-    git sparse-checkout set openmp cmake && \
-    cd openmp && \
-    cmake . \
-        -B build \
-        -DCMAKE_BUILD_TYPE=Release \
-        -DCMAKE_OSX_ARCHITECTURES=${TARGET_ARCH} \
-        -DLIBOMP_ENABLE_SHARED=ON && \
-    cmake --build build -j && \
-    cmake --install build && \
-    cd ../../..
+function build_libomp() {
+    echo "Building libomp"
+    git clone \
+            --depth 1 \
+            --filter=blob:none \
+            --sparse \
+            --branch ${LLVM_VERSION:-"llvmorg-19.1.1"} \
+            https://github.com/llvm/llvm-project.git \
+            third-party/llvm-project && \
+        cd third-party/llvm-project && \
+        git sparse-checkout set openmp cmake && \
+        cd openmp && \
+        cmake . \
+            -B build \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DCMAKE_OSX_ARCHITECTURES=${TARGET_ARCH} \
+            -DLIBOMP_ENABLE_SHARED=ON && \
+        cmake --build build -j && \
+        cmake --install build && \
+        cd ../../..
+}
+
+build_libomp
 
 # Build and patch faiss
 echo "Building faiss"

diff --git a/setup.py b/setup.py
@@ -9,19 +9,23 @@
 # Faiss conifgurations
 FAISS_INSTALL_PREFIX = os.getenv("FAISS_INSTALL_PREFIX", "/usr/local")
 FAISS_OPT_LEVEL = os.getenv("FAISS_OPT_LEVEL", "generic")
-FAISS_ENABLE_GPU = (os.getenv("FAISS_ENABLE_GPU", "").lower() in ("on", "true"))
+FAISS_ENABLE_GPU = os.getenv("FAISS_ENABLE_GPU", "").lower() in ("on", "true")
 
 # Common configurations
 FAISS_ROOT = "faiss"  # relative to the setup.py file
 
-DEFINE_MACROS = [("FINTEGER", "int")]
+DEFINE_MACROS = []
 INCLUDE_DIRS = [
-    np.get_include(), FAISS_ROOT, os.path.join(FAISS_INSTALL_PREFIX, "include")]
+    np.get_include(),
+    FAISS_ROOT,
+    os.path.join(FAISS_INSTALL_PREFIX, "include"),
+]
 LIBRARY_DIRS: List[str] = [os.path.join(FAISS_INSTALL_PREFIX, "lib")]
 EXTRA_COMPILE_ARGS: List[str] = []
 EXTRA_LINK_ARGS: List[str] = []
 SWIG_OPTS = ["-c++", "-Doverride=", "-doxygen", f"-I{FAISS_ROOT}"] + [
-    f"-I{x}" for x in INCLUDE_DIRS]
+    f"-I{x}" for x in INCLUDE_DIRS
+]
 
 # GPU options
 if FAISS_ENABLE_GPU:
@@ -39,15 +43,16 @@ def win32_options(
 ) -> dict:
     """Windows options."""
     return dict(
-        extra_compile_args=extra_compile_args + [
+        extra_compile_args=extra_compile_args
+        + [
             "/openmp:llvm",
             "/std:c++17",
             "/Zc:inline",
             "/wd4101",  # unreferenced local variable.
             "/MD",  # Bugfix: https://bugs.python.org/issue38597
         ],
-        extra_link_args=["/OPT:ICF", "/OPT:REF"] + (
-            extra_link_args or ["faiss.lib", "openblas.lib"]),
+        extra_link_args=["/OPT:ICF", "/OPT:REF"]
+        + (extra_link_args or ["faiss.lib", "openblas.lib"]),
         swig_opts=swig_opts + ["-DSWIGWIN"],
     )
 
@@ -71,7 +76,8 @@ def linux_options(
             "-lculibos",
         ]
     return dict(
-        extra_compile_args=extra_compile_args + [
+        extra_compile_args=extra_compile_args
+        + [
             "-std=c++17",
             "-Wno-sign-compare",
             "-fopenmp",
@@ -83,7 +89,8 @@ def linux_options(
             "-lrt",
             "-s",
             "-Wl,--gc-sections",
-        ] + (extra_link_args or default_link_args),
+        ]
+        + (extra_link_args or default_link_args),
         swig_opts=swig_opts + ["-DSWIGWORDSIZE64"],
     )
 
@@ -95,7 +102,8 @@ def darwin_options(
 ) -> dict:
     """macOS options."""
     return dict(
-        extra_compile_args=extra_compile_args + [
+        extra_compile_args=extra_compile_args
+        + [
             "-std=c++17",
             "-Wno-sign-compare",
             "-Xpreprocessor",
@@ -105,12 +113,16 @@ def darwin_options(
             "-Xpreprocessor",
             "-fopenmp",
             "-dead_strip",
-        ] + (extra_link_args or [
-            "-lfaiss",
-            "-lomp",
-            "-framework",
-            "Accelerate",
-        ]),
+        ]
+        + (
+            extra_link_args
+            or [
+                "-lfaiss",
+                "-lomp",
+                "-framework",
+                "Accelerate",
+            ]
+        ),
         swig_opts=swig_opts,
     )
 
@@ -165,8 +177,15 @@ def avx512_options(
         flags = ["/arch:AVX512", "/bigobj"]
     else:
         flags = [
-            "-mavx2", "-mfma", "-mf16c", "-mavx512f", "-mavx512cd", "-mavx512vl",
-            "-mavx512dq", "-mavx512bw", "-mpopcnt"
+            "-mavx2",
+            "-mfma",
+            "-mf16c",
+            "-mavx512f",
+            "-mavx512cd",
+            "-mavx512vl",
+            "-mavx512dq",
+            "-mavx512bw",
+            "-mpopcnt",
         ]
     return dict(
         name="faiss._swigfaiss_avx512",
@@ -183,7 +202,8 @@ def avx512_options(
 }
 
 platform_config = PLATFORM_CONFIGS[sys.platform](
-    EXTRA_COMPILE_ARGS, EXTRA_LINK_ARGS, SWIG_OPTS)
+    EXTRA_COMPILE_ARGS, EXTRA_LINK_ARGS, SWIG_OPTS
+)
 
 ext_modules = [
     Extension(
@@ -204,16 +224,18 @@ def avx512_options(
 
 class CustomBuildPy(build_py):
     """Run build_ext before build_py to compile swig code."""
+
     def run(self):
         self.run_command("build_ext")
         return build_py.run(self)
 
 
 setup(
-    packages=["faiss", "faiss.contrib"],
+    packages=["faiss", "faiss.contrib", "faiss.contrib.torch"],
     package_dir={
         "faiss": os.path.join(FAISS_ROOT, "faiss", "python"),
         "faiss.contrib": os.path.join(FAISS_ROOT, "contrib"),
+        "faiss.contrib.torch": os.path.join(FAISS_ROOT, "contrib", "torch"),
     },
     include_package_data=False,
     package_data={"": ["*.i", "*.h"]},