Add a new architecture mode: 'avx512-sr'.

.github/workflows/build.yml

@@ -67,6 +67,17 @@ jobs:
         uses: ./.github/actions/build_cmake
         with:
           opt_level: avx512
+  linux-x86_64-AVX512-sr-cmake:
+    name: Linux x86_64 AVX512-SR  (cmake)
+    needs: linux-x86_64-cmake
+    runs-on: faiss-aws-m7i.large
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Build and Test (cmake)
+        uses: ./.github/actions/build_cmake
+        with:
+          opt_level: avx512-sr
   linux-x86_64-GPU-cmake:
     name: Linux x86_64 GPU (cmake)
     needs: linux-x86_64-cmake

.gitignore

@@ -19,4 +19,5 @@
 /tests/gtest/
 faiss/python/swigfaiss_avx2.swig
 faiss/python/swigfaiss_avx512.swig
+faiss/python/swigfaiss_avx512_sr.swig
 faiss/python/swigfaiss_sve.swig

CMakeLists.txt

@@ -57,7 +57,7 @@ set(CMAKE_CXX_STANDARD 17)
 
 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
 
-# Valid values are "generic", "avx2", "avx512", "sve".
+# Valid values are "generic", "avx2", "avx512", "avx512-sr", "sve".
 option(FAISS_OPT_LEVEL "" "generic")
 option(FAISS_ENABLE_GPU "Enable support for GPU indexes." ON)
 option(FAISS_ENABLE_RAFT "Enable RAFT for GPU indexes." OFF)

INSTALL.md

@@ -131,7 +131,7 @@ Several options can be passed to CMake, among which:
   optimization options (enables `-O3` on gcc for instance),
   - `-DFAISS_OPT_LEVEL=avx2` in order to enable the required compiler flags to
   generate code using optimized SIMD/Vector instructions. Possible values are below:
-    - On x86-64, `generic`, `avx2` and `avx512`, by increasing order of optimization,
+    - On x86-64, `generic`, `avx2`, 'avx512', and `avx512-sr`, by increasing order of optimization,
     - On aarch64, `generic` and `sve`, by increasing order of optimization,
   - `-DFAISS_USE_LTO=ON` in order to enable [Link-Time Optimization](https://en.wikipedia.org/wiki/Link-time_optimization) (default is `OFF`, possible values are `ON` and `OFF`).
 - BLAS-related options:
@@ -180,6 +180,12 @@ For AVX512:
 $ make -C build -j faiss_avx512
 ```
 
+For AVX512 (advanced):
+
+``` shell
+$ make -C build -j faiss_avx512_sr
+```
+
 This will ensure the creation of neccesary files when building and installing the python package.
 
 ## Step 3: Building the python bindings (optional)

cmake/link_to_faiss_lib.cmake

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 function(link_to_faiss_lib target)
-  if(NOT FAISS_OPT_LEVEL STREQUAL "avx2" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512" AND NOT FAISS_OPT_LEVEL STREQUAL "sve")
+  if(NOT FAISS_OPT_LEVEL STREQUAL "avx2" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512-sr" AND NOT FAISS_OPT_LEVEL STREQUAL "sve")
     target_link_libraries(${target} PRIVATE faiss)
   endif()
 
@@ -27,6 +27,17 @@ function(link_to_faiss_lib target)
     target_link_libraries(${target} PRIVATE faiss_avx512)
   endif()
 
+  if(FAISS_OPT_LEVEL STREQUAL "avx512-sr")
+    if(NOT WIN32)
+      # Architecture mode to support AVX512 extensions available since Intel (R) Sapphire Rapids.
+      # Ref: https://networkbuilders.intel.com/solutionslibrary/intel-avx-512-fp16-instruction-set-for-intel-xeon-processor-based-products-technology-guide
+      target_compile_options(${target} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-march=sapphirerapids -mtune=sapphirerapids>)
+    else()
+      target_compile_options(${target} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/arch:AVX512>)
+    endif()
+    target_link_libraries(${target} PRIVATE faiss_avx512_sr)
+  endif()
+
   if(FAISS_OPT_LEVEL STREQUAL "sve")
     if(NOT WIN32)
       if("${CMAKE_CXX_FLAGS} ${CMAKE_CXX_FLAGS_DEBUG} " MATCHES "(^| )-march=native")

faiss/CMakeLists.txt

@@ -234,7 +234,7 @@ set(FAISS_HEADERS ${FAISS_HEADERS} PARENT_SCOPE)
 add_library(faiss ${FAISS_SRC})
 
 add_library(faiss_avx2 ${FAISS_SRC})
-if(NOT FAISS_OPT_LEVEL STREQUAL "avx2" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512")
+if(NOT FAISS_OPT_LEVEL STREQUAL "avx2" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512-sr")
   set_target_properties(faiss_avx2 PROPERTIES EXCLUDE_FROM_ALL TRUE)
 endif()
 if(NOT WIN32)
@@ -263,6 +263,20 @@ else()
   add_compile_options(/bigobj)
 endif()
 
+add_library(faiss_avx512_sr ${FAISS_SRC})
+if(NOT FAISS_OPT_LEVEL STREQUAL "avx512-sr")
+  set_target_properties(faiss_avx512_sr PROPERTIES EXCLUDE_FROM_ALL TRUE)
+endif()
+if(NOT WIN32)
+  # Architecture mode to support AVX512 extensions available since Intel(R) Sapphire Rapids.
+  # Ref: https://networkbuilders.intel.com/solutionslibrary/intel-avx-512-fp16-instruction-set-for-intel-xeon-processor-based-products-technology-guide
+  target_compile_options(faiss_avx512_sr PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-march=sapphirerapids -mtune=sapphirerapids>)
+else()
+  target_compile_options(faiss_avx512_sr PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/arch:AVX512>)
+  # we need bigobj for the swig wrapper
+  add_compile_options(/bigobj)
+endif()
+
 add_library(faiss_sve ${FAISS_SRC})
 if(NOT FAISS_OPT_LEVEL STREQUAL "sve")
   set_target_properties(faiss_sve PROPERTIES EXCLUDE_FROM_ALL TRUE)
@@ -298,10 +312,13 @@ target_include_directories(faiss_avx2 PUBLIC
 target_include_directories(faiss_avx512 PUBLIC
   $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>)
 # Handle `#include <faiss/foo.h>`.
+target_include_directories(faiss_avx512_sr PUBLIC
+  $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>)
+# Handle `#include <faiss/foo.h>`.
 target_include_directories(faiss_sve PUBLIC
   $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>)
 
-set_target_properties(faiss faiss_avx2 faiss_avx512 faiss_sve PROPERTIES
+set_target_properties(faiss faiss_avx2 faiss_avx512 faiss_avx512_sr faiss_sve PROPERTIES
   POSITION_INDEPENDENT_CODE ON
   WINDOWS_EXPORT_ALL_SYMBOLS ON
 )
@@ -310,6 +327,7 @@ if(WIN32)
   target_compile_definitions(faiss PRIVATE FAISS_MAIN_LIB)
   target_compile_definitions(faiss_avx2 PRIVATE FAISS_MAIN_LIB)
   target_compile_definitions(faiss_avx512 PRIVATE FAISS_MAIN_LIB)
+  target_compile_definitions(faiss_avx512_sr PRIVATE FAISS_MAIN_LIB)
   target_compile_definitions(faiss_sve PRIVATE FAISS_MAIN_LIB)
 endif()
 
@@ -319,6 +337,7 @@ if (${finteger_idx} EQUAL -1)
 endif()
 target_compile_definitions(faiss_avx2 PRIVATE FINTEGER=int)
 target_compile_definitions(faiss_avx512 PRIVATE FINTEGER=int)
+target_compile_definitions(faiss_avx512_sr PRIVATE FINTEGER=int)
 target_compile_definitions(faiss_sve PRIVATE FINTEGER=int)
 
 if(FAISS_USE_LTO)
@@ -330,6 +349,7 @@ if(FAISS_USE_LTO)
     set_property(TARGET faiss PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
     set_property(TARGET faiss_avx2 PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
     set_property(TARGET faiss_avx512 PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
+    set_property(TARGET faiss_avx512_sr PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
   else()
     message(STATUS "LTO not supported: <${ipo_error}>")
   endif()
@@ -339,24 +359,28 @@ find_package(OpenMP REQUIRED)
 target_link_libraries(faiss PRIVATE OpenMP::OpenMP_CXX)
 target_link_libraries(faiss_avx2 PRIVATE OpenMP::OpenMP_CXX)
 target_link_libraries(faiss_avx512 PRIVATE OpenMP::OpenMP_CXX)
+target_link_libraries(faiss_avx512_sr PRIVATE OpenMP::OpenMP_CXX)
 target_link_libraries(faiss_sve PRIVATE OpenMP::OpenMP_CXX)
 
 find_package(MKL)
 if(MKL_FOUND)
   target_link_libraries(faiss PRIVATE ${MKL_LIBRARIES})
   target_link_libraries(faiss_avx2 PRIVATE ${MKL_LIBRARIES})
   target_link_libraries(faiss_avx512 PRIVATE ${MKL_LIBRARIES})
+  target_link_libraries(faiss_avx512_sr PRIVATE ${MKL_LIBRARIES})
 else()
   find_package(BLAS REQUIRED)
   target_link_libraries(faiss PRIVATE ${BLAS_LIBRARIES})
   target_link_libraries(faiss_avx2 PRIVATE ${BLAS_LIBRARIES})
   target_link_libraries(faiss_avx512 PRIVATE ${BLAS_LIBRARIES})
+  target_link_libraries(faiss_avx512_sr 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_avx512_sr PRIVATE ${LAPACK_LIBRARIES})
   target_link_libraries(faiss_sve PRIVATE ${LAPACK_LIBRARIES})
 endif()
 
@@ -381,6 +405,13 @@ if(FAISS_OPT_LEVEL STREQUAL "avx512")
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
   )
 endif()
+if(FAISS_OPT_LEVEL STREQUAL "avx512-sr")
+  install(TARGETS faiss_avx2 faiss_avx512 faiss_avx512_sr
+    EXPORT faiss-targets
+    ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  )
+endif()
 if(FAISS_OPT_LEVEL STREQUAL "sve")
   install(TARGETS faiss_sve
     EXPORT faiss-targets

faiss/gpu/CMakeLists.txt

@@ -270,6 +270,7 @@ if(FAISS_ENABLE_RAFT)
   target_compile_definitions(faiss PUBLIC USE_NVIDIA_RAFT=1)
   target_compile_definitions(faiss_avx2 PUBLIC USE_NVIDIA_RAFT=1)
   target_compile_definitions(faiss_avx512 PUBLIC USE_NVIDIA_RAFT=1)
+  target_compile_definitions(faiss_avx512_sr PUBLIC USE_NVIDIA_RAFT=1)
 
   # Mark all functions as hidden so that we don't generate
   # global 'public' functions that also exist in libraft.so
@@ -305,6 +306,7 @@ set(FAISS_GPU_HEADERS ${FAISS_GPU_HEADERS} PARENT_SCOPE)
 target_link_libraries(faiss PRIVATE  "$<LINK_LIBRARY:WHOLE_ARCHIVE,faiss_gpu>")
 target_link_libraries(faiss_avx2 PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,faiss_gpu>")
 target_link_libraries(faiss_avx512 PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,faiss_gpu>")
+target_link_libraries(faiss_avx512_sr PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,faiss_gpu>")
 target_link_libraries(faiss_sve PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,faiss_gpu>")
 
 foreach(header ${FAISS_GPU_HEADERS})

faiss/python/CMakeLists.txt

@@ -54,11 +54,13 @@ endmacro()
 # we duplicate the source in order to override the module name.
 configure_file(swigfaiss.swig ${CMAKE_CURRENT_SOURCE_DIR}/swigfaiss_avx2.swig COPYONLY)
 configure_file(swigfaiss.swig ${CMAKE_CURRENT_SOURCE_DIR}/swigfaiss_avx512.swig COPYONLY)
+configure_file(swigfaiss.swig ${CMAKE_CURRENT_SOURCE_DIR}/swigfaiss_avx512_sr.swig COPYONLY)
 configure_file(swigfaiss.swig ${CMAKE_CURRENT_SOURCE_DIR}/swigfaiss_sve.swig COPYONLY)
 
 configure_swigfaiss(swigfaiss.swig)
 configure_swigfaiss(swigfaiss_avx2.swig)
 configure_swigfaiss(swigfaiss_avx512.swig)
+configure_swigfaiss(swigfaiss_avx512_sr.swig)
 configure_swigfaiss(swigfaiss_sve.swig)
 configure_swigfaiss(faiss_example_external_module.swig)
 
@@ -72,6 +74,8 @@ if(TARGET faiss)
       "${faiss_SOURCE_DIR}/faiss/${h}")
     list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
       "${faiss_SOURCE_DIR}/faiss/${h}")
+    list(APPEND SWIG_MODULE_swigfaiss_avx512_sr_EXTRA_DEPS
+      "${faiss_SOURCE_DIR}/faiss/${h}")
     list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS
       "${faiss_SOURCE_DIR}/faiss/${h}")
     list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
@@ -85,6 +89,8 @@ if(TARGET faiss)
         "${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
       list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
         "${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
+      list(APPEND SWIG_MODULE_swigfaiss_avx512_sr_EXTRA_DEPS
+        "${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
       list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
         "${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
     endforeach()
@@ -96,6 +102,8 @@ if(TARGET faiss)
         "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
       list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
         "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
+      list(APPEND SWIG_MODULE_swigfaiss_avx512_sr_EXTRA_DEPS
+        "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
       list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS
         "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
       list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
@@ -146,6 +154,18 @@ if(NOT FAISS_OPT_LEVEL STREQUAL "avx512")
   set_target_properties(swigfaiss_avx512 PROPERTIES EXCLUDE_FROM_ALL TRUE)
 endif()
 
+set_property(SOURCE swigfaiss_avx512_sr.swig
+  PROPERTY SWIG_MODULE_NAME swigfaiss_avx512_sr)
+swig_add_library(swigfaiss_avx512_sr
+  TYPE SHARED
+  LANGUAGE python
+  SOURCES swigfaiss_avx512_sr.swig
+)
+set_property(TARGET swigfaiss_avx512_sr PROPERTY SWIG_COMPILE_OPTIONS -doxygen)
+if(NOT FAISS_OPT_LEVEL STREQUAL "avx512-sr")
+  set_target_properties(swigfaiss_avx512_sr PROPERTIES EXCLUDE_FROM_ALL TRUE)
+endif()
+
 set_property(SOURCE swigfaiss_sve.swig
   PROPERTY SWIG_MODULE_NAME swigfaiss_sve)
 swig_add_library(swigfaiss_sve
@@ -172,13 +192,15 @@ if(NOT WIN32)
   set_target_properties(swigfaiss PROPERTIES SUFFIX .so)
   set_target_properties(swigfaiss_avx2 PROPERTIES SUFFIX .so)
   set_target_properties(swigfaiss_avx512 PROPERTIES SUFFIX .so)
+  set_target_properties(swigfaiss_avx512_sr PROPERTIES SUFFIX .so)
   set_target_properties(swigfaiss_sve PROPERTIES SUFFIX .so)
   set_target_properties(faiss_example_external_module PROPERTIES SUFFIX .so)
 else()
   # we need bigobj for the swig wrapper
   target_compile_options(swigfaiss PRIVATE /bigobj)
   target_compile_options(swigfaiss_avx2 PRIVATE /bigobj)
   target_compile_options(swigfaiss_avx512 PRIVATE /bigobj)
+  target_compile_options(swigfaiss_avx512_sr PRIVATE /bigobj)
   target_compile_options(swigfaiss_sve PRIVATE /bigobj)
   target_compile_options(faiss_example_external_module PRIVATE /bigobj)
 endif()
@@ -188,6 +210,7 @@ if(FAISS_ENABLE_GPU)
     target_link_libraries(swigfaiss PRIVATE hip::host)
     target_link_libraries(swigfaiss_avx2 PRIVATE hip::host)
     target_link_libraries(swigfaiss_avx512 PRIVATE hip::host)
+    target_link_libraries(swigfaiss_avx512_sr PRIVATE hip::host)
     target_link_libraries(faiss_example_external_module PRIVATE hip::host)
   else()
     find_package(CUDAToolkit REQUIRED)
@@ -203,6 +226,9 @@ if(FAISS_ENABLE_GPU)
     target_link_libraries(swigfaiss_avx512 PRIVATE CUDA::cudart
       $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft>
       $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
+    target_link_libraries(swigfaiss_avx512_sr PRIVATE CUDA::cudart
+      $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft>
+      $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
     target_link_libraries(swigfaiss_sve PRIVATE CUDA::cudart
       $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft>
       $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
@@ -232,6 +258,13 @@ target_link_libraries(swigfaiss_avx512 PRIVATE
   OpenMP::OpenMP_CXX
 )
 
+target_link_libraries(swigfaiss_avx512_sr PRIVATE
+  faiss_avx512_sr
+  Python::Module
+  Python::NumPy
+  OpenMP::OpenMP_CXX
+)
+
 target_link_libraries(swigfaiss_sve PRIVATE
   faiss_sve
   Python::Module
@@ -252,6 +285,7 @@ target_link_libraries(faiss_example_external_module PRIVATE
 target_include_directories(swigfaiss PRIVATE ${PROJECT_SOURCE_DIR}/../..)
 target_include_directories(swigfaiss_avx2 PRIVATE ${PROJECT_SOURCE_DIR}/../..)
 target_include_directories(swigfaiss_avx512 PRIVATE ${PROJECT_SOURCE_DIR}/../..)
+target_include_directories(swigfaiss_avx512_sr PRIVATE ${PROJECT_SOURCE_DIR}/../..)
 target_include_directories(swigfaiss_sve PRIVATE ${PROJECT_SOURCE_DIR}/../..)
 target_include_directories(faiss_example_external_module PRIVATE ${PROJECT_SOURCE_DIR}/../..)
 
@@ -278,6 +312,7 @@ target_include_directories(faiss_python_callbacks PRIVATE ${Python_INCLUDE_DIRS}
 target_link_libraries(swigfaiss PRIVATE faiss_python_callbacks)
 target_link_libraries(swigfaiss_avx2 PRIVATE faiss_python_callbacks)
 target_link_libraries(swigfaiss_avx512 PRIVATE faiss_python_callbacks)
+target_link_libraries(swigfaiss_avx512_sr PRIVATE faiss_python_callbacks)
 target_link_libraries(swigfaiss_sve PRIVATE faiss_python_callbacks)
 target_link_libraries(faiss_example_external_module PRIVATE faiss_python_callbacks)
 

faiss/python/loader.py

@@ -67,6 +67,8 @@ def is_sve_supported():
             result.add("AVX2")
         if "avx512" in numpy.distutils.cpuinfo.cpu.info[0].get('flags', ""):
             result.add("AVX512")
+        if "avx512-sr" in numpy.distutils.cpuinfo.cpu.info[0].get('flags', ""):
+            result.add("AVX512_SR")
         if is_sve_supported():
             result.add("SVE")
         for f in os.getenv("FAISS_DISABLE_CPU_FEATURES", "").split(", \t\n\r"):
@@ -92,6 +94,18 @@ def is_sve_supported():
     instruction_sets.add(opt_level)
 
 loaded = False
+has_AVX512_SR = any("AVX512_SR" in x.upper() for x in instruction_sets)
+if has_AVX512_SR:
+    try:
+        logger.info("Loading faiss with AVX512-SR support.")
+        from .swigfaiss_avx512_sr import *
+        logger.info("Successfully loaded faiss with AVX512-SR support.")
+        loaded = True
+    except ImportError as e:
+        logger.info(f"Could not load library with AVX512-SR support due to:\n{e!r}")
+        # reset so that we load without AVX512 below
+        loaded = False
+
 has_AVX512 = any("AVX512" in x.upper() for x in instruction_sets)
 if has_AVX512:
     try: