From 34f5250398a9366fd8854d1b50412f0488ebb0dd Mon Sep 17 00:00:00 2001 From: vorj Date: Wed, 20 Jul 2022 01:31:14 +0000 Subject: [PATCH 1/5] add sve targets, faiss_sve and swigfaiss_sve --- .gitignore | 1 + CMakeLists.txt | 2 +- INSTALL.md | 6 +++-- faiss/CMakeLists.txt | 44 +++++++++++++++++++++++++++++++++++++ faiss/gpu/CMakeLists.txt | 1 + faiss/python/CMakeLists.txt | 28 +++++++++++++++++++++++ faiss/python/loader.py | 32 +++++++++++++++++++++++++++ faiss/python/setup.py | 12 ++++++++-- faiss/utils/utils.cpp | 2 ++ tests/CMakeLists.txt | 28 ++++++++++++++++++++++- 10 files changed, 150 insertions(+), 6 deletions(-) diff --git a/.gitignore b/.gitignore index d6df432fa5..01b98f0a9c 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,4 @@ /tests/gtest/ faiss/python/swigfaiss_avx2.swig faiss/python/swigfaiss_avx512.swig +faiss/python/swigfaiss_sve.swig diff --git a/CMakeLists.txt b/CMakeLists.txt index cedee9c456..7743f0c024 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -50,7 +50,7 @@ set(CMAKE_CXX_STANDARD 17) list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake") -# Valid values are "generic", "avx2", "avx512". +# Valid values are "generic", "avx2", "avx512", "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) diff --git a/INSTALL.md b/INSTALL.md index 5bd4f6d448..158efde091 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -122,8 +122,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`, - `avx2` and `avx512`, by increasing order of optimization), + generate code using optimized SIMD/Vector instructions. possible values are + below: + - On x86\_64, `generic`, `avx2` and `avx512`, 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 diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 1b0860f3fb..db941adc59 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -262,6 +262,31 @@ 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 $<$,$>:${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 $<$,$>:-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 $<$,$>:${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 $<$,$>:-march=armv8-a+sve>) + endif() +endif() + # Handle `#include `. target_include_directories(faiss PUBLIC $) @@ -271,6 +296,9 @@ target_include_directories(faiss_avx2 PUBLIC # Handle `#include `. target_include_directories(faiss_avx512 PUBLIC $) +# Handle `#include `. +target_include_directories(faiss_sve PUBLIC + $) set_target_properties(faiss PROPERTIES POSITION_INDEPENDENT_CODE ON @@ -284,11 +312,16 @@ set_target_properties(faiss_avx512 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_avx512 PRIVATE FAISS_MAIN_LIB) + target_compile_definitions(faiss_sve PRIVATE FAISS_MAIN_LIB) endif() string(FIND "${CMAKE_CXX_FLAGS}" "FINTEGER" finteger_idx) @@ -297,11 +330,13 @@ 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_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_avx512 PRIVATE OpenMP::OpenMP_CXX) +target_link_libraries(faiss_sve PRIVATE OpenMP::OpenMP_CXX) find_package(MKL) if(MKL_FOUND) @@ -313,11 +348,13 @@ else() 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_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 @@ -341,6 +378,13 @@ if(FAISS_OPT_LEVEL STREQUAL "avx512") LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} ) endif() +if(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}) get_filename_component(dir ${header} DIRECTORY ) diff --git a/faiss/gpu/CMakeLists.txt b/faiss/gpu/CMakeLists.txt index d20f3b7f8e..5b29987379 100644 --- a/faiss/gpu/CMakeLists.txt +++ b/faiss/gpu/CMakeLists.txt @@ -296,6 +296,7 @@ set(FAISS_GPU_HEADERS ${FAISS_GPU_HEADERS} PARENT_SCOPE) target_link_libraries(faiss PRIVATE "$") target_link_libraries(faiss_avx2 PRIVATE "$") target_link_libraries(faiss_avx512 PRIVATE "$") +target_link_libraries(faiss_sve PRIVATE "$") foreach(header ${FAISS_GPU_HEADERS}) get_filename_component(dir ${header} DIRECTORY ) diff --git a/faiss/python/CMakeLists.txt b/faiss/python/CMakeLists.txt index 0073c20e04..d8d933497e 100644 --- a/faiss/python/CMakeLists.txt +++ b/faiss/python/CMakeLists.txt @@ -50,10 +50,12 @@ 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_sve.swig COPYONLY) configure_swigfaiss(swigfaiss.swig) configure_swigfaiss(swigfaiss_avx2.swig) configure_swigfaiss(swigfaiss_avx512.swig) +configure_swigfaiss(swigfaiss_sve.swig) if(TARGET faiss) # Manually add headers as extra dependencies of swigfaiss. @@ -62,11 +64,13 @@ if(TARGET faiss) 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_avx512_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_avx512_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) @@ -112,16 +116,30 @@ if(NOT FAISS_OPT_LEVEL STREQUAL "avx512") set_target_properties(swigfaiss_avx512 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_avx512 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_avx512 PRIVATE /bigobj) + target_compile_options(swigfaiss_sve PRIVATE /bigobj) endif() if(FAISS_ENABLE_GPU) @@ -132,6 +150,7 @@ if(FAISS_ENABLE_GPU) target_link_libraries(swigfaiss PRIVATE CUDA::cudart $<$:raft::raft> $<$:nvidia::cutlass::cutlass>) target_link_libraries(swigfaiss_avx2 PRIVATE CUDA::cudart $<$:raft::raft> $<$:nvidia::cutlass::cutlass>) target_link_libraries(swigfaiss_avx512 PRIVATE CUDA::cudart $<$:raft::raft> $<$:nvidia::cutlass::cutlass>) + target_link_libraries(swigfaiss_sve PRIVATE CUDA::cudart $<$:raft::raft> $<$:nvidia::cutlass::cutlass>) endif() find_package(OpenMP REQUIRED) @@ -157,11 +176,19 @@ target_link_libraries(swigfaiss_avx512 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 `. 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_sve PRIVATE ${PROJECT_SOURCE_DIR}/../..) find_package(Python REQUIRED COMPONENTS Development NumPy @@ -186,6 +213,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_sve PRIVATE faiss_python_callbacks) configure_file(setup.py setup.py COPYONLY) configure_file(__init__.py __init__.py COPYONLY) diff --git a/faiss/python/loader.py b/faiss/python/loader.py index 8cc97f2f44..166d456921 100644 --- a/faiss/python/loader.py +++ b/faiss/python/loader.py @@ -49,6 +49,26 @@ def supported_instruction_sets(): return result 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__) @@ -92,6 +112,18 @@ def supported_instruction_sets(): # reset so that we load without AVX2 below loaded = False +has_SVE = is_sve_supported() and not "SVE" in os.getenv("FAISS_DISABLE_CPU_FEATURES", "").split(", \t\n\r") +if has_SVE and not loaded: + try: + logger.info("Loading faiss with SVE support.") + from .swigfaiss_sve import * + logger.info("Successfully loaded faiss with SVE support.") + loaded = True + 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 + loaded = False + if not loaded: # we import * so that the symbol X can be accessed as faiss.X logger.info("Loading faiss.") diff --git a/faiss/python/setup.py b/faiss/python/setup.py index 939aeeffbe..8d2150e7f7 100644 --- a/faiss/python/setup.py +++ b/faiss/python/setup.py @@ -26,14 +26,17 @@ swigfaiss_generic_lib = f"{prefix}_swigfaiss{ext}" swigfaiss_avx2_lib = f"{prefix}_swigfaiss_avx2{ext}" swigfaiss_avx512_lib = f"{prefix}_swigfaiss_avx512{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_avx512 = os.path.exists(swigfaiss_avx512_lib) +found_swigfaiss_sve = os.path.exists(swigfaiss_sve_lib) -assert (found_swigfaiss_generic or found_swigfaiss_avx2 or found_swigfaiss_avx512), \ +assert (found_swigfaiss_generic or found_swigfaiss_avx2 or found_swigfaiss_avx512 or found_swigfaiss_sve), \ f"Could not find {swigfaiss_generic_lib} or " \ - f"{swigfaiss_avx2_lib} or {swigfaiss_avx512_lib}. Faiss may not be compiled yet." + f"{swigfaiss_avx2_lib} or {swigfaiss_avx512_lib} or {swigfaiss_sve_lib}. " \ + f"Faiss may not be compiled yet." if found_swigfaiss_generic: print(f"Copying {swigfaiss_generic_lib}") @@ -50,6 +53,11 @@ shutil.copyfile("swigfaiss_avx512.py", "faiss/swigfaiss_avx512.py") shutil.copyfile(swigfaiss_avx512_lib, f"faiss/_swigfaiss_avx512{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, diff --git a/faiss/utils/utils.cpp b/faiss/utils/utils.cpp index dc6faddaf5..42d6fe2cc5 100644 --- a/faiss/utils/utils.cpp +++ b/faiss/utils/utils.cpp @@ -118,6 +118,8 @@ std::string get_compile_options() { options += "AVX2 "; #elif __AVX512F__ options += "AVX512 "; +#elif defined(__ARM_FEATURE_SVE) + options += "SVE NEON "; #elif defined(__aarch64__) options += "NEON "; #else diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3980d7dd7c..26f07398b8 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -39,7 +39,7 @@ set(FAISS_TEST_SRC add_executable(faiss_test ${FAISS_TEST_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 "sve") target_link_libraries(faiss_test PRIVATE faiss) endif() @@ -61,6 +61,32 @@ if(FAISS_OPT_LEVEL STREQUAL "avx512") target_link_libraries(faiss_test PRIVATE faiss_avx512) endif() +if(FAISS_OPT_LEVEL STREQUAL "sve") + 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_test PRIVATE $<$,$>:${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_test PRIVATE $<$,$>:-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_test PRIVATE $<$,$>:${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_test PRIVATE $<$,$>:-march=armv8-a+sve>) + endif() + else() + # TODO: support Windows + endif() + target_link_libraries(faiss_test PRIVATE faiss_sve) +endif() + include(FetchContent) FetchContent_Declare( googletest From 16658e1dc7ade298476239b6520898eb96cb5208 Mon Sep 17 00:00:00 2001 From: vorj <40021161+vorj@users.noreply.github.com> Date: Wed, 19 Jun 2024 20:00:14 +0900 Subject: [PATCH 2/5] stop using is_sve_supported() directly --- faiss/python/loader.py | 49 +++++++++++++++++++++++------------------- 1 file changed, 27 insertions(+), 22 deletions(-) diff --git a/faiss/python/loader.py b/faiss/python/loader.py index 166d456921..a0e1443250 100644 --- a/faiss/python/loader.py +++ b/faiss/python/loader.py @@ -24,6 +24,28 @@ def supported_instruction_sets(): >>> supported_instruction_sets() # for ARM {"NEON", "ASIMD", ...} """ + + # 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 + import numpy if Version(numpy.__version__) >= Version("1.19"): # use private API as next-best thing until numpy/numpy#18058 is solved @@ -31,6 +53,8 @@ def supported_instruction_sets(): # __cpu_features__ is a dictionary with CPU features # as keys, and True / False as values supported = {k for k, v in __cpu_features__.items() if v} + if is_sve_supported(): + supported.add("SVE") for f in os.getenv("FAISS_DISABLE_CPU_FEATURES", "").split(", \t\n\r"): supported.discard(f) return supported @@ -46,30 +70,11 @@ def supported_instruction_sets(): result.add("AVX2") if "avx512" in numpy.distutils.cpuinfo.cpu.info[0].get('flags', ""): result.add("AVX512") + if is_sve_supported(): + result.add("SVE") return result 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__) instruction_sets = None @@ -112,7 +117,7 @@ def is_sve_supported(): # reset so that we load without AVX2 below loaded = False -has_SVE = is_sve_supported() and not "SVE" in os.getenv("FAISS_DISABLE_CPU_FEATURES", "").split(", \t\n\r") +has_SVE = "SVE" in instruction_sets if has_SVE and not loaded: try: logger.info("Loading faiss with SVE support.") From e42c4c8ee4cd5fe83f138eecb02558fb7c24ec98 Mon Sep 17 00:00:00 2001 From: vorj <40021161+vorj@users.noreply.github.com> Date: Wed, 19 Jun 2024 20:08:56 +0900 Subject: [PATCH 3/5] use numpy.distutils.cpuinfo instead of reading /proc/cpuinfo directly --- faiss/python/loader.py | 17 ++++------------- 1 file changed, 4 insertions(+), 13 deletions(-) diff --git a/faiss/python/loader.py b/faiss/python/loader.py index a0e1443250..5abcb0d885 100644 --- a/faiss/python/loader.py +++ b/faiss/python/loader.py @@ -26,25 +26,16 @@ def supported_instruction_sets(): """ # Currently numpy.core._multiarray_umath.__cpu_features__ doesn't support Arm SVE, - # so let's read Features in /proc/cpuinfo and search 'sve' entry + # so let's read Features in numpy.distutils.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 + # platform-dependent legacy fallback using numpy.distutils.cpuinfo + import numpy.distutils.cpuinfo + return "sve" in numpy.distutils.cpuinfo.cpu.info[0].get('Features', "").split() import numpy if Version(numpy.__version__) >= Version("1.19"): From 241cd6c5da5dd51fc8b83c37a503dc11c36065ce Mon Sep 17 00:00:00 2001 From: vorj <40021161+vorj@users.noreply.github.com> Date: Wed, 19 Jun 2024 20:09:31 +0900 Subject: [PATCH 4/5] skip fallback detection if numpy >= 2.0 --- faiss/python/loader.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/faiss/python/loader.py b/faiss/python/loader.py index 5abcb0d885..f409ac173d 100644 --- a/faiss/python/loader.py +++ b/faiss/python/loader.py @@ -25,7 +25,7 @@ def supported_instruction_sets(): {"NEON", "ASIMD", ...} """ - # Currently numpy.core._multiarray_umath.__cpu_features__ doesn't support Arm SVE, + # Old numpy.core._multiarray_umath.__cpu_features__ doesn't support Arm SVE, # so let's read Features in numpy.distutils.cpuinfo and search 'sve' entry def is_sve_supported(): if platform.machine() != "aarch64": @@ -33,6 +33,10 @@ def is_sve_supported(): # Currently SVE is only supported on Linux if platform.system() != "Linux": return False + # Numpy 2.0 supports SVE detection by __cpu_features__, so just skip + import numpy + if Version(numpy.__version__) >= Version("2.0"): + return False # platform-dependent legacy fallback using numpy.distutils.cpuinfo import numpy.distutils.cpuinfo return "sve" in numpy.distutils.cpuinfo.cpu.info[0].get('Features', "").split() From bea1a8ad8e5135c7995deeaec7362f6fe3b7d213 Mon Sep 17 00:00:00 2001 From: vorj <40021161+vorj@users.noreply.github.com> Date: Wed, 19 Jun 2024 20:09:52 +0900 Subject: [PATCH 5/5] support FAISS_DISABLE_CPU_FEATURES with old numpy --- faiss/python/loader.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/faiss/python/loader.py b/faiss/python/loader.py index f409ac173d..fa75edb468 100644 --- a/faiss/python/loader.py +++ b/faiss/python/loader.py @@ -67,6 +67,8 @@ def is_sve_supported(): result.add("AVX512") if is_sve_supported(): result.add("SVE") + for f in os.getenv("FAISS_DISABLE_CPU_FEATURES", "").split(", \t\n\r"): + result.discard(f) return result return set()