Add AVX512 support to k-nn FAISS

build.gradle

@@ -19,6 +19,7 @@ buildscript {
         isSnapshot = "true" == System.getProperty("build.snapshot", "true")
         simd_enabled = System.getProperty("simd.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
@@ -317,6 +318,7 @@ task cmakeJniLib(type:Exec) {
     args.add(".")
     args.add("-DKNN_PLUGIN_VERSION=${opensearch_version}")
     args.add("-DSIMD_ENABLED=${simd_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

@@ -85,9 +85,17 @@ if(NOT DEFINED SIMD_ENABLED)
     set(SIMD_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 SIMD_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 SIMD_ENABLED 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.
+    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)

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;
         }
@@ -461,7 +473,8 @@ public List<Setting<?>> getSettings() {
             KNN_FAISS_AVX2_DISABLED_SETTING,
             KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING,
             QUANTIZATION_STATE_CACHE_SIZE_LIMIT_SETTING,
-            QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES_SETTING
+            QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES_SETTING,
+            KNN_FAISS_AVX512_DISABLED_SETTING
         );
         return Stream.concat(settings.stream(), Stream.concat(getFeatureFlags().stream(), dynamicCacheSettings.values().stream()))
             .collect(Collectors.toList());
@@ -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

@@ -21,6 +21,7 @@
 
 import static org.opensearch.knn.index.KNNSettings.isFaissAVX2Disabled;
 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
@@ -35,10 +36,16 @@ class FaissService {
     static {
         AccessController.doPrivileged((PrivilegedAction<Void>) () -> {
 
-            // Even if the underlying system supports AVX2, users can override and disable it by using the
+            // Even if the underlying system supports AVX512 and 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()) {
-                System.loadLibrary(KNNConstants.FAISS_AVX2_JNI_LIBRARY_NAME);
+            if (!isFaissAVX2Disabled()) {
+                if (isAVX512SupportedBySystem()) {
+                    System.loadLibrary(KNNConstants.FAISS_AVX512_JNI_LIBRARY_NAME);
+                } else if (isAVX2SupportedBySystem()) {
+                    System.loadLibrary(KNNConstants.FAISS_AVX2_JNI_LIBRARY_NAME);
+                } else {
+                    System.loadLibrary(KNNConstants.FAISS_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 {
 
@@ -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()) {
+            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;
+    }
 }

src/test/java/org/opensearch/knn/jni/PlatformUtilTests.java

@@ -22,6 +22,7 @@
 
 import static org.mockito.Mockito.mockStatic;
 import static org.opensearch.knn.jni.PlatformUtils.isAVX2SupportedBySystem;
+import static org.opensearch.knn.jni.PlatformUtils.isAVX512SupportedBySystem;
 
 public class PlatformUtilTests extends KNNTestCase {
     public static final String MAC_CPU_FEATURES = "machdep.cpu.leaf7_features";
@@ -124,4 +125,61 @@ public void testIsAVX2SupportedBySystem_platformIsLinux_throwsExceptionReturnsFa
 
     }
 
+    // AVX512 tests
+
+    public void testIsAVX512SupportedBySystem_platformIsNotIntel_returnsFalse() {
+        try (MockedStatic<Platform> mockedPlatform = mockStatic(Platform.class)) {
+            mockedPlatform.when(Platform::isIntel).thenReturn(false);
+            assertFalse(isAVX512SupportedBySystem());
+        }
+    }
+
+    public void testIsAVX512SupportedBySystem_platformIsMac_returnsFalse() {
+        try (MockedStatic<Platform> mockedPlatform = mockStatic(Platform.class)) {
+            mockedPlatform.when(Platform::isMac).thenReturn(false);
+            assertFalse(isAVX512SupportedBySystem());
+        }
+    }
+
+    public void testIsAVX512SupportedBySystem_platformIsIntelMac_returnsFalse() {
+        try (MockedStatic<Platform> mockedPlatform = mockStatic(Platform.class)) {
+            mockedPlatform.when(Platform::isIntel).thenReturn(true);
+            mockedPlatform.when(Platform::isMac).thenReturn(true);
+            assertFalse(isAVX512SupportedBySystem());
+        }
+    }
+
+    public void testIsAVX512SupportedBySystem_platformIsIntelWithOSAsWindows_returnsFalse() {
+        try (MockedStatic<Platform> mockedPlatform = mockStatic(Platform.class)) {
+            mockedPlatform.when(Platform::isIntel).thenReturn(true);
+            mockedPlatform.when(Platform::isWindows).thenReturn(true);
+            assertFalse(isAVX512SupportedBySystem());
+        }
+    }
+
+    public void testIsAVX512SupportedBySystem_platformIsLinuxAllAVX512FlagsPresent_returnsTrue() {
+        try (MockedStatic<Platform> mockedPlatform = mockStatic(Platform.class)) {
+            mockedPlatform.when(Platform::isIntel).thenReturn(true);
+            mockedPlatform.when(Platform::isLinux).thenReturn(true);
+
+            try (MockedStatic<Files> mockedFiles = mockStatic(Files.class)) {
+                mockedFiles.when(() -> Files.lines(Paths.get(LINUX_PROC_CPU_INFO)))
+                    .thenReturn(Stream.of("flags: AVX2 avx512f avx512cd avx512vl avx512dq avx512bw", "dummy string"));
+                assertTrue(isAVX512SupportedBySystem());
+            }
+        }
+    }
+
+    public void testIsAVX512SupportedBySystem_platformIsLinuxSomeAVX512FlagsPresent_returnsFalse() {
+        try (MockedStatic<Platform> mockedPlatform = mockStatic(Platform.class)) {
+            mockedPlatform.when(Platform::isIntel).thenReturn(true);
+            mockedPlatform.when(Platform::isLinux).thenReturn(true);
+
+            try (MockedStatic<Files> mockedFiles = mockStatic(Files.class)) {
+                mockedFiles.when(() -> Files.lines(Paths.get(LINUX_PROC_CPU_INFO)))
+                    .thenReturn(Stream.of("flags: AVX2 avx512vl avx512dq avx512bw avx512vbmi umip pku ospke avx512_vbmi2", "dummy string"));
+                assertFalse(isAVX512SupportedBySystem());
+            }
+        }
+    }
 }