avx512 CI + conda packages (#3197)

Summary: Pull Request resolved: #3197 Reviewed By: mlomeli1 Differential Revision: D52689379 Pulled By: algoriddle fbshipit-source-id: 54e27c6d310d6da14777ae10ae62f46e7076cacf
facebookresearch · Jan 11, 2024 · 0013c70 · 0013c70
1 parent 0710cbd
commit 0013c70
Show file tree

Hide file tree

Showing 17 changed files with 43 additions and 35 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -314,6 +314,10 @@ workflows:
           name: Linux x86_64 AVX2 (cmake)
           exec: linux-x86_64-cpu
           opt_level: "avx2"
+      - build_cmake:
+          name: Linux x86_64 AVX512 (cmake)
+          exec: linux-x86_64-cpu
+          opt_level: "avx512"
       - build_cmake:
           name: Linux x86_64 GPU (cmake)
           exec: linux-x86_64-gpu

diff --git a/.gitignore b/.gitignore
@@ -17,3 +17,4 @@
 /tests/test
 /tests/gtest/
 faiss/python/swigfaiss_avx2.swig
+faiss/python/swigfaiss_avx512.swig
diff --git a/conda/faiss-gpu-raft/build-lib.sh b/conda/faiss-gpu-raft/build-lib.sh
@@ -7,11 +7,11 @@
 set -e
 
 
-# Build libfaiss.so/libfaiss_avx2.so.
+# Build libfaiss.so/libfaiss_avx2.so/libfaiss_avx512.so
 cmake -B _build \
       -DBUILD_SHARED_LIBS=ON \
       -DBUILD_TESTING=OFF \
-      -DFAISS_OPT_LEVEL=avx2 \
+      -DFAISS_OPT_LEVEL=avx512 \
       -DFAISS_ENABLE_GPU=ON \
       -DFAISS_ENABLE_RAFT=ON \
       -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHS}" \
@@ -20,7 +20,7 @@ cmake -B _build \
       -DCMAKE_INSTALL_LIBDIR=lib \
       -DCMAKE_BUILD_TYPE=Release .
 
-make -C _build -j$(nproc) faiss faiss_avx2
+make -C _build -j$(nproc) faiss faiss_avx2 faiss_avx512
 
 cmake --install _build --prefix $PREFIX
 cmake --install _build --prefix _libfaiss_stage/
diff --git a/conda/faiss-gpu-raft/build-pkg.sh b/conda/faiss-gpu-raft/build-pkg.sh
@@ -7,17 +7,17 @@
 set -e
 
 
-# Build swigfaiss.so/swigfaiss_avx2.so.
+# Build swigfaiss.so/swigfaiss_avx2.so/swigfaiss_avx512.so
 cmake -B _build_python_${PY_VER} \
       -Dfaiss_ROOT=_libfaiss_stage/ \
-      -DFAISS_OPT_LEVEL=avx2 \
+      -DFAISS_OPT_LEVEL=avx512 \
       -DFAISS_ENABLE_GPU=ON \
       -DFAISS_ENABLE_RAFT=ON \
       -DCMAKE_BUILD_TYPE=Release \
       -DPython_EXECUTABLE=$PYTHON \
       faiss/python
 
-make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2
+make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2 swigfaiss_avx512
 
 # Build actual python module.
 cd _build_python_${PY_VER}/

diff --git a/conda/faiss-gpu-raft/test_cpu_dispatch.sh b/conda/faiss-gpu-raft/test_cpu_dispatch.sh
@@ -6,5 +6,6 @@
 
 set -e
 
-FAISS_DISABLE_CPU_FEATURES=AVX2 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss.so
-LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx2.so
+FAISS_OPT_LEVEL= LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss.so
+FAISS_OPT_LEVEL=AVX2 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx2.so
+FAISS_OPT_LEVEL=AVX512 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx512.so
diff --git a/conda/faiss-gpu/build-lib.sh b/conda/faiss-gpu/build-lib.sh
@@ -7,11 +7,11 @@
 set -e
 
 
-# Build libfaiss.so/libfaiss_avx2.so.
+# Build libfaiss.so/libfaiss_avx2.so/libfaiss_avx512.so
 cmake -B _build \
       -DBUILD_SHARED_LIBS=ON \
       -DBUILD_TESTING=OFF \
-      -DFAISS_OPT_LEVEL=avx2 \
+      -DFAISS_OPT_LEVEL=avx512 \
       -DFAISS_ENABLE_GPU=ON \
       -DFAISS_ENABLE_RAFT=OFF \
       -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHS}" \
@@ -20,7 +20,7 @@ cmake -B _build \
       -DCMAKE_INSTALL_LIBDIR=lib \
       -DCMAKE_BUILD_TYPE=Release .
 
-make -C _build -j$(nproc) faiss faiss_avx2
+make -C _build -j$(nproc) faiss faiss_avx2 faiss_avx512
 
 cmake --install _build --prefix $PREFIX
 cmake --install _build --prefix _libfaiss_stage/
diff --git a/conda/faiss-gpu/build-pkg.sh b/conda/faiss-gpu/build-pkg.sh
@@ -7,17 +7,17 @@
 set -e
 
 
-# Build swigfaiss.so/swigfaiss_avx2.so.
+# Build swigfaiss.so/swigfaiss_avx2.so/swigfaiss_avx512.so
 cmake -B _build_python_${PY_VER} \
       -Dfaiss_ROOT=_libfaiss_stage/ \
-      -DFAISS_OPT_LEVEL=avx2 \
+      -DFAISS_OPT_LEVEL=avx512 \
       -DFAISS_ENABLE_GPU=ON \
       -DFAISS_ENABLE_RAFT=OFF \
       -DCMAKE_BUILD_TYPE=Release \
       -DPython_EXECUTABLE=$PYTHON \
       faiss/python
 
-make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2
+make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2 swigfaiss_avx512
 
 # Build actual python module.
 cd _build_python_${PY_VER}/

diff --git a/conda/faiss-gpu/test_cpu_dispatch.sh b/conda/faiss-gpu/test_cpu_dispatch.sh
@@ -6,5 +6,6 @@
 
 set -e
 
-FAISS_DISABLE_CPU_FEATURES=AVX2 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss.so
-LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx2.so
+FAISS_OPT_LEVEL= LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss.so
+FAISS_OPT_LEVEL=AVX2 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx2.so
+FAISS_OPT_LEVEL=AVX512 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx512.so
diff --git a/conda/faiss/build-lib-osx.sh b/conda/faiss/build-lib-osx.sh
@@ -7,11 +7,11 @@
 set -e
 
 
-# Build libfaiss.so/libfaiss_avx2.so.
+# Build libfaiss.so/libfaiss_avx2.so/libfaiss_avx512.so
 cmake -B _build \
       -DBUILD_SHARED_LIBS=ON \
       -DBUILD_TESTING=OFF \
-      -DFAISS_OPT_LEVEL=avx2 \
+      -DFAISS_OPT_LEVEL=avx512 \
       -DFAISS_ENABLE_GPU=OFF \
       -DFAISS_ENABLE_PYTHON=OFF \
       -DBLA_VENDOR=Intel10_64lp \
@@ -21,7 +21,7 @@ cmake -B _build \
       -DCMAKE_INSTALL_LIBDIR=lib \
       -DCMAKE_BUILD_TYPE=Release .
 
-make -C _build -j$(nproc) faiss faiss_avx2
+make -C _build -j$(nproc) faiss faiss_avx2 faiss_avx512
 
 cmake --install _build --prefix $PREFIX
 cmake --install _build --prefix _libfaiss_stage/
diff --git a/conda/faiss/build-lib.sh b/conda/faiss/build-lib.sh
@@ -7,18 +7,18 @@
 set -e
 
 
-# Build libfaiss.so/libfaiss_avx2.so.
+# Build libfaiss.so/libfaiss_avx2.so/libfaiss_avx512.so
 cmake -B _build \
       -DBUILD_SHARED_LIBS=ON \
       -DBUILD_TESTING=OFF \
-      -DFAISS_OPT_LEVEL=avx2 \
+      -DFAISS_OPT_LEVEL=avx512 \
       -DFAISS_ENABLE_GPU=OFF \
       -DFAISS_ENABLE_PYTHON=OFF \
       -DBLA_VENDOR=Intel10_64lp \
       -DCMAKE_INSTALL_LIBDIR=lib \
       -DCMAKE_BUILD_TYPE=Release .
 
-make -C _build -j$(nproc) faiss faiss_avx2
+make -C _build -j$(nproc) faiss faiss_avx2 faiss_avx512
 
 cmake --install _build --prefix $PREFIX
 cmake --install _build --prefix _libfaiss_stage/
diff --git a/conda/faiss/build-pkg-arm64.sh b/conda/faiss/build-pkg-arm64.sh
@@ -7,7 +7,7 @@
 set -e
 
 
-# Build swigfaiss.so/swigfaiss_avx2.so.
+# Build swigfaiss.so
 cmake -B _build_python_${PY_VER} \
       -Dfaiss_ROOT=_libfaiss_stage/ \
       -DFAISS_ENABLE_GPU=OFF \

diff --git a/conda/faiss/build-pkg-osx.sh b/conda/faiss/build-pkg-osx.sh
@@ -7,10 +7,10 @@
 set -e
 
 
-# Build swigfaiss.so/swigfaiss_avx2.so.
+# Build swigfaiss.so/swigfaiss_avx2.so/swigfaiss_avx512
 cmake -B _build_python_${PY_VER} \
       -Dfaiss_ROOT=_libfaiss_stage/ \
-      -DFAISS_OPT_LEVEL=avx2 \
+      -DFAISS_OPT_LEVEL=avx512 \
       -DFAISS_ENABLE_GPU=OFF \
       -DOpenMP_CXX_FLAGS=-fopenmp=libiomp5 \
       -DOpenMP_CXX_LIB_NAMES=libiomp5 \
@@ -19,7 +19,7 @@ cmake -B _build_python_${PY_VER} \
       -DPython_EXECUTABLE=$PYTHON \
       faiss/python
 
-make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2
+make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2 swigfaiss_avx512
 
 # Build actual python module.
 cd _build_python_${PY_VER}/

diff --git a/conda/faiss/build-pkg.sh b/conda/faiss/build-pkg.sh
@@ -7,16 +7,16 @@
 set -e
 
 
-# Build swigfaiss.so/swigfaiss_avx2.so.
+# Build swigfaiss.so/swigfaiss_avx2.so/swigfaiss_avx512.so
 cmake -B _build_python_${PY_VER} \
       -Dfaiss_ROOT=_libfaiss_stage/ \
-      -DFAISS_OPT_LEVEL=avx2 \
+      -DFAISS_OPT_LEVEL=avx512 \
       -DFAISS_ENABLE_GPU=OFF \
       -DCMAKE_BUILD_TYPE=Release \
       -DPython_EXECUTABLE=$PYTHON \
       faiss/python
 
-make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2
+make -C _build_python_${PY_VER} -j$(nproc) swigfaiss swigfaiss_avx2 swigfaiss_avx512
 
 # Build actual python module.
 cd _build_python_${PY_VER}/

diff --git a/conda/faiss/test_cpu_dispatch.sh b/conda/faiss/test_cpu_dispatch.sh
@@ -6,5 +6,6 @@
 
 set -e
 
-FAISS_DISABLE_CPU_FEATURES=AVX2 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss.so
-LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx2.so
+FAISS_OPT_LEVEL= LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss.so
+FAISS_OPT_LEVEL=AVX2 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx2.so
+FAISS_OPT_LEVEL=AVX512 LD_DEBUG=libs python -c "import faiss" 2>&1 | grep libfaiss_avx512.so
diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt
@@ -328,7 +328,7 @@ if(FAISS_OPT_LEVEL STREQUAL "avx2")
   )
 endif()
 if(FAISS_OPT_LEVEL STREQUAL "avx512")
-  install(TARGETS faiss_avx512
+  install(TARGETS faiss_avx2 faiss_avx512
     EXPORT faiss-targets
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}

diff --git a/faiss/utils/utils.cpp b/faiss/utils/utils.cpp
@@ -117,7 +117,7 @@ std::string get_compile_options() {
 #ifdef __AVX2__
     options += "AVX2 ";
 #elif __AVX512F__
-    options += "AVX512";
+    options += "AVX512 ";
 #elif defined(__aarch64__)
     options += "NEON ";
 #else

diff --git a/tests/test_index_accuracy.py b/tests/test_index_accuracy.py
@@ -639,14 +639,14 @@ def test_OIVFPQ(self):
         d = ev.d
         quantizer = faiss.IndexFlatL2(d)
         index = faiss.IndexIVFPQ(quantizer, d, ncentroids, M, 8)
-        index.nprobe = 5
+        index.nprobe = 12
 
         res = ev.launch("IVFPQ", index)
         e_ivfpq = ev.evalres(res)
 
         quantizer = faiss.IndexFlatL2(d)
         index_ivfpq = faiss.IndexIVFPQ(quantizer, d, ncentroids, M, 8)
-        index_ivfpq.nprobe = 5
+        index_ivfpq.nprobe = 12
         opq_matrix = faiss.OPQMatrix(d, M)
         opq_matrix.niter = 10
         index = faiss.IndexPreTransform(opq_matrix, index_ivfpq)