Add a FAISS_OPT_LEVEL=avx512_spr build mode.

.github/workflows/CI.yml

@@ -75,8 +75,14 @@ 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 -Dnproc.count=`nproc` -Davx512_spr.enabled=true"
+            else
+              echo "avx512 available on system"
+              su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc`"
+            fi
           elif lscpu  | grep -i avx2
           then
             echo "avx2 available on system"

.github/workflows/backwards_compatibility_tests_workflow.yml

@@ -108,8 +108,14 @@ 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
-            echo "avx512 available on system"
-            ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -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"
+	      ./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`
+	    fi
           elif lscpu  | grep -i avx2
           then
             echo "avx2 available on system"

.github/workflows/test_security.yml

@@ -73,8 +73,14 @@ 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 -Dnproc.count=`nproc`" -Davx512_spr.enabled=true
+            else
+              echo "avx512 available on system"
+              su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc`"
+            fi
           elif lscpu  | grep -i avx2
           then
             echo "avx2 available on system"

CHANGELOG.md

@@ -19,6 +19,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]
 ### 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]
 - Allow method parameter override for training based indices (#2290) https://github.com/opensearch-project/k-NN/pull/2290]

DEVELOPER_GUIDE.md

@@ -285,21 +285,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

build.gradle

@@ -20,6 +20,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
@@ -334,6 +335,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}")
     if (Os.isFamily(Os.FAMILY_WINDOWS)) {

jni/cmake/init-faiss.cmake

@@ -89,9 +89,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)

scripts/build.sh

@@ -122,7 +122,7 @@ 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"
@@ -131,7 +131,10 @@ if [ "$PLATFORM" != "windows" ] && [ "$ARCHITECTURE" = "x64" ]; then
   ./gradlew :buildJniLib -Davx2.enabled=true -Davx512.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

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;

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<Boolean> 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<Setting<?>> 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()

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<Void>) () -> {
 
             // 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);

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<Boolean>) () -> {
                     Stream<String> linestream = Files.lines(Paths.get(fileName));

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";