diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index ae9ec2e56..f9dfffdbf 100644 --- a/.github/workflows/CI.yml +++ b/.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' diff --git a/.github/workflows/backwards_compatibility_tests_workflow.yml b/.github/workflows/backwards_compatibility_tests_workflow.yml index c20201dab..69af91c6d 100644 --- a/.github/workflows/backwards_compatibility_tests_workflow.yml +++ b/.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 diff --git a/.github/workflows/test_security.yml b/.github/workflows/test_security.yml index 2f8df8526..f3deb096e 100644 --- a/.github/workflows/test_security.yml +++ b/.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 diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a8d808f4..15f5a7139 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index c2545c886..c652209f1 100644 --- a/DEVELOPER_GUIDE.md +++ b/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 diff --git a/build.gradle b/build.gradle index 983573990..f38d0c165 100644 --- a/build.gradle +++ b/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)) { diff --git a/jni/cmake/init-faiss.cmake b/jni/cmake/init-faiss.cmake index 08f3ccc40..4492d9f45 100644 --- a/jni/cmake/init-faiss.cmake +++ b/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) diff --git a/scripts/build.sh b/scripts/build.sh index 12798633b..be7304ee5 100644 --- a/scripts/build.sh +++ b/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 diff --git a/src/main/java/org/opensearch/knn/common/KNNConstants.java b/src/main/java/org/opensearch/knn/common/KNNConstants.java index 7ba900671..16084499c 100644 --- a/src/main/java/org/opensearch/knn/common/KNNConstants.java +++ b/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; diff --git a/src/main/java/org/opensearch/knn/index/KNNSettings.java b/src/main/java/org/opensearch/knn/index/KNNSettings.java index 4da11a2ad..5fcc51bb5 100644 --- a/src/main/java/org/opensearch/knn/index/KNNSettings.java +++ b/src/main/java/org/opensearch/knn/index/KNNSettings.java @@ -14,6 +14,7 @@ import org.opensearch.client.Client; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.Booleans; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; @@ -86,11 +87,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 +305,12 @@ public class KNNSettings { Dynamic ); + public static final Setting KNN_FAISS_AVX512_DISABLED_SETTING = Setting.boolSetting( + KNN_FAISS_AVX512_DISABLED, + KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE, + NodeScope + ); + /** * Dynamic settings */ @@ -429,6 +438,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 +473,7 @@ public List> 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 +513,13 @@ public static boolean isFaissAVX2Disabled() { } } + public static boolean isFaissAVX512Disabled() { + return Booleans.parseBoolean( + KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX512_DISABLED).toString(), + KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE + ); + } + 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 037171b98..4bceed015 100644 --- a/src/main/java/org/opensearch/knn/jni/FaissService.java +++ b/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 @@ -35,9 +37,11 @@ class FaissService { static { AccessController.doPrivileged((PrivilegedAction) () -> { - // 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); diff --git a/src/main/java/org/opensearch/knn/jni/PlatformUtils.java b/src/main/java/org/opensearch/knn/jni/PlatformUtils.java index 8a5549dec..445862f24 100644 --- a/src/main/java/org/opensearch/knn/jni/PlatformUtils.java +++ b/src/main/java/org/opensearch/knn/jni/PlatformUtils.java @@ -20,8 +20,11 @@ import java.nio.file.Files; import java.nio.file.Paths; import java.security.AccessController; +import java.security.PrivilegedActionException; import java.security.PrivilegedExceptionAction; +import java.util.Arrays; import java.util.Locale; +import java.util.stream.Stream; public class PlatformUtils { @@ -37,7 +40,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 +61,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 +82,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) () -> { + Stream linestream = Files.lines(Paths.get(fileName)); + String flags = linestream.filter(line -> line.startsWith("flags")).findFirst().orElse(""); + return Arrays.stream(avx512).allMatch(flags::contains); + }); + + } catch (PrivilegedActionException e) { + logger.error("[KNN] Error reading file [{}]. [{}]", fileName, e.getMessage(), e); + } + } + return false; + } } diff --git a/src/main/plugin-metadata/plugin-security.policy b/src/main/plugin-metadata/plugin-security.policy index d5ab0be21..ed329740f 100644 --- a/src/main/plugin-metadata/plugin-security.policy +++ b/src/main/plugin-metadata/plugin-security.policy @@ -3,6 +3,7 @@ grant { permission java.lang.RuntimePermission "loadLibrary.opensearchknn_faiss"; 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.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 7816505de..19c0abb07 100644 --- a/src/test/java/org/opensearch/knn/jni/PlatformUtilTests.java +++ b/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 mockedPlatform = mockStatic(Platform.class)) { + mockedPlatform.when(Platform::isIntel).thenReturn(false); + assertFalse(isAVX512SupportedBySystem()); + } + } + + public void testIsAVX512SupportedBySystem_platformIsMac_returnsFalse() { + try (MockedStatic mockedPlatform = mockStatic(Platform.class)) { + mockedPlatform.when(Platform::isMac).thenReturn(false); + assertFalse(isAVX512SupportedBySystem()); + } + } + + public void testIsAVX512SupportedBySystem_platformIsIntelMac_returnsFalse() { + try (MockedStatic 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 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 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: AVX2 avx512f avx512cd avx512vl avx512dq avx512bw", "dummy string")); + assertTrue(isAVX512SupportedBySystem()); + } + } + } + + public void testIsAVX512SupportedBySystem_platformIsLinuxSomeAVX512FlagsPresent_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: AVX2 avx512vl avx512dq avx512bw avx512vbmi umip pku ospke avx512_vbmi2", "dummy string")); + assertFalse(isAVX512SupportedBySystem()); + } + } + } }