From 90fdad4f9c810bd011bc3c78666dad71a2186221 Mon Sep 17 00:00:00 2001 From: mulugetam Date: Fri, 24 Jan 2025 13:49:37 -0700 Subject: [PATCH] Add a FAISS_OPT_LEVEL=avx512_spr build mode. (#2404) * Add avx512_spr option to FAISS_OPT_LEVEL. - avx512_spr enables advanced avx512 instructions available since Intel(R) Sapphire Rapids. Signed-off-by: Mulugeta Mammo * Fix documentation. Signed-off-by: Assane Diop * Update scripts/build.sh Co-authored-by: Naveen Tatikonda Signed-off-by: mulugetam * Update scripts/build.sh Co-authored-by: Naveen Tatikonda Signed-off-by: mulugetam * Fix formatting issues. Signed-off-by: Mulugeta Mammo * Update CHANGELOG.md for avx512_spr build mode. Signed-off-by: Mulugeta Mammo * Fix bugs in build options and avx512_spr flag testing. Signed-off-by: Mulugeta Mammo * Upgrade gcc version to 12.4. Signed-off-by: Mulugeta Mammo --------- Signed-off-by: Mulugeta Mammo Signed-off-by: Assane Diop Signed-off-by: mulugetam Signed-off-by: Naveen Tatikonda Co-authored-by: Assane Diop Co-authored-by: Naveen Tatikonda --- .github/workflows/CI.yml | 14 +++-- ...backwards_compatibility_tests_workflow.yml | 12 +++- .github/workflows/test_security.yml | 14 +++-- CHANGELOG.md | 1 + DEVELOPER_GUIDE.md | 37 ++++++++---- build.gradle | 2 + jni/cmake/init-faiss.cmake | 16 ++++- scripts/build.sh | 11 ++-- .../opensearch/knn/common/KNNConstants.java | 1 + .../org/opensearch/knn/index/KNNSettings.java | 22 +++++++ .../org/opensearch/knn/jni/FaissService.java | 9 ++- .../org/opensearch/knn/jni/PlatformUtils.java | 16 +++-- .../plugin-metadata/plugin-security.policy | 1 + .../opensearch/knn/jni/PlatformUtilTests.java | 59 +++++++++++++++++++ 14 files changed, 182 insertions(+), 33 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 7969a9d59..b65aea6e9 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -75,15 +75,21 @@ jobs: chown -R 1000:1000 `pwd` if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw then - echo "avx512 available on system" - su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc`" + if lscpu | grep -q "GenuineIntel" && lscpu | grep -i avx512_fp16 | grep -i avx512_bf16 | grep -i avx512_vpopcntdq + then + echo "the system is an Intel(R) Sapphire Rapids or a newer-generation processor" + su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx512_spr.enabled=true -Dnproc.count=`nproc`" + else + echo "avx512 available on system" + su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx512_spr.enabled=false -Dnproc.count=`nproc`" + fi elif lscpu | grep -i avx2 then echo "avx2 available on system" - su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc` -Davx512.enabled=false" + su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx512.enabled=false -Davx512_spr.enabled=false -Dnproc.count=`nproc`" else echo "avx512 and avx2 not available on system" - su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Dnproc.count=`nproc`" + su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Davx512_spr.enabled=false -Dnproc.count=`nproc`" fi diff --git a/.github/workflows/backwards_compatibility_tests_workflow.yml b/.github/workflows/backwards_compatibility_tests_workflow.yml index 5a90d5852..55d4960c7 100644 --- a/.github/workflows/backwards_compatibility_tests_workflow.yml +++ b/.github/workflows/backwards_compatibility_tests_workflow.yml @@ -108,15 +108,21 @@ jobs: echo "Running restart-upgrade backwards compatibility tests ..." if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw then + if lscpu | grep -q "GenuineIntel" && lscpu | grep -i avx512_fp16 | grep -i avx512_bf16 | grep -i avx512_vpopcntdq + then + echo "the system is an Intel(R) Sapphire Rapids or a newer-generation processor" + ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Dnproc.count=`nproc` -Davx512_spr.enabled=true + else echo "avx512 available on system" - ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Dnproc.count=`nproc` + ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Dnproc.count=`nproc` -Davx512_spr.enabled=false + fi elif lscpu | grep -i avx2 then echo "avx2 available on system" - ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Dnproc.count=`nproc` -Davx512.enabled=false + ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Davx512.enabled=false -Davx512_spr.enabled=false -Dnproc.count=`nproc` else echo "avx512 and avx2 not available on system" - ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Davx2.enabled=false -Davx512.enabled=false -Dsimd.enabled=false -Dnproc.count=`nproc` + ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Davx2.enabled=false -Davx512.enabled=false -Davx512_spr.enabled=false -Dsimd.enabled=false -Dnproc.count=`nproc` fi Rolling-Upgrade-BWCTests-k-NN: diff --git a/.github/workflows/test_security.yml b/.github/workflows/test_security.yml index c9017972b..4cfbbf47c 100644 --- a/.github/workflows/test_security.yml +++ b/.github/workflows/test_security.yml @@ -73,13 +73,19 @@ jobs: chown -R 1000:1000 `pwd` if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw then - echo "avx512 available on system" - su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc`" + if lscpu | grep -q "GenuineIntel" && lscpu | grep -i avx512_fp16 | grep -i avx512_bf16 | grep -i avx512_vpopcntdq + then + echo "the system is an Intel(R) Sapphire Rapids or a newer-generation processor" + su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx512_spr.enabled=true -Dnproc.count=`nproc`" + else + echo "avx512 available on system" + su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx512_spr.enabled=false -Dnproc.count=`nproc`" + fi elif lscpu | grep -i avx2 then echo "avx2 available on system" - su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc` -Davx512.enabled=false" + su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx512.enabled=false -Davx512_spr.enabled=false -Dnproc.count=`nproc`" else echo "avx512 and avx2 not available on system" - su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Dnproc.count=`nproc`" + su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Davx512_spr.enabled=false -Dnproc.count=`nproc`" fi diff --git a/CHANGELOG.md b/CHANGELOG.md index 6d6a2eb79..9b7d6214b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add Support for Multi Values in innerHit for Nested k-NN Fields in Lucene and FAISS (#2283)[https://github.com/opensearch-project/k-NN/pull/2283] - Add binary index support for Lucene engine. (#2292)[https://github.com/opensearch-project/k-NN/pull/2292] - Add expand_nested_docs Parameter support to NMSLIB engine (#2331)[https://github.com/opensearch-project/k-NN/pull/2331] +- Add a new build mode, `FAISS_OPT_LEVEL=avx512_spr`, which enables the use of advanced AVX-512 instructions introduced with Intel(R) Sapphire Rapids (#2404)[https://github.com/opensearch-project/k-NN/pull/2404] - Add cosine similarity support for faiss engine (#2376)[https://github.com/opensearch-project/k-NN/pull/2376] ### Enhancements - Introduced a writing layer in native engines where relies on the writing interface to process IO. (#2241)[https://github.com/opensearch-project/k-NN/pull/2241] diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 8df0dc17d..dc635a53c 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -297,21 +297,38 @@ make -j 4 ### Enable SIMD Optimization SIMD(Single Instruction/Multiple Data) Optimization is enabled by default on Linux and Mac which boosts the performance by enabling `AVX2` and `AVX512` on `x86 architecture` and `NEON` on `ARM64 architecture` where applicable while building the Faiss library. But to enable SIMD, -the underlying processor should support these capabilities (AVX512, AVX2 or NEON). It can be disabled by setting the parameter `avx2.enabled` to `false` and -`avx512.enabled` to `false`. If your processor supports `AVX512` or `AVX2`, they can be set by enabling the setting . By default, these values are enabled on -OpenSearch. Some exceptions: As of now, SIMD support is not supported on Windows OS, and AVX512 is not present on MAC systems due to hardware not supporting the -feature. +the underlying processor should support these capabilities (AVX512, AVX2 or NEON). It can be disabled by setting the parameter `avx2.enabled`, `avx512.enabled`, +and `avx512_spr.enabled` to `false`. If your processor supports `AVX512` or `AVX2`, they can be set by enabling the setting. On Intel(R) Sapphire Rapids and +newer-generation systems, enabling `avx512_spr` offers support for `AVX512-FP16` and other features. By default, these values are enabled on OpenSearch. +Some exceptions: As of now, SIMD support is not supported on Windows OS, and AVX512 is not present on MAC systems due to hardware not supporting the feature. ``` -# While building OpenSearch k-NN -./gradlew build -Davx2.enabled=true -Davx512.enabled=true +# if (system_supports_avx512_spr) generate_avx512_spr_binaries() +# else if (system_supports_avx512) generate_avx512_binaries() +# else if (system_supports_ avx2) generate_avx2_binaries() +# else() generate_generic_binaries() -# While running OpenSearch k-NN -./gradlew run -Davx2.enabled=true -Davx512.enabled=true +# generate avx2 binaries +./gradlew build -Davx2.enabled=true -Davx512.enabled=false -Davx512_spr.enabled=false + +# if (system_supports_avx512_spr) generate_avx512_spr_binaries() +# else if (system_supports_avx512) generate_avx512_binaries() +# else() generate_generic_binaries() +./gradlew build -Davx2.enabled=false -Davx512.enabled=true + +# if (system_supports_avx512_spr) generate_avx512_spr_binaries() +# else if (system_supports_avx2) generate_avx2_binaries() +# else() generate_generic_binaries() +./gradlew build -Davx512.enabled=false -Davx512_spr.enabled=true + +# if (system_supports_avx512) generate_avx512_binaries() +# else if (system_supports_avx2) generate_avx2_binaries() +# else() generate_generic_binaries() +./gradlew build -Davx512.enabled=true -Davx512_spr.enabled=false -# While building the JNI libraries +# similar logic applies for jni cd jni -cmake . -DAVX2_ENABLED=true -DAVX512_ENABLED=true +cmake . -DAVX2_ENABLED=true -DAVX512_ENABLED=true -DAVX512_SPR_ENABLED=true ``` ## Run OpenSearch k-NN diff --git a/build.gradle b/build.gradle index 9a3ca673f..512508022 100644 --- a/build.gradle +++ b/build.gradle @@ -22,6 +22,7 @@ buildscript { avx2_enabled = System.getProperty("avx2.enabled", "true") nproc_count = System.getProperty("nproc.count", "1") avx512_enabled = System.getProperty("avx512.enabled", "true") + avx512_spr_enabled = System.getProperty("avx512_spr.enabled", "true") // This flag determines whether the CMake build system should apply a custom patch. It prevents build failures // when the cmakeJniLib task is run multiple times. If the build.lib.commit_patches is true, the CMake build // system skips applying the patch if the patches have been applied already. If build.lib.commit_patches is @@ -336,6 +337,7 @@ task cmakeJniLib(type:Exec) { args.add("-DKNN_PLUGIN_VERSION=${opensearch_version}") args.add("-DAVX2_ENABLED=${avx2_enabled}") args.add("-DAVX512_ENABLED=${avx512_enabled}") + args.add("-DAVX512_SPR_ENABLED=${avx512_spr_enabled}") args.add("-DCOMMIT_LIB_PATCHES=${commit_lib_patches}") args.add("-DAPPLY_LIB_PATCHES=${apply_lib_patches}") def javaHome = Jvm.current().getJavaHome() diff --git a/jni/cmake/init-faiss.cmake b/jni/cmake/init-faiss.cmake index 8243c26aa..3cb90b767 100644 --- a/jni/cmake/init-faiss.cmake +++ b/jni/cmake/init-faiss.cmake @@ -108,9 +108,23 @@ if(NOT DEFINED AVX512_ENABLED) set(AVX512_ENABLED true) # set default value as true if the argument is not set endif() -if(${CMAKE_SYSTEM_NAME} STREQUAL Windows OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64" OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm64" OR ( NOT AVX2_ENABLED AND NOT AVX512_ENABLED)) +if(NOT DEFINED AVX512_SPR_ENABLED) + # Check if the system is Intel(R) Sapphire Rapids or a newer-generation processor + execute_process(COMMAND bash -c "lscpu | grep -q 'GenuineIntel' && lscpu | grep -i 'avx512_fp16' | grep -i 'avx512_bf16' | grep -i 'avx512_vpopcntdq'" OUTPUT_VARIABLE SPR_FLAGS OUTPUT_STRIP_TRAILING_WHITESPACE) + if (AND NOT "${SPR_FLAGS}" STREQUAL "") + set(AVX512_SPR_ENABLED true) + else() + set(AVX512_SPR_ENABLED false) + endif() +endif() + +if(${CMAKE_SYSTEM_NAME} STREQUAL Windows OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64" OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm64" OR ( NOT AVX2_ENABLED AND NOT AVX512_ENABLED AND NOT AVX512_SPR_ENABLED)) set(FAISS_OPT_LEVEL generic) # Keep optimization level as generic on Windows OS as it is not supported due to MINGW64 compiler issue. Also, on aarch64 avx2 is not supported. set(TARGET_LINK_FAISS_LIB faiss) +elseif(${CMAKE_SYSTEM_NAME} STREQUAL Linux AND AVX512_SPR_ENABLED) + set(FAISS_OPT_LEVEL avx512_spr) + set(TARGET_LINK_FAISS_LIB faiss_avx512_spr) + string(PREPEND LIB_EXT "_avx512_spr") elseif(${CMAKE_SYSTEM_NAME} STREQUAL Linux AND AVX512_ENABLED) set(FAISS_OPT_LEVEL avx512) # Keep optimization level as avx512 to improve performance on Linux. This is not present on mac systems, and presently not supported on Windows OS. set(TARGET_LINK_FAISS_LIB faiss_avx512) diff --git a/scripts/build.sh b/scripts/build.sh index 203b76c99..c0ad56b30 100755 --- a/scripts/build.sh +++ b/scripts/build.sh @@ -112,7 +112,7 @@ fi # https://github.com/opensearch-project/k-NN/issues/1138 # https://github.com/opensearch-project/opensearch-build/issues/4386 GCC_VERSION=`gcc --version | head -n 1 | cut -d ' ' -f3` -GCC_REQUIRED_VERSION=9.0.0 +GCC_REQUIRED_VERSION=12.4 COMPARE_VERSION=`echo $GCC_REQUIRED_VERSION $GCC_VERSION | tr ' ' '\n' | sort -V | uniq | head -n 1` if [ "$COMPARE_VERSION" != "$GCC_REQUIRED_VERSION" ]; then echo "gcc version on this env is older than $GCC_REQUIRED_VERSION, exit 1" @@ -122,16 +122,19 @@ fi # Build k-NN lib and plugin through gradle tasks cd $work_dir ./gradlew build --no-daemon --refresh-dependencies -x integTest -x test -Dopensearch.version=$VERSION -Dbuild.snapshot=$SNAPSHOT -Dbuild.version_qualifier=$QUALIFIER -Dbuild.lib.commit_patches=false -./gradlew :buildJniLib -Davx512.enabled=false -Davx2.enabled=false -Dbuild.lib.commit_patches=false -Dnproc.count=${NPROC_COUNT:-1} +./gradlew :buildJniLib -Davx512.enabled=false -Davx512_spr.enabled=false -Davx2.enabled=false -Dbuild.lib.commit_patches=false -Dnproc.count=${NPROC_COUNT:-1} if [ "$PLATFORM" != "windows" ] && [ "$ARCHITECTURE" = "x64" ]; then echo "Building k-NN library after enabling AVX2" # Skip applying patches as patches were applied already from previous :buildJniLib task # If we apply patches again, it fails with conflict - ./gradlew :buildJniLib -Davx2.enabled=true -Davx512.enabled=false -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false + ./gradlew :buildJniLib -Davx2.enabled=true -Davx512.enabled=false -Davx512_spr.enabled=false -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false echo "Building k-NN library after enabling AVX512" - ./gradlew :buildJniLib -Davx512.enabled=true -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false + ./gradlew :buildJniLib -Davx512.enabled=true -Davx512_spr.enabled=false -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false + + echo "Building k-NN library after enabling AVX512_SPR" + ./gradlew :buildJniLib -Davx512_spr.enabled=true -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false fi ./gradlew publishPluginZipPublicationToZipStagingRepository -Dopensearch.version=$VERSION -Dbuild.snapshot=$SNAPSHOT -Dbuild.version_qualifier=$QUALIFIER diff --git a/src/main/java/org/opensearch/knn/common/KNNConstants.java b/src/main/java/org/opensearch/knn/common/KNNConstants.java index ce6095fd0..4479099e8 100644 --- a/src/main/java/org/opensearch/knn/common/KNNConstants.java +++ b/src/main/java/org/opensearch/knn/common/KNNConstants.java @@ -142,6 +142,7 @@ public class KNNConstants { public static final String FAISS_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + FAISS_NAME; public static final String FAISS_AVX2_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + FAISS_NAME + "_avx2"; public static final String FAISS_AVX512_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + FAISS_NAME + "_avx512"; + public static final String FAISS_AVX512_SPR_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + FAISS_NAME + "_avx512_spr"; public static final String NMSLIB_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + NMSLIB_NAME; public static final String COMMON_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + COMMONS_NAME; diff --git a/src/main/java/org/opensearch/knn/index/KNNSettings.java b/src/main/java/org/opensearch/knn/index/KNNSettings.java index 097329d81..d60b500c1 100644 --- a/src/main/java/org/opensearch/knn/index/KNNSettings.java +++ b/src/main/java/org/opensearch/knn/index/KNNSettings.java @@ -90,6 +90,7 @@ public class KNNSettings { public static final String QUANTIZATION_STATE_CACHE_SIZE_LIMIT = "knn.quantization.cache.size.limit"; public static final String QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES = "knn.quantization.cache.expiry.minutes"; public static final String KNN_FAISS_AVX512_DISABLED = "knn.faiss.avx512.disabled"; + public static final String KNN_FAISS_AVX512_SPR_DISABLED = "knn.faiss.avx512_spr.disabled"; public static final String KNN_DISK_VECTOR_SHARD_LEVEL_RESCORING_DISABLED = "index.knn.disk.vector.shard_level_rescoring_disabled"; /** @@ -98,6 +99,7 @@ public class KNNSettings { */ public static final boolean KNN_DEFAULT_FAISS_AVX2_DISABLED_VALUE = false; public static final boolean KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE = false; + public static final boolean KNN_DEFAULT_FAISS_AVX512_SPR_DISABLED_VALUE = false; public static final String INDEX_KNN_DEFAULT_SPACE_TYPE = "l2"; public static final Integer INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD_DEFAULT_VALUE = 15_000; public static final Integer INDEX_KNN_BUILD_VECTOR_DATA_STRUCTURE_THRESHOLD_MIN = -1; @@ -353,6 +355,12 @@ public class KNNSettings { NodeScope ); + public static final Setting KNN_FAISS_AVX512_SPR_DISABLED_SETTING = Setting.boolSetting( + KNN_FAISS_AVX512_SPR_DISABLED, + KNN_DEFAULT_FAISS_AVX512_SPR_DISABLED_VALUE, + NodeScope + ); + /** * Dynamic settings */ @@ -484,6 +492,10 @@ private Setting getSetting(String key) { return KNN_FAISS_AVX512_DISABLED_SETTING; } + if (KNN_FAISS_AVX512_SPR_DISABLED.equals(key)) { + return KNN_FAISS_AVX512_SPR_DISABLED_SETTING; + } + if (KNN_VECTOR_STREAMING_MEMORY_LIMIT_IN_MB.equals(key)) { return KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING; } @@ -521,6 +533,7 @@ public List> getSettings() { KNN_FAISS_AVX2_DISABLED_SETTING, KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING, KNN_FAISS_AVX512_DISABLED_SETTING, + KNN_FAISS_AVX512_SPR_DISABLED_SETTING, QUANTIZATION_STATE_CACHE_SIZE_LIMIT_SETTING, QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES_SETTING, KNN_DISK_VECTOR_SHARD_LEVEL_RESCORING_DISABLED_SETTING @@ -570,6 +583,15 @@ public static boolean isFaissAVX512Disabled() { ); } + public static boolean isFaissAVX512SPRDisabled() { + return Booleans.parseBoolean( + Objects.requireNonNullElse( + KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX512_SPR_DISABLED), + KNN_DEFAULT_FAISS_AVX512_SPR_DISABLED_VALUE + ).toString() + ); + } + public static Integer getFilteredExactSearchThreshold(final String indexName) { return KNNSettings.state().clusterService.state() .getMetadata() diff --git a/src/main/java/org/opensearch/knn/jni/FaissService.java b/src/main/java/org/opensearch/knn/jni/FaissService.java index dcc7b180d..3ae8bbb92 100644 --- a/src/main/java/org/opensearch/knn/jni/FaissService.java +++ b/src/main/java/org/opensearch/knn/jni/FaissService.java @@ -23,8 +23,10 @@ import static org.opensearch.knn.index.KNNSettings.isFaissAVX2Disabled; import static org.opensearch.knn.index.KNNSettings.isFaissAVX512Disabled; +import static org.opensearch.knn.index.KNNSettings.isFaissAVX512SPRDisabled; import static org.opensearch.knn.jni.PlatformUtils.isAVX2SupportedBySystem; import static org.opensearch.knn.jni.PlatformUtils.isAVX512SupportedBySystem; +import static org.opensearch.knn.jni.PlatformUtils.isAVX512SPRSupportedBySystem; /** * Service to interact with faiss jni layer. Class dependencies should be minimal @@ -40,8 +42,11 @@ class FaissService { AccessController.doPrivileged((PrivilegedAction) () -> { // Even if the underlying system supports AVX512 and AVX2, users can override and disable it by setting - // 'knn.faiss.avx2.disabled' or 'knn.faiss.avx512.disabled' to true in the opensearch.yml configuration - if (!isFaissAVX512Disabled() && isAVX512SupportedBySystem()) { + // 'knn.faiss.avx2.disabled', 'knn.faiss.avx512.disabled', or 'knn.faiss.avx512_spr.disabled' to true in the opensearch.yml + // configuration + if (!isFaissAVX512SPRDisabled() && isAVX512SPRSupportedBySystem()) { + System.loadLibrary(KNNConstants.FAISS_AVX512_SPR_JNI_LIBRARY_NAME); + } else if (!isFaissAVX512Disabled() && isAVX512SupportedBySystem()) { System.loadLibrary(KNNConstants.FAISS_AVX512_JNI_LIBRARY_NAME); } else if (!isFaissAVX2Disabled() && isAVX2SupportedBySystem()) { System.loadLibrary(KNNConstants.FAISS_AVX2_JNI_LIBRARY_NAME); diff --git a/src/main/java/org/opensearch/knn/jni/PlatformUtils.java b/src/main/java/org/opensearch/knn/jni/PlatformUtils.java index 445862f24..a67a88487 100644 --- a/src/main/java/org/opensearch/knn/jni/PlatformUtils.java +++ b/src/main/java/org/opensearch/knn/jni/PlatformUtils.java @@ -84,6 +84,17 @@ public static boolean isAVX2SupportedBySystem() { } public static boolean isAVX512SupportedBySystem() { + return areAVX512FlagsAvailable(new String[] { "avx512f", "avx512cd", "avx512vl", "avx512dq", "avx512bw" }); + } + + public static boolean isAVX512SPRSupportedBySystem() { + return areAVX512FlagsAvailable(new String[] { "avx512_fp16", "avx512_bf16", "avx512_vpopcntdq" }); + } + + private static boolean areAVX512FlagsAvailable(String[] avx512) { + // AVX512 has multiple flags, which control various features. k-nn requires the same set of flags as faiss to compile + // using avx512. Please update these if faiss updates their compilation instructions in the future. + // https://github.com/facebookresearch/faiss/blob/main/faiss/CMakeLists.txt if (!Platform.isIntel() || Platform.isMac() || Platform.isWindows()) { return false; @@ -98,11 +109,6 @@ public static boolean isAVX512SupportedBySystem() { // supports AVX512 instructions supported by faiss. String fileName = "/proc/cpuinfo"; - // AVX512 has multiple flags, which control various features. k-nn requires the same set of flags as faiss to compile - // using avx512. Please update these if faiss updates their compilation instructions in the future. - // https://github.com/facebookresearch/faiss/blob/main/faiss/CMakeLists.txt - String[] avx512 = { "avx512f", "avx512cd", "avx512vl", "avx512dq", "avx512bw" }; - try { return AccessController.doPrivileged((PrivilegedExceptionAction) () -> { Stream linestream = Files.lines(Paths.get(fileName)); diff --git a/src/main/plugin-metadata/plugin-security.policy b/src/main/plugin-metadata/plugin-security.policy index ed329740f..7ebf8f137 100644 --- a/src/main/plugin-metadata/plugin-security.policy +++ b/src/main/plugin-metadata/plugin-security.policy @@ -4,6 +4,7 @@ grant { permission java.lang.RuntimePermission "loadLibrary.opensearchknn_common"; permission java.lang.RuntimePermission "loadLibrary.opensearchknn_faiss_avx2"; permission java.lang.RuntimePermission "loadLibrary.opensearchknn_faiss_avx512"; + permission java.lang.RuntimePermission "loadLibrary.opensearchknn_faiss_avx512_spr"; permission java.net.SocketPermission "*", "connect,resolve"; permission java.lang.RuntimePermission "accessDeclaredMembers"; permission java.io.FilePermission "/proc/cpuinfo", "read"; diff --git a/src/test/java/org/opensearch/knn/jni/PlatformUtilTests.java b/src/test/java/org/opensearch/knn/jni/PlatformUtilTests.java index 19c0abb07..c524d211d 100644 --- a/src/test/java/org/opensearch/knn/jni/PlatformUtilTests.java +++ b/src/test/java/org/opensearch/knn/jni/PlatformUtilTests.java @@ -23,6 +23,7 @@ import static org.mockito.Mockito.mockStatic; import static org.opensearch.knn.jni.PlatformUtils.isAVX2SupportedBySystem; import static org.opensearch.knn.jni.PlatformUtils.isAVX512SupportedBySystem; +import static org.opensearch.knn.jni.PlatformUtils.isAVX512SPRSupportedBySystem; public class PlatformUtilTests extends KNNTestCase { public static final String MAC_CPU_FEATURES = "machdep.cpu.leaf7_features"; @@ -182,4 +183,62 @@ public void testIsAVX512SupportedBySystem_platformIsLinuxSomeAVX512FlagsPresent_ } } } + + // Tests AVX512 instructions available since Intel(R) Sapphire Rapids. + + public void testIsAVX512SPRSupportedBySystem_platformIsNotIntel_returnsFalse() { + try (MockedStatic mockedPlatform = mockStatic(Platform.class)) { + mockedPlatform.when(Platform::isIntel).thenReturn(false); + assertFalse(isAVX512SPRSupportedBySystem()); + } + } + + public void testIsAVX512SPRSupportedBySystem_platformIsMac_returnsFalse() { + try (MockedStatic mockedPlatform = mockStatic(Platform.class)) { + mockedPlatform.when(Platform::isMac).thenReturn(false); + assertFalse(isAVX512SPRSupportedBySystem()); + } + } + + public void testIsAVX512SPRSupportedBySystem_platformIsIntelMac_returnsFalse() { + try (MockedStatic mockedPlatform = mockStatic(Platform.class)) { + mockedPlatform.when(Platform::isIntel).thenReturn(true); + mockedPlatform.when(Platform::isMac).thenReturn(true); + assertFalse(isAVX512SPRSupportedBySystem()); + } + } + + public void testIsAVX512SPRSupportedBySystem_platformIsIntelWithOSAsWindows_returnsFalse() { + try (MockedStatic mockedPlatform = mockStatic(Platform.class)) { + mockedPlatform.when(Platform::isIntel).thenReturn(true); + mockedPlatform.when(Platform::isWindows).thenReturn(true); + assertFalse(isAVX512SPRSupportedBySystem()); + } + } + + public void testIsAVX512SPRSupportedBySystem_platformIsLinuxAllAVX512SPRFlagsPresent_returnsTrue() { + try (MockedStatic mockedPlatform = mockStatic(Platform.class)) { + mockedPlatform.when(Platform::isIntel).thenReturn(true); + mockedPlatform.when(Platform::isLinux).thenReturn(true); + + try (MockedStatic mockedFiles = mockStatic(Files.class)) { + mockedFiles.when(() -> Files.lines(Paths.get(LINUX_PROC_CPU_INFO))) + .thenReturn(Stream.of("flags: avx512_fp16 avx512_bf16 avx512_vpopcntdq", "dummy string")); + assertTrue(isAVX512SPRSupportedBySystem()); + } + } + } + + public void testIsAVX512SPRSupportedBySystem_platformIsLinuxSomeAVX512SPRFlagsPresent_returnsFalse() { + try (MockedStatic mockedPlatform = mockStatic(Platform.class)) { + mockedPlatform.when(Platform::isIntel).thenReturn(true); + mockedPlatform.when(Platform::isLinux).thenReturn(true); + + try (MockedStatic mockedFiles = mockStatic(Files.class)) { + mockedFiles.when(() -> Files.lines(Paths.get(LINUX_PROC_CPU_INFO))) + .thenReturn(Stream.of("flags: avx512_fp16 avx512_vpopcntdq", "dummy string")); + assertFalse(isAVX512SPRSupportedBySystem()); + } + } + } }