Add changes for AVX-512 support in k-NN.

.github/workflows/CI.yml

@@ -73,13 +73,17 @@ jobs:
         # switching the user, as OpenSearch cluster can only be started as root/Administrator on linux-deb/linux-rpm/windows-zip.
         run: |
           chown -R 1000:1000 `pwd`
-          if lscpu  | grep -i avx2
+          if lscpu  | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw          
           then
-            echo "avx2 available on system"
+            echo "avx512 available on system"
             su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc`"
+          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"
           else
-            echo "avx2 not available on system"
-            su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dsimd.enabled=false -Dnproc.count=`nproc`"
+            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`"
           fi  
 
 
@@ -126,7 +130,7 @@ jobs:
               ./gradlew build -Dnproc.count=3
           else
               echo "avx2 not available on system"
-              ./gradlew build -Dsimd.enabled=false -Dnproc.count=3
+              ./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Dnproc.count=3
           fi
 
   Build-k-NN-Windows:
@@ -187,4 +191,4 @@ jobs:
       # TODO: Detect processor count and set the value of nproc.count
       - name: Run build
         run: |
-          ./gradlew.bat build -D'simd.enabled=false' -D'nproc.count=4'
+          ./gradlew.bat build -D'avx2.enabled=false' -D'avx512.enabled=false' -D'nproc.count=4'

.github/workflows/backwards_compatibility_tests_workflow.yml

@@ -105,9 +105,19 @@ jobs:
       - if: startsWith(matrix.os,'ubuntu')
         name: Run k-NN Restart-Upgrade BWC Tests from BWCVersion-${{ matrix.bwc_version }} to OpenSearch Version-${{ matrix.opensearch_version }} on Ubuntu
         run: |
-          echo "Running restart-upgrade backwards compatibility tests ..."
-          ./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE
-
+          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`
+          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
+          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`
+          fi 
 
   Rolling-Upgrade-BWCTests-k-NN:
     strategy:
@@ -176,4 +186,15 @@ jobs:
         name: Run k-NN Rolling-Upgrade BWC Tests from BWCVersion-${{ matrix.bwc_version }} to OpenSearch Version-${{ matrix.opensearch_version }} on Ubuntu
         run: |
           echo "Running rolling-upgrade backwards compatibility tests ..."
-          ./gradlew :qa:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE
+          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:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE -Dnproc.count=`nproc`
+          elif lscpu  | grep -i avx2
+          then
+            echo "avx2 available on system"
+            ./gradlew :qa:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE -Dnproc.count=`nproc` -Davx512.enabled=false
+          else
+            echo "avx512 and avx2 not available on system"
+            ./gradlew :qa:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE -Davx2.enabled=false -Davx512.enabled=false -Dsimd.enabled=false -Dnproc.count=`nproc`
+          fi  

.github/workflows/test_security.yml

@@ -70,4 +70,15 @@ jobs:
         # switching the user, as OpenSearch cluster can only be started as root/Administrator on linux-deb/linux-rpm/windows-zip.
         run: |
           chown -R 1000:1000 `pwd`
-          su `id -un 1000` -c "whoami && java -version && ./gradlew integTest -Dsecurity.enabled=true -Dsimd.enabled=true -Dnproc.count=`nproc`"
+          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`"
+          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"
+          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`"
+          fi  

CHANGELOG.md

@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 2.x](https://github.com/opensearch-project/k-NN/compare/2.17...2.x)
 ### Features
+* Add AVX512 support to k-NN for FAISS library [#2069](https://github.com/opensearch-project/k-NN/pull/2069)
 ### Enhancements
 * Add short circuit if no live docs are in segments [#2059](https://github.com/opensearch-project/k-NN/pull/2059)
 ### Bug Fixes

DEVELOPER_GUIDE.md

@@ -283,19 +283,22 @@ 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` on `x86 architecture` and `NEON` on `ARM64 architecture` while building the Faiss library. But to enable SIMD, the underlying processor
-should support this (AVX2 or NEON). It can be disabled by setting the parameter `simd.enabled` to `false`. As of now, it is not supported on Windows OS.
+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.
 
 ```
 # While building OpenSearch k-NN
-./gradlew build -Dsimd.enabled=true
+./gradlew build -Davx2.enabled=true -Davx512.enabled=true
 
 # While running OpenSearch k-NN
-./gradlew run -Dsimd.enabled=true
+./gradlew run -Davx2.enabled=true -Davx512.enabled=true
 
 # While building the JNI libraries
 cd jni
-cmake . -DSIMD_ENABLED=true
+cmake . -DAVX2_ENABLED=true -DAVX512_ENABLED=true
 ```
 
 ## Run OpenSearch k-NN

build.gradle

@@ -17,8 +17,9 @@ buildscript {
         version_qualifier = System.getProperty("build.version_qualifier", "")
         opensearch_group = "org.opensearch"
         isSnapshot = "true" == System.getProperty("build.snapshot", "true")
-        simd_enabled = System.getProperty("simd.enabled", "true")
+        avx2_enabled = System.getProperty("avx2.enabled", "true")
         nproc_count = System.getProperty("nproc.count", "1")
+        avx512_enabled = System.getProperty("avx512.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
@@ -316,7 +317,8 @@ task cmakeJniLib(type:Exec) {
     args.add("cmake")
     args.add(".")
     args.add("-DKNN_PLUGIN_VERSION=${opensearch_version}")
-    args.add("-DSIMD_ENABLED=${simd_enabled}")
+    args.add("-DAVX2_ENABLED=${avx2_enabled}")
+    args.add("-DAVX512_ENABLED=${avx512_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

@@ -81,13 +81,21 @@ set(BUILD_TESTING OFF)          # Avoid building faiss tests
 set(FAISS_ENABLE_GPU OFF)
 set(FAISS_ENABLE_PYTHON OFF)
 
-if(NOT DEFINED SIMD_ENABLED)
-    set(SIMD_ENABLED true)   # set default value as true if the argument is not set
+if(NOT DEFINED AVX2_ENABLED)
+    set(AVX2_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 ${SIMD_ENABLED})
+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))
     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_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)
+    string(PREPEND LIB_EXT "_avx512") # Prepend "_avx512" to lib extension to create the library as "libopensearchknn_faiss_avx512.so" on linux
 else()
     set(FAISS_OPT_LEVEL avx2)       # Keep optimization level as avx2 to improve performance on Linux and Mac.
     set(TARGET_LINK_FAISS_LIB faiss_avx2)

scripts/build.sh

@@ -118,13 +118,16 @@ 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 -Dsimd.enabled=false -Dbuild.lib.commit_patches=false
+./gradlew :buildJniLib -Davx512.enabled=false -Davx2.enabled=false -Dbuild.lib.commit_patches=false
 
 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 -Dsimd.enabled=true -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false
+  ./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
 fi
 
 ./gradlew publishPluginZipPublicationToZipStagingRepository -Dopensearch.version=$VERSION -Dbuild.snapshot=$SNAPSHOT -Dbuild.version_qualifier=$QUALIFIER

src/main/java/org/opensearch/knn/common/KNNConstants.java

@@ -140,6 +140,7 @@ public class KNNConstants {
     private static final String JNI_LIBRARY_PREFIX = "opensearchknn_";
     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 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

@@ -86,11 +86,13 @@ public class KNNSettings {
     public static final String KNN_FAISS_AVX2_DISABLED = "knn.faiss.avx2.disabled";
     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";
 
     /**
      * Default setting values
      */
     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 String INDEX_KNN_DEFAULT_SPACE_TYPE = "l2";
     public static final String INDEX_KNN_DEFAULT_SPACE_TYPE_FOR_BINARY = "hamming";
     public static final Integer INDEX_KNN_DEFAULT_ALGO_PARAM_M = 16;
@@ -302,6 +304,12 @@ public class KNNSettings {
         Dynamic
     );
 
+    public static final Setting<Boolean> KNN_FAISS_AVX512_DISABLED_SETTING = Setting.boolSetting(
+        KNN_FAISS_AVX512_DISABLED,
+        KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE,
+        NodeScope
+    );
+
     /**
      * Dynamic settings
      */
@@ -429,6 +437,10 @@ private Setting<?> getSetting(String key) {
             return KNN_FAISS_AVX2_DISABLED_SETTING;
         }
 
+        if (KNN_FAISS_AVX512_DISABLED.equals(key)) {
+            return KNN_FAISS_AVX512_DISABLED_SETTING;
+        }
+
         if (KNN_VECTOR_STREAMING_MEMORY_LIMIT_IN_MB.equals(key)) {
             return KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING;
         }
@@ -460,6 +472,7 @@ public List<Setting<?>> getSettings() {
             ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD_SETTING,
             KNN_FAISS_AVX2_DISABLED_SETTING,
             KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING,
+            KNN_FAISS_AVX512_DISABLED_SETTING,
             QUANTIZATION_STATE_CACHE_SIZE_LIMIT_SETTING,
             QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES_SETTING
         );
@@ -499,6 +512,22 @@ public static boolean isFaissAVX2Disabled() {
         }
     }
 
+    public static boolean isFaissAVX512Disabled() {
+        try {
+            return KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX512_DISABLED);
+        } catch (Exception e) {
+            // In some UTs we identified that cluster setting is not set properly an leads to NPE. This check will avoid
+            // those cases and will still return the default value.
+            log.warn(
+                "Unable to get setting value {} from cluster settings. Using default value as {}",
+                KNN_FAISS_AVX512_DISABLED,
+                KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE,
+                e
+            );
+            return KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE;
+        }
+    }
+
     public static Integer getFilteredExactSearchThreshold(final String indexName) {
         return KNNSettings.state().clusterService.state()
             .getMetadata()

src/main/java/org/opensearch/knn/jni/FaissService.java

@@ -20,7 +20,9 @@
 import java.util.Map;
 
 import static org.opensearch.knn.index.KNNSettings.isFaissAVX2Disabled;
+import static org.opensearch.knn.index.KNNSettings.isFaissAVX512Disabled;
 import static org.opensearch.knn.jni.PlatformUtils.isAVX2SupportedBySystem;
+import static org.opensearch.knn.jni.PlatformUtils.isAVX512SupportedBySystem;;
 
 /**
  * Service to interact with faiss jni layer. Class dependencies should be minimal
@@ -30,14 +32,17 @@
  *      src/main/java/org/opensearch/knn/index/query/KNNQueryResult.java
  *      src/main/java/org/opensearch/knn/common/KNNConstants.java
  */
+@SuppressWarnings("deprecation")
 class FaissService {
 
     static {
         AccessController.doPrivileged((PrivilegedAction<Void>) () -> {
 
-            // Even if the underlying system supports AVX2, users can override and disable it by using the
-            // 'knn.faiss.avx2.disabled' setting by setting it to true in the opensearch.yml configuration
-            if (!isFaissAVX2Disabled() && isAVX2SupportedBySystem()) {
+            // 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()) {
+                System.loadLibrary(KNNConstants.FAISS_AVX512_JNI_LIBRARY_NAME);
+            } else if (!isFaissAVX2Disabled() && isAVX2SupportedBySystem()) {
                 System.loadLibrary(KNNConstants.FAISS_AVX2_JNI_LIBRARY_NAME);
             } else {
                 System.loadLibrary(KNNConstants.FAISS_JNI_LIBRARY_NAME);

src/main/java/org/opensearch/knn/jni/PlatformUtils.java

@@ -21,7 +21,9 @@
 import java.nio.file.Paths;
 import java.security.AccessController;
 import java.security.PrivilegedExceptionAction;
+import java.util.Arrays;
 import java.util.Locale;
+import java.util.stream.Stream;
 
 public class PlatformUtils {
 
@@ -37,7 +39,7 @@ public class PlatformUtils {
      *    the flags contains 'avx2' and return true if it exists else false.
      */
     public static boolean isAVX2SupportedBySystem() {
-        if (!Platform.isIntel()) {
+        if (!Platform.isIntel() || Platform.isWindows()) {
             return false;
         }
 
@@ -58,7 +60,6 @@ public static boolean isAVX2SupportedBySystem() {
             }
 
         } else if (Platform.isLinux()) {
-
             // The "/proc/cpuinfo" is a virtual file which identifies and provides the processor details used
             // by system. This info contains "flags" for each processor which determines the qualities of that processor
             // and it's ability to process different instruction sets like mmx, avx, avx2 and so on.
@@ -80,4 +81,38 @@ public static boolean isAVX2SupportedBySystem() {
         }
         return false;
     }
+
+    public static boolean isAVX512SupportedBySystem() {
+
+        if (!Platform.isIntel() || Platform.isMac() || Platform.isWindows()) {
+            return false;
+        }
+
+        if (Platform.isLinux()) {
+            // The "/proc/cpuinfo" is a virtual file which identifies and provides the processor details used
+            // by system. This info contains "flags" for each processor which determines the qualities of that processor
+            // and it's ability to process different instruction sets like mmx, avx, avx2, avx512 and so on.
+            // https://access.redhat.com/documentation/en-us/red_hat_enterprise_linux/6/html/deployment_guide/s2-proc-cpuinfo
+            // Here, we are trying to read the details of all processors used by system and find if any of the processor
+            // 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));
+                    String flags = linestream.filter(line -> line.startsWith("flags")).findFirst().orElse("");
+                    return Arrays.stream(avx512).allMatch(flags::contains);
+                });
+
+            } catch (Exception e) {
+                logger.error("[KNN] Error reading file [{}]. [{}]", fileName, e.getMessage(), e);
+            }
+        }
+        return false;
+    }
 }