add sve targets, faiss_sve and swigfaiss_sve

CMakeLists.txt

@@ -52,7 +52,7 @@ endif()
 
 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
 
-# Valid values are "generic", "avx2".
+# Valid values are "generic", "avx2", "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

@@ -116,8 +116,10 @@ Several options can be passed to CMake, among which:
   - `-DCMAKE_BUILD_TYPE=Release` in order to enable generic compiler
   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 instructions (possible values are `generic`
-  and `avx2`, by increasing order of optimization),
+  generate code using optimized SIMD/Vector instructions. possible values are
+  below:
+    - On x86\_64, `generic` and `avx2` , by increasing order of optimization,
+    - On aarch64, `generic` and `sve` , by increasing order of optimization,
 - BLAS-related options:
   - `-DBLA_VENDOR=Intel10_64_dyn -DMKL_LIBRARIES=/path/to/mkl/libs` to use the
   Intel MKL BLAS implementation, which is significantly faster than OpenBLAS

faiss/CMakeLists.txt

@@ -242,12 +242,39 @@ else()
   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)
+endif()
+if(NOT WIN32)
+  if("${CMAKE_CXX_FLAGS} ${CMAKE_CXX_FLAGS_DEBUG} " MATCHES "(^| )-march=native")
+    # Do nothing, expect SVE to be enabled by -march=native
+  elseif("${CMAKE_CXX_FLAGS} ${CMAKE_CXX_FLAGS_DEBUG} " MATCHES "(^| )(-march=armv[0-9]+(\\.[1-9]+)?-[^+ ](\\+[^+$ ]+)*)")
+    # Add +sve
+    target_compile_options(faiss_sve PRIVATE $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CONFIG:DEBUG>>:${CMAKE_MATCH_2}+sve>)
+  elseif(NOT "${CMAKE_CXX_FLAGS} ${CMAKE_CXX_FLAGS_DEBUG} " MATCHES "(^| )-march=armv")
+    # No valid -march, so specify -march=armv8-a+sve as the default
+    target_compile_options(faiss_sve PRIVATE $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CONFIG:DEBUG>>:-march=armv8-a+sve>)
+  endif()
+  if("${CMAKE_CXX_FLAGS} ${CMAKE_CXX_FLAGS_RELEASE} " MATCHES "(^| )-march=native")
+    # Do nothing, expect SVE to be enabled by -march=native
+  elseif("${CMAKE_CXX_FLAGS} ${CMAKE_CXX_FLAGS_RELEASE} " MATCHES "(^| )(-march=armv[0-9]+(\\.[1-9]+)?-[^+ ](\\+[^+$ ]+)*)")
+    # Add +sve
+    target_compile_options(faiss_sve PRIVATE $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CONFIG:RELEASE>>:${CMAKE_MATCH_2}+sve>)
+  elseif(NOT "${CMAKE_CXX_FLAGS} ${CMAKE_CXX_FLAGS_RELEASE} " MATCHES "(^| )-march=armv")
+    # No valid -march, so specify -march=armv8-a+sve as the default
+    target_compile_options(faiss_sve PRIVATE $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CONFIG:RELEASE>>:-march=armv8-a+sve>)
+  endif()
+endif()
+
 # Handle `#include <faiss/foo.h>`.
 target_include_directories(faiss PUBLIC
   $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>)
 # Handle `#include <faiss/foo.h>`.
 target_include_directories(faiss_avx2 PUBLIC
   $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>)
+target_include_directories(faiss_sve PUBLIC
+  $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>)
 
 set_target_properties(faiss PROPERTIES
   POSITION_INDEPENDENT_CODE ON
@@ -257,18 +284,25 @@ set_target_properties(faiss_avx2 PROPERTIES
   POSITION_INDEPENDENT_CODE ON
   WINDOWS_EXPORT_ALL_SYMBOLS ON
 )
+set_target_properties(faiss_sve PROPERTIES
+  POSITION_INDEPENDENT_CODE ON
+  WINDOWS_EXPORT_ALL_SYMBOLS ON
+)
 
 if(WIN32)
   target_compile_definitions(faiss PRIVATE FAISS_MAIN_LIB)
   target_compile_definitions(faiss_avx2 PRIVATE FAISS_MAIN_LIB)
+  target_compile_definitions(faiss_sve PRIVATE FAISS_MAIN_LIB)
 endif()
 
 target_compile_definitions(faiss PRIVATE FINTEGER=int)
 target_compile_definitions(faiss_avx2 PRIVATE FINTEGER=int)
+target_compile_definitions(faiss_sve PRIVATE FINTEGER=int)
 
 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_sve PRIVATE OpenMP::OpenMP_CXX)
 
 find_package(MKL)
 if(MKL_FOUND)
@@ -278,10 +312,12 @@ else()
   find_package(BLAS REQUIRED)
   target_link_libraries(faiss PRIVATE ${BLAS_LIBRARIES})
   target_link_libraries(faiss_avx2 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_sve PRIVATE ${LAPACK_LIBRARIES})
 endif()
 
 install(TARGETS faiss
@@ -297,6 +333,12 @@ if(FAISS_OPT_LEVEL STREQUAL "avx2")
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
   )
+elseif(FAISS_OPT_LEVEL STREQUAL "sve")
+  install(TARGETS faiss_sve
+    EXPORT faiss-targets
+    ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  )
 endif()
 
 foreach(header ${FAISS_HEADERS})

faiss/gpu/CMakeLists.txt

@@ -191,6 +191,7 @@ set(FAISS_GPU_HEADERS ${FAISS_GPU_HEADERS} PARENT_SCOPE)
 
 target_sources(faiss PRIVATE ${FAISS_GPU_SRC})
 target_sources(faiss_avx2 PRIVATE ${FAISS_GPU_SRC})
+target_sources(faiss_sve PRIVATE ${FAISS_GPU_SRC})
 
 foreach(header ${FAISS_GPU_HEADERS})
   get_filename_component(dir ${header} DIRECTORY )
@@ -202,5 +203,7 @@ endforeach()
 find_package(CUDAToolkit REQUIRED)
 target_link_libraries(faiss PRIVATE CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
 target_link_libraries(faiss_avx2 PRIVATE CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
+target_link_libraries(faiss_sve PRIVATE CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
 target_compile_options(faiss PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr>)
 target_compile_options(faiss_avx2 PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr>)
+target_compile_options(faiss_sve PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr>)

faiss/python/CMakeLists.txt

@@ -44,20 +44,24 @@ endmacro()
 # CMake's SWIG wrappers only allow tweaking certain settings at source level, so
 # 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_sve.swig COPYONLY)
 
 configure_swigfaiss(swigfaiss.swig)
 configure_swigfaiss(swigfaiss_avx2.swig)
+configure_swigfaiss(swigfaiss_sve.swig)
 
 if(TARGET faiss)
   # Manually add headers as extra dependencies of swigfaiss.
   set(SWIG_MODULE_swigfaiss_EXTRA_DEPS)
   foreach(h ${FAISS_HEADERS})
     list(APPEND SWIG_MODULE_swigfaiss_EXTRA_DEPS "${faiss_SOURCE_DIR}/faiss/${h}")
     list(APPEND SWIG_MODULE_swigfaiss_avx2_EXTRA_DEPS "${faiss_SOURCE_DIR}/faiss/${h}")
+    list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS "${faiss_SOURCE_DIR}/faiss/${h}")
   endforeach()
   foreach(h ${FAISS_GPU_HEADERS})
     list(APPEND SWIG_MODULE_swigfaiss_EXTRA_DEPS "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
     list(APPEND SWIG_MODULE_swigfaiss_avx2_EXTRA_DEPS "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
+    list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
   endforeach()
 else()
   find_package(faiss REQUIRED)
@@ -82,20 +86,35 @@ if(NOT FAISS_OPT_LEVEL STREQUAL "avx2")
   set_target_properties(swigfaiss_avx2 PROPERTIES EXCLUDE_FROM_ALL TRUE)
 endif()
 
+set_property(SOURCE swigfaiss_sve.swig
+  PROPERTY SWIG_MODULE_NAME swigfaiss_sve)
+swig_add_library(swigfaiss_sve
+  TYPE SHARED
+  LANGUAGE python
+  SOURCES swigfaiss_sve.swig
+)
+set_property(TARGET swigfaiss_sve PROPERTY SWIG_COMPILE_OPTIONS -doxygen)
+if(NOT FAISS_OPT_LEVEL STREQUAL "sve")
+  set_target_properties(swigfaiss_sve PROPERTIES EXCLUDE_FROM_ALL TRUE)
+endif()
+
 if(NOT WIN32)
   # NOTE: Python does not recognize the dylib extension.
   set_target_properties(swigfaiss PROPERTIES SUFFIX .so)
   set_target_properties(swigfaiss_avx2 PROPERTIES SUFFIX .so)
+  set_target_properties(swigfaiss_sve 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_sve PRIVATE /bigobj)
 endif()
 
 if(FAISS_ENABLE_GPU)
   find_package(CUDAToolkit REQUIRED)
   target_link_libraries(swigfaiss PRIVATE CUDA::cudart)
   target_link_libraries(swigfaiss_avx2 PRIVATE CUDA::cudart)
+  target_link_libraries(swigfaiss_sve PRIVATE CUDA::cudart)
 endif()
 
 find_package(OpenMP REQUIRED)
@@ -114,10 +133,18 @@ target_link_libraries(swigfaiss_avx2 PRIVATE
   OpenMP::OpenMP_CXX
 )
 
+target_link_libraries(swigfaiss_sve PRIVATE
+  faiss_sve
+  Python::Module
+  Python::NumPy
+  OpenMP::OpenMP_CXX
+)
+
 # Hack so that python_callbacks.h can be included as
 # `#include <faiss/python/python_callbacks.h>`.
 target_include_directories(swigfaiss PRIVATE ${PROJECT_SOURCE_DIR}/../..)
 target_include_directories(swigfaiss_avx2 PRIVATE ${PROJECT_SOURCE_DIR}/../..)
+target_include_directories(swigfaiss_sve PRIVATE ${PROJECT_SOURCE_DIR}/../..)
 
 find_package(Python REQUIRED
   COMPONENTS Development NumPy
@@ -137,6 +164,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_sve PRIVATE faiss_python_callbacks)
 
 configure_file(setup.py setup.py COPYONLY)
 configure_file(__init__.py __init__.py COPYONLY)

faiss/python/loader.py

@@ -45,6 +45,26 @@ def supported_instruction_sets():
             return {"AVX2"}
     return set()
 
+# Currently numpy.core._multiarray_umath.__cpu_features__ doesn't support Arm SVE,
+# so let's read Features in /proc/cpuinfo and search 'sve' entry
+def is_sve_supported():
+    if platform.machine() != "aarch64":
+        return False
+    # Currently SVE is only supported on Linux
+    if platform.system() != "Linux":
+        return False
+    if not os.path.exists('/proc/cpuinfo'):
+        return False
+    proc = subprocess.Popen(['cat', '/proc/cpuinfo'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
+    so, _se = proc.communicate()
+    if proc.returncode != 0:
+        return False
+    for line in so.decode(encoding='UTF-8').splitlines():
+        if ':' in line:
+            l, r = line.split(':', 1)
+            if l.strip() == 'Features' and "sve" in r.strip().split():
+                return True
+    return False
 
 logger = logging.getLogger(__name__)
 
@@ -59,7 +79,18 @@ def supported_instruction_sets():
         # reset so that we load without AVX2 below
         has_AVX2 = False
 
-if not has_AVX2:
+has_SVE = is_sve_supported() and not "SVE" in os.getenv("FAISS_DISABLE_CPU_FEATURES", "").split(", \t\n\r")
+if has_SVE:
+    try:
+        logger.info("Loading faiss with SVE support.")
+        from .swigfaiss_sve import *
+        logger.info("Successfully loaded faiss with SVE support.")
+    except ImportError as e:
+        logger.info(f"Could not load library with SVE support due to:\n{e!r}")
+        # reset so that we load without SVE below
+        has_SVE = False
+
+if not has_AVX2 and not has_SVE:
     # we import * so that the symbol X can be accessed as faiss.X
     logger.info("Loading faiss.")
     from .swigfaiss import *

faiss/python/setup.py

@@ -25,13 +25,16 @@
 
 swigfaiss_generic_lib = f"{prefix}_swigfaiss{ext}"
 swigfaiss_avx2_lib = f"{prefix}_swigfaiss_avx2{ext}"
+swigfaiss_sve_lib = f"{prefix}_swigfaiss_sve{ext}"
 
 found_swigfaiss_generic = os.path.exists(swigfaiss_generic_lib)
 found_swigfaiss_avx2 = os.path.exists(swigfaiss_avx2_lib)
+found_swigfaiss_sve = os.path.exists(swigfaiss_sve_lib)
 
-assert (found_swigfaiss_generic or found_swigfaiss_avx2), \
-    f"Could not find {swigfaiss_generic_lib} or " \
-    f"{swigfaiss_avx2_lib}. Faiss may not be compiled yet."
+assert (found_swigfaiss_generic or found_swigfaiss_avx2 or found_swigfaiss_sve), \
+    f"Could not find {swigfaiss_generic_lib}, " \
+    f"{swigfaiss_avx2_lib}, or {swigfaiss_sve_lib}." \
+    f" Faiss may not be compiled yet."
 
 if found_swigfaiss_generic:
     print(f"Copying {swigfaiss_generic_lib}")
@@ -43,6 +46,11 @@
     shutil.copyfile("swigfaiss_avx2.py", "faiss/swigfaiss_avx2.py")
     shutil.copyfile(swigfaiss_avx2_lib, f"faiss/_swigfaiss_avx2{ext}")
 
+if found_swigfaiss_sve:
+    print(f"Copying {swigfaiss_sve_lib}")
+    shutil.copyfile("swigfaiss_sve.py", "faiss/swigfaiss_sve.py")
+    shutil.copyfile(swigfaiss_sve_lib, f"faiss/_swigfaiss_sve{ext}")
+
 long_description="""
 Faiss is a library for efficient similarity search and clustering of dense
 vectors. It contains algorithms that search in sets of vectors of any size,