Skip to content

Commit b74644c

Browse files
mulugetamfacebook-github-bot
authored andcommitted
Add a new architecture mode: 'avx512_spr'. (facebookresearch#4025)
Summary: This PR adds a new architecture mode to support the new extensions to AVX512, namely [AVX512-FP16](https://networkbuilders.intel.com/solutionslibrary/intel-avx-512-fp16-instruction-set-for-intel-xeon-processor-based-products-technology-guide), which have been available since Intel® Sapphire Rapids. This PR is a prerequisite for [PR#4020](facebookresearch#4020) that speeds up hamming distance evaluations. Pull Request resolved: facebookresearch#4025 Reviewed By: pankajsingh88 Differential Revision: D67524575 Pulled By: mengdilin fbshipit-source-id: f3a09943b062d720b241f95aef2f390923ffd779
1 parent 3433184 commit b74644c

11 files changed

Lines changed: 125 additions & 9 deletions

File tree

.github/workflows/build-pull-request.yml

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,17 @@ jobs:
6060
uses: ./.github/actions/build_cmake
6161
with:
6262
opt_level: avx512
63+
linux-x86_64-AVX512_SPR-cmake:
64+
name: Linux x86_64 AVX512_SPR (cmake)
65+
needs: linux-x86_64-cmake
66+
runs-on: faiss-aws-m7i.large
67+
steps:
68+
- name: Checkout
69+
uses: actions/checkout@v4
70+
- name: Build and Test (cmake)
71+
uses: ./.github/actions/build_cmake
72+
with:
73+
opt_level: avx512_spr
6374
linux-x86_64-GPU-cmake:
6475
name: Linux x86_64 GPU (cmake)
6576
needs: linux-x86_64-cmake

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,4 +19,5 @@
1919
/tests/gtest/
2020
faiss/python/swigfaiss_avx2.swig
2121
faiss/python/swigfaiss_avx512.swig
22+
faiss/python/swigfaiss_avx512_spr.swig
2223
faiss/python/swigfaiss_sve.swig

CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ set(CMAKE_CXX_STANDARD 17)
5757

5858
list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
5959

60-
# Valid values are "generic", "avx2", "avx512", "sve".
60+
# Valid values are "generic", "avx2", "avx512", "avx512_spr", "sve".
6161
option(FAISS_OPT_LEVEL "" "generic")
6262
option(FAISS_ENABLE_GPU "Enable support for GPU indexes." ON)
6363
option(FAISS_ENABLE_CUVS "Enable cuVS for GPU indexes." OFF)

INSTALL.md

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ Several options can be passed to CMake, among which:
131131
optimization options (enables `-O3` on gcc for instance),
132132
- `-DFAISS_OPT_LEVEL=avx2` in order to enable the required compiler flags to
133133
generate code using optimized SIMD/Vector instructions. Possible values are below:
134-
- On x86-64, `generic`, `avx2` and `avx512`, by increasing order of optimization,
134+
- On x86-64, `generic`, `avx2`, 'avx512', and `avx512_spr` (for avx512 features available since Intel(R) Sapphire Rapids), by increasing order of optimization,
135135
- On aarch64, `generic` and `sve`, by increasing order of optimization,
136136
- `-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`).
137137
- BLAS-related options:
@@ -180,6 +180,12 @@ For AVX512:
180180
$ make -C build -j faiss_avx512
181181
```
182182

183+
For AVX512 features available since Intel(R) Sapphire Rapids.
184+
185+
``` shell
186+
$ make -C build -j faiss_avx512_spr
187+
```
188+
183189
This will ensure the creation of neccesary files when building and installing the python package.
184190

185191
## Step 3: Building the python bindings (optional)

cmake/link_to_faiss_lib.cmake

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
# LICENSE file in the root directory of this source tree.
66

77
function(link_to_faiss_lib target)
8-
if(NOT FAISS_OPT_LEVEL STREQUAL "avx2" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512" AND NOT FAISS_OPT_LEVEL STREQUAL "sve")
8+
if(NOT FAISS_OPT_LEVEL STREQUAL "avx2" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512_spr" AND NOT FAISS_OPT_LEVEL STREQUAL "sve")
99
target_link_libraries(${target} PRIVATE faiss)
1010
endif()
1111

@@ -27,6 +27,17 @@ function(link_to_faiss_lib target)
2727
target_link_libraries(${target} PRIVATE faiss_avx512)
2828
endif()
2929

30+
if(FAISS_OPT_LEVEL STREQUAL "avx512_spr")
31+
if(NOT WIN32)
32+
# Architecture mode to support AVX512 extensions available since Intel (R) Sapphire Rapids.
33+
# Ref: https://networkbuilders.intel.com/solutionslibrary/intel-avx-512-fp16-instruction-set-for-intel-xeon-processor-based-products-technology-guide
34+
target_compile_options(${target} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-march=sapphirerapids -mtune=sapphirerapids>)
35+
else()
36+
target_compile_options(${target} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/arch:AVX512>)
37+
endif()
38+
target_link_libraries(${target} PRIVATE faiss_avx512_spr)
39+
endif()
40+
3041
if(FAISS_OPT_LEVEL STREQUAL "sve")
3142
if(NOT WIN32)
3243
if("${CMAKE_CXX_FLAGS} ${CMAKE_CXX_FLAGS_DEBUG} " MATCHES "(^| )-march=native")

faiss/CMakeLists.txt

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,7 @@ set(FAISS_HEADERS ${FAISS_HEADERS} PARENT_SCOPE)
243243
add_library(faiss ${FAISS_SRC})
244244

245245
add_library(faiss_avx2 ${FAISS_SRC})
246-
if(NOT FAISS_OPT_LEVEL STREQUAL "avx2" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512")
246+
if(NOT FAISS_OPT_LEVEL STREQUAL "avx2" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512_spr")
247247
set_target_properties(faiss_avx2 PROPERTIES EXCLUDE_FROM_ALL TRUE)
248248
endif()
249249
if(NOT WIN32)
@@ -272,6 +272,20 @@ else()
272272
add_compile_options(/bigobj)
273273
endif()
274274

275+
add_library(faiss_avx512_spr ${FAISS_SRC})
276+
if(NOT FAISS_OPT_LEVEL STREQUAL "avx512_spr")
277+
set_target_properties(faiss_avx512_spr PROPERTIES EXCLUDE_FROM_ALL TRUE)
278+
endif()
279+
if(NOT WIN32)
280+
# Architecture mode to support AVX512 extensions available since Intel(R) Sapphire Rapids.
281+
# Ref: https://networkbuilders.intel.com/solutionslibrary/intel-avx-512-fp16-instruction-set-for-intel-xeon-processor-based-products-technology-guide
282+
target_compile_options(faiss_avx512_spr PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-march=sapphirerapids -mtune=sapphirerapids>)
283+
else()
284+
target_compile_options(faiss_avx512_spr PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/arch:AVX512>)
285+
# we need bigobj for the swig wrapper
286+
add_compile_options(/bigobj)
287+
endif()
288+
275289
add_library(faiss_sve ${FAISS_SRC})
276290
if(NOT FAISS_OPT_LEVEL STREQUAL "sve")
277291
set_target_properties(faiss_sve PROPERTIES EXCLUDE_FROM_ALL TRUE)
@@ -307,10 +321,13 @@ target_include_directories(faiss_avx2 PUBLIC
307321
target_include_directories(faiss_avx512 PUBLIC
308322
$<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>)
309323
# Handle `#include <faiss/foo.h>`.
324+
target_include_directories(faiss_avx512_spr PUBLIC
325+
$<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>)
326+
# Handle `#include <faiss/foo.h>`.
310327
target_include_directories(faiss_sve PUBLIC
311328
$<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>)
312329

313-
set_target_properties(faiss faiss_avx2 faiss_avx512 faiss_sve PROPERTIES
330+
set_target_properties(faiss faiss_avx2 faiss_avx512 faiss_avx512_spr faiss_sve PROPERTIES
314331
POSITION_INDEPENDENT_CODE ON
315332
WINDOWS_EXPORT_ALL_SYMBOLS ON
316333
)
@@ -319,6 +336,7 @@ if(WIN32)
319336
target_compile_definitions(faiss PRIVATE FAISS_MAIN_LIB)
320337
target_compile_definitions(faiss_avx2 PRIVATE FAISS_MAIN_LIB)
321338
target_compile_definitions(faiss_avx512 PRIVATE FAISS_MAIN_LIB)
339+
target_compile_definitions(faiss_avx512_spr PRIVATE FAISS_MAIN_LIB)
322340
target_compile_definitions(faiss_sve PRIVATE FAISS_MAIN_LIB)
323341
endif()
324342

@@ -328,6 +346,7 @@ if (${finteger_idx} EQUAL -1)
328346
endif()
329347
target_compile_definitions(faiss_avx2 PRIVATE FINTEGER=int)
330348
target_compile_definitions(faiss_avx512 PRIVATE FINTEGER=int)
349+
target_compile_definitions(faiss_avx512_spr PRIVATE FINTEGER=int)
331350
target_compile_definitions(faiss_sve PRIVATE FINTEGER=int)
332351

333352
if(FAISS_USE_LTO)
@@ -339,6 +358,7 @@ if(FAISS_USE_LTO)
339358
set_property(TARGET faiss PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
340359
set_property(TARGET faiss_avx2 PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
341360
set_property(TARGET faiss_avx512 PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
361+
set_property(TARGET faiss_avx512_spr PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
342362
else()
343363
message(STATUS "LTO not supported: <${ipo_error}>")
344364
endif()
@@ -348,24 +368,28 @@ find_package(OpenMP REQUIRED)
348368
target_link_libraries(faiss PRIVATE OpenMP::OpenMP_CXX)
349369
target_link_libraries(faiss_avx2 PRIVATE OpenMP::OpenMP_CXX)
350370
target_link_libraries(faiss_avx512 PRIVATE OpenMP::OpenMP_CXX)
371+
target_link_libraries(faiss_avx512_spr PRIVATE OpenMP::OpenMP_CXX)
351372
target_link_libraries(faiss_sve PRIVATE OpenMP::OpenMP_CXX)
352373

353374
find_package(MKL)
354375
if(MKL_FOUND)
355376
target_link_libraries(faiss PRIVATE ${MKL_LIBRARIES})
356377
target_link_libraries(faiss_avx2 PRIVATE ${MKL_LIBRARIES})
357378
target_link_libraries(faiss_avx512 PRIVATE ${MKL_LIBRARIES})
379+
target_link_libraries(faiss_avx512_spr PRIVATE ${MKL_LIBRARIES})
358380
else()
359381
find_package(BLAS REQUIRED)
360382
target_link_libraries(faiss PRIVATE ${BLAS_LIBRARIES})
361383
target_link_libraries(faiss_avx2 PRIVATE ${BLAS_LIBRARIES})
362384
target_link_libraries(faiss_avx512 PRIVATE ${BLAS_LIBRARIES})
385+
target_link_libraries(faiss_avx512_spr PRIVATE ${BLAS_LIBRARIES})
363386
target_link_libraries(faiss_sve PRIVATE ${BLAS_LIBRARIES})
364387

365388
find_package(LAPACK REQUIRED)
366389
target_link_libraries(faiss PRIVATE ${LAPACK_LIBRARIES})
367390
target_link_libraries(faiss_avx2 PRIVATE ${LAPACK_LIBRARIES})
368391
target_link_libraries(faiss_avx512 PRIVATE ${LAPACK_LIBRARIES})
392+
target_link_libraries(faiss_avx512_spr PRIVATE ${LAPACK_LIBRARIES})
369393
target_link_libraries(faiss_sve PRIVATE ${LAPACK_LIBRARIES})
370394
endif()
371395

@@ -390,6 +414,13 @@ if(FAISS_OPT_LEVEL STREQUAL "avx512")
390414
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
391415
)
392416
endif()
417+
if(FAISS_OPT_LEVEL STREQUAL "avx512_spr")
418+
install(TARGETS faiss_avx2 faiss_avx512 faiss_avx512_spr
419+
EXPORT faiss-targets
420+
ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
421+
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
422+
)
423+
endif()
393424
if(FAISS_OPT_LEVEL STREQUAL "sve")
394425
install(TARGETS faiss_sve
395426
EXPORT faiss-targets

faiss/gpu/CMakeLists.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,7 @@ if(FAISS_ENABLE_CUVS)
270270
target_compile_definitions(faiss PUBLIC USE_NVIDIA_CUVS=1)
271271
target_compile_definitions(faiss_avx2 PUBLIC USE_NVIDIA_CUVS=1)
272272
target_compile_definitions(faiss_avx512 PUBLIC USE_NVIDIA_CUVS=1)
273+
target_compile_definitions(faiss_avx512_spr PUBLIC USE_NVIDIA_CUVS=1)
273274

274275
# Mark all functions as hidden so that we don't generate
275276
# global 'public' functions that also exist in libraft.so
@@ -305,6 +306,7 @@ set(FAISS_GPU_HEADERS ${FAISS_GPU_HEADERS} PARENT_SCOPE)
305306
target_link_libraries(faiss PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,faiss_gpu>")
306307
target_link_libraries(faiss_avx2 PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,faiss_gpu>")
307308
target_link_libraries(faiss_avx512 PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,faiss_gpu>")
309+
target_link_libraries(faiss_avx512_spr PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,faiss_gpu>")
308310
target_link_libraries(faiss_sve PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,faiss_gpu>")
309311

310312
foreach(header ${FAISS_GPU_HEADERS})

faiss/python/CMakeLists.txt

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,11 +54,13 @@ endmacro()
5454
# we duplicate the source in order to override the module name.
5555
configure_file(swigfaiss.swig ${CMAKE_CURRENT_SOURCE_DIR}/swigfaiss_avx2.swig COPYONLY)
5656
configure_file(swigfaiss.swig ${CMAKE_CURRENT_SOURCE_DIR}/swigfaiss_avx512.swig COPYONLY)
57+
configure_file(swigfaiss.swig ${CMAKE_CURRENT_SOURCE_DIR}/swigfaiss_avx512_spr.swig COPYONLY)
5758
configure_file(swigfaiss.swig ${CMAKE_CURRENT_SOURCE_DIR}/swigfaiss_sve.swig COPYONLY)
5859

5960
configure_swigfaiss(swigfaiss.swig)
6061
configure_swigfaiss(swigfaiss_avx2.swig)
6162
configure_swigfaiss(swigfaiss_avx512.swig)
63+
configure_swigfaiss(swigfaiss_avx512_spr.swig)
6264
configure_swigfaiss(swigfaiss_sve.swig)
6365
configure_swigfaiss(faiss_example_external_module.swig)
6466

@@ -72,6 +74,8 @@ if(TARGET faiss)
7274
"${faiss_SOURCE_DIR}/faiss/${h}")
7375
list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
7476
"${faiss_SOURCE_DIR}/faiss/${h}")
77+
list(APPEND SWIG_MODULE_swigfaiss_avx512_spr_EXTRA_DEPS
78+
"${faiss_SOURCE_DIR}/faiss/${h}")
7579
list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS
7680
"${faiss_SOURCE_DIR}/faiss/${h}")
7781
list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
@@ -85,6 +89,8 @@ if(TARGET faiss)
8589
"${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
8690
list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
8791
"${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
92+
list(APPEND SWIG_MODULE_swigfaiss_avx512_spr_EXTRA_DEPS
93+
"${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
8894
list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
8995
"${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
9096
endforeach()
@@ -96,6 +102,8 @@ if(TARGET faiss)
96102
"${faiss_SOURCE_DIR}/faiss/gpu/${h}")
97103
list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
98104
"${faiss_SOURCE_DIR}/faiss/gpu/${h}")
105+
list(APPEND SWIG_MODULE_swigfaiss_avx512_spr_EXTRA_DEPS
106+
"${faiss_SOURCE_DIR}/faiss/gpu/${h}")
99107
list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS
100108
"${faiss_SOURCE_DIR}/faiss/gpu/${h}")
101109
list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
@@ -146,6 +154,18 @@ if(NOT FAISS_OPT_LEVEL STREQUAL "avx512")
146154
set_target_properties(swigfaiss_avx512 PROPERTIES EXCLUDE_FROM_ALL TRUE)
147155
endif()
148156

157+
set_property(SOURCE swigfaiss_avx512_spr.swig
158+
PROPERTY SWIG_MODULE_NAME swigfaiss_avx512_spr)
159+
swig_add_library(swigfaiss_avx512_spr
160+
TYPE SHARED
161+
LANGUAGE python
162+
SOURCES swigfaiss_avx512_spr.swig
163+
)
164+
set_property(TARGET swigfaiss_avx512_spr PROPERTY SWIG_COMPILE_OPTIONS -doxygen)
165+
if(NOT FAISS_OPT_LEVEL STREQUAL "avx512_spr")
166+
set_target_properties(swigfaiss_avx512_spr PROPERTIES EXCLUDE_FROM_ALL TRUE)
167+
endif()
168+
149169
set_property(SOURCE swigfaiss_sve.swig
150170
PROPERTY SWIG_MODULE_NAME swigfaiss_sve)
151171
swig_add_library(swigfaiss_sve
@@ -172,13 +192,15 @@ if(NOT WIN32)
172192
set_target_properties(swigfaiss PROPERTIES SUFFIX .so)
173193
set_target_properties(swigfaiss_avx2 PROPERTIES SUFFIX .so)
174194
set_target_properties(swigfaiss_avx512 PROPERTIES SUFFIX .so)
195+
set_target_properties(swigfaiss_avx512_spr PROPERTIES SUFFIX .so)
175196
set_target_properties(swigfaiss_sve PROPERTIES SUFFIX .so)
176197
set_target_properties(faiss_example_external_module PROPERTIES SUFFIX .so)
177198
else()
178199
# we need bigobj for the swig wrapper
179200
target_compile_options(swigfaiss PRIVATE /bigobj)
180201
target_compile_options(swigfaiss_avx2 PRIVATE /bigobj)
181202
target_compile_options(swigfaiss_avx512 PRIVATE /bigobj)
203+
target_compile_options(swigfaiss_avx512_spr PRIVATE /bigobj)
182204
target_compile_options(swigfaiss_sve PRIVATE /bigobj)
183205
target_compile_options(faiss_example_external_module PRIVATE /bigobj)
184206
endif()
@@ -188,6 +210,7 @@ if(FAISS_ENABLE_GPU)
188210
target_link_libraries(swigfaiss PRIVATE hip::host)
189211
target_link_libraries(swigfaiss_avx2 PRIVATE hip::host)
190212
target_link_libraries(swigfaiss_avx512 PRIVATE hip::host)
213+
target_link_libraries(swigfaiss_avx512_spr PRIVATE hip::host)
191214
target_link_libraries(faiss_example_external_module PRIVATE hip::host)
192215
else()
193216
find_package(CUDAToolkit REQUIRED)
@@ -197,6 +220,7 @@ if(FAISS_ENABLE_GPU)
197220
target_link_libraries(swigfaiss PRIVATE CUDA::cudart $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs>)
198221
target_link_libraries(swigfaiss_avx2 PRIVATE CUDA::cudart $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs>)
199222
target_link_libraries(swigfaiss_avx512 PRIVATE CUDA::cudart $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs>)
223+
target_link_libraries(swigfaiss_avx512_spr PRIVATE CUDA::cudart $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs>)
200224
target_link_libraries(swigfaiss_sve PRIVATE CUDA::cudart $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs>)
201225
endif()
202226
endif()
@@ -224,6 +248,13 @@ target_link_libraries(swigfaiss_avx512 PRIVATE
224248
OpenMP::OpenMP_CXX
225249
)
226250

251+
target_link_libraries(swigfaiss_avx512_spr PRIVATE
252+
faiss_avx512_spr
253+
Python::Module
254+
Python::NumPy
255+
OpenMP::OpenMP_CXX
256+
)
257+
227258
target_link_libraries(swigfaiss_sve PRIVATE
228259
faiss_sve
229260
Python::Module
@@ -244,6 +275,7 @@ target_link_libraries(faiss_example_external_module PRIVATE
244275
target_include_directories(swigfaiss PRIVATE ${PROJECT_SOURCE_DIR}/../..)
245276
target_include_directories(swigfaiss_avx2 PRIVATE ${PROJECT_SOURCE_DIR}/../..)
246277
target_include_directories(swigfaiss_avx512 PRIVATE ${PROJECT_SOURCE_DIR}/../..)
278+
target_include_directories(swigfaiss_avx512_spr PRIVATE ${PROJECT_SOURCE_DIR}/../..)
247279
target_include_directories(swigfaiss_sve PRIVATE ${PROJECT_SOURCE_DIR}/../..)
248280
target_include_directories(faiss_example_external_module PRIVATE ${PROJECT_SOURCE_DIR}/../..)
249281

@@ -270,6 +302,7 @@ target_include_directories(faiss_python_callbacks PRIVATE ${Python_INCLUDE_DIRS}
270302
target_link_libraries(swigfaiss PRIVATE faiss_python_callbacks)
271303
target_link_libraries(swigfaiss_avx2 PRIVATE faiss_python_callbacks)
272304
target_link_libraries(swigfaiss_avx512 PRIVATE faiss_python_callbacks)
305+
target_link_libraries(swigfaiss_avx512_spr PRIVATE faiss_python_callbacks)
273306
target_link_libraries(swigfaiss_sve PRIVATE faiss_python_callbacks)
274307
target_link_libraries(faiss_example_external_module PRIVATE faiss_python_callbacks)
275308

faiss/python/loader.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,9 @@ def is_sve_supported():
6767
result.add("AVX2")
6868
if "avx512" in numpy.distutils.cpuinfo.cpu.info[0].get('flags', ""):
6969
result.add("AVX512")
70+
if "avx512_fp16" in numpy.distutils.cpuinfo.cpu.info[0].get('flags', ""):
71+
# avx512_fp16 is supported starting SPR
72+
result.add("AVX512_SPR")
7073
if is_sve_supported():
7174
result.add("SVE")
7275
for f in os.getenv("FAISS_DISABLE_CPU_FEATURES", "").split(", \t\n\r"):
@@ -92,6 +95,18 @@ def is_sve_supported():
9295
instruction_sets.add(opt_level)
9396

9497
loaded = False
98+
has_AVX512_SPR = any("AVX512_SPR" in x.upper() for x in instruction_sets)
99+
if has_AVX512_SPR:
100+
try:
101+
logger.info("Loading faiss with AVX512-SPR support.")
102+
from .swigfaiss_avx512_spr import *
103+
logger.info("Successfully loaded faiss with AVX512-SPR support.")
104+
loaded = True
105+
except ImportError as e:
106+
logger.info(f"Could not load library with AVX512-SPR support due to:\n{e!r}")
107+
# reset so that we load without AVX512 below
108+
loaded = False
109+
95110
has_AVX512 = any("AVX512" in x.upper() for x in instruction_sets)
96111
if has_AVX512:
97112
try:

0 commit comments

Comments
 (0)