Skip to content

Commit 550e8a1

Browse files
Add changes for AVX-512 support in k-NN. (#2110) (#2130)
* changes for AVX-512. Signed-off by: Akash Shankaran <akash.shankaran@intel.com> * add cpu detection logic to security workflow. Signed-off by: Akash Shankaran <akash.shankaran@intel.com> * add cpu detection logic to backward compat test workflow. Signed-off by: Akash Shankaran <akash.shankaran@intel.com> * fix bwc workflow. Signed-off by: Akash Shankaran <akash.shankaran@intel.com> * address PR feedback. Signed-off by: Akash Shankaran <akash.shankaran@intel.com> * fix a bug in KNNSettings. Signed-off by: Akash Shankaran <akash.shankaran@intel.com> * fix a bug in KNNSettings. Signed-off by: Akash Shankaran <akash.shankaran@intel.com> * update KNNSettings. Signed-off by: Akash Shankaran <akash.shankaran@intel.com> --------- (cherry picked from commit 5423cc1) Signed-off-by: Akash Shankaran <akash.shankaran@intel.com> Signed-off-by: Ryan Bogan <rbogan@amazon.com> Co-authored-by: akashsha1 <113050768+akashsha1@users.noreply.github.com>
1 parent 9601677 commit 550e8a1

14 files changed

Lines changed: 202 additions & 28 deletions

File tree

.github/workflows/CI.yml

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -73,13 +73,17 @@ jobs:
7373
# switching the user, as OpenSearch cluster can only be started as root/Administrator on linux-deb/linux-rpm/windows-zip.
7474
run: |
7575
chown -R 1000:1000 `pwd`
76-
if lscpu | grep -i avx2
76+
if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw
7777
then
78-
echo "avx2 available on system"
78+
echo "avx512 available on system"
7979
su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc`"
80+
elif lscpu | grep -i avx2
81+
then
82+
echo "avx2 available on system"
83+
su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc` -Davx512.enabled=false"
8084
else
81-
echo "avx2 not available on system"
82-
su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dsimd.enabled=false -Dnproc.count=`nproc`"
85+
echo "avx512 and avx2 not available on system"
86+
su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Dnproc.count=`nproc`"
8387
fi
8488
8589
@@ -126,7 +130,7 @@ jobs:
126130
./gradlew build -Dnproc.count=3
127131
else
128132
echo "avx2 not available on system"
129-
./gradlew build -Dsimd.enabled=false -Dnproc.count=3
133+
./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Dnproc.count=3
130134
fi
131135
132136
Build-k-NN-Windows:
@@ -187,4 +191,4 @@ jobs:
187191
# TODO: Detect processor count and set the value of nproc.count
188192
- name: Run build
189193
run: |
190-
./gradlew.bat build -D'simd.enabled=false' -D'nproc.count=4'
194+
./gradlew.bat build -D'avx2.enabled=false' -D'avx512.enabled=false' -D'nproc.count=4'

.github/workflows/backwards_compatibility_tests_workflow.yml

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -64,9 +64,19 @@ jobs:
6464
6565
- name: Run k-NN Restart-Upgrade BWC Tests from BWCVersion-${{ matrix.bwc_version }} to OpenSearch Version-${{ matrix.opensearch_version }}
6666
run: |
67-
echo "Running restart-upgrade backwards compatibility tests ..."
68-
./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE
69-
67+
echo "Running restart-upgrade backwards compatibility tests ..."
68+
if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw
69+
then
70+
echo "avx512 available on system"
71+
./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Dnproc.count=`nproc`
72+
elif lscpu | grep -i avx2
73+
then
74+
echo "avx2 available on system"
75+
./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Dnproc.count=`nproc` -Davx512.enabled=false
76+
else
77+
echo "avx512 and avx2 not available on system"
78+
./gradlew :qa:restart-upgrade:testRestartUpgrade -Dtests.bwc.version=$BWC_VERSION_RESTART_UPGRADE -Davx2.enabled=false -Davx512.enabled=false -Dsimd.enabled=false -Dnproc.count=`nproc`
79+
fi
7080
7181
Rolling-Upgrade-BWCTests-k-NN:
7282
strategy:
@@ -102,4 +112,15 @@ jobs:
102112
- name: Run k-NN Rolling-Upgrade BWC Tests from BWCVersion-${{ matrix.bwc_version }} to OpenSearch Version-${{ matrix.opensearch_version }}
103113
run: |
104114
echo "Running rolling-upgrade backwards compatibility tests ..."
105-
./gradlew :qa:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE
115+
if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw
116+
then
117+
echo "avx512 available on system"
118+
./gradlew :qa:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE -Dnproc.count=`nproc`
119+
elif lscpu | grep -i avx2
120+
then
121+
echo "avx2 available on system"
122+
./gradlew :qa:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE -Dnproc.count=`nproc` -Davx512.enabled=false
123+
else
124+
echo "avx512 and avx2 not available on system"
125+
./gradlew :qa:rolling-upgrade:testRollingUpgrade -Dtests.bwc.version=$BWC_VERSION_ROLLING_UPGRADE -Davx2.enabled=false -Davx512.enabled=false -Dsimd.enabled=false -Dnproc.count=`nproc`
126+
fi

.github/workflows/test_security.yml

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,4 +70,15 @@ jobs:
7070
# switching the user, as OpenSearch cluster can only be started as root/Administrator on linux-deb/linux-rpm/windows-zip.
7171
run: |
7272
chown -R 1000:1000 `pwd`
73-
su `id -un 1000` -c "whoami && java -version && ./gradlew integTest -Dsecurity.enabled=true -Dsimd.enabled=true -Dnproc.count=`nproc`"
73+
if lscpu | grep -i avx512f | grep -i avx512cd | grep -i avx512vl | grep -i avx512dq | grep -i avx512bw
74+
then
75+
echo "avx512 available on system"
76+
su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc`"
77+
elif lscpu | grep -i avx2
78+
then
79+
echo "avx2 available on system"
80+
su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dnproc.count=`nproc` -Davx512.enabled=false"
81+
else
82+
echo "avx512 and avx2 not available on system"
83+
su `id -un 1000` -c "whoami && java -version && ./gradlew build -Davx2.enabled=false -Davx512.enabled=false -Dnproc.count=`nproc`"
84+
fi

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
1717

1818
## [Unreleased 2.x](https://github.com/opensearch-project/k-NN/compare/2.17...2.x)
1919
### Features
20+
* Add AVX512 support to k-NN for FAISS library [#2069](https://github.com/opensearch-project/k-NN/pull/2069)
2021
### Enhancements
2122
* Add short circuit if no live docs are in segments [#2059](https://github.com/opensearch-project/k-NN/pull/2059)
2223
### Bug Fixes

DEVELOPER_GUIDE.md

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -278,19 +278,22 @@ make -j 4
278278

279279
### Enable SIMD Optimization
280280
SIMD(Single Instruction/Multiple Data) Optimization is enabled by default on Linux and Mac which boosts the performance
281-
by enabling `AVX2` on `x86 architecture` and `NEON` on `ARM64 architecture` while building the Faiss library. But to enable SIMD, the underlying processor
282-
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.
281+
by enabling `AVX2` and `AVX512` on `x86 architecture` and `NEON` on `ARM64 architecture` where applicable while building the Faiss library. But to enable SIMD,
282+
the underlying processor should support these capabilities (AVX512, AVX2 or NEON). It can be disabled by setting the parameter `avx2.enabled` to `false` and
283+
`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
284+
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
285+
feature.
283286

284287
```
285288
# While building OpenSearch k-NN
286-
./gradlew build -Dsimd.enabled=true
289+
./gradlew build -Davx2.enabled=true -Davx512.enabled=true
287290
288291
# While running OpenSearch k-NN
289-
./gradlew run -Dsimd.enabled=true
292+
./gradlew run -Davx2.enabled=true -Davx512.enabled=true
290293
291294
# While building the JNI libraries
292295
cd jni
293-
cmake . -DSIMD_ENABLED=true
296+
cmake . -DAVX2_ENABLED=true -DAVX512_ENABLED=true
294297
```
295298

296299
## Run OpenSearch k-NN

build.gradle

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,9 @@ buildscript {
1717
version_qualifier = System.getProperty("build.version_qualifier", "")
1818
opensearch_group = "org.opensearch"
1919
isSnapshot = "true" == System.getProperty("build.snapshot", "true")
20-
simd_enabled = System.getProperty("simd.enabled", "true")
20+
avx2_enabled = System.getProperty("avx2.enabled", "true")
2121
nproc_count = System.getProperty("nproc.count", "1")
22+
avx512_enabled = System.getProperty("avx512.enabled", "true")
2223
// This flag determines whether the CMake build system should apply a custom patch. It prevents build failures
2324
// when the cmakeJniLib task is run multiple times. If the build.lib.commit_patches is true, the CMake build
2425
// 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) {
316317
args.add("cmake")
317318
args.add(".")
318319
args.add("-DKNN_PLUGIN_VERSION=${opensearch_version}")
319-
args.add("-DSIMD_ENABLED=${simd_enabled}")
320+
args.add("-DAVX2_ENABLED=${avx2_enabled}")
321+
args.add("-DAVX512_ENABLED=${avx512_enabled}")
320322
args.add("-DCOMMIT_LIB_PATCHES=${commit_lib_patches}")
321323
args.add("-DAPPLY_LIB_PATCHES=${apply_lib_patches}")
322324
if (Os.isFamily(Os.FAMILY_WINDOWS)) {

jni/cmake/init-faiss.cmake

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,13 +81,21 @@ set(BUILD_TESTING OFF) # Avoid building faiss tests
8181
set(FAISS_ENABLE_GPU OFF)
8282
set(FAISS_ENABLE_PYTHON OFF)
8383

84-
if(NOT DEFINED SIMD_ENABLED)
85-
set(SIMD_ENABLED true) # set default value as true if the argument is not set
84+
if(NOT DEFINED AVX2_ENABLED)
85+
set(AVX2_ENABLED true) # set default value as true if the argument is not set
8686
endif()
8787

88-
if(${CMAKE_SYSTEM_NAME} STREQUAL Windows OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64" OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm64" OR NOT ${SIMD_ENABLED})
88+
if(NOT DEFINED AVX512_ENABLED)
89+
set(AVX512_ENABLED true) # set default value as true if the argument is not set
90+
endif()
91+
92+
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))
8993
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.
9094
set(TARGET_LINK_FAISS_LIB faiss)
95+
elseif(${CMAKE_SYSTEM_NAME} STREQUAL Linux AND AVX512_ENABLED)
96+
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.
97+
set(TARGET_LINK_FAISS_LIB faiss_avx512)
98+
string(PREPEND LIB_EXT "_avx512") # Prepend "_avx512" to lib extension to create the library as "libopensearchknn_faiss_avx512.so" on linux
9199
else()
92100
set(FAISS_OPT_LEVEL avx2) # Keep optimization level as avx2 to improve performance on Linux and Mac.
93101
set(TARGET_LINK_FAISS_LIB faiss_avx2)

scripts/build.sh

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -118,13 +118,16 @@ fi
118118
# Build k-NN lib and plugin through gradle tasks
119119
cd $work_dir
120120
./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
121-
./gradlew :buildJniLib -Dsimd.enabled=false -Dbuild.lib.commit_patches=false
121+
./gradlew :buildJniLib -Davx512.enabled=false -Davx2.enabled=false -Dbuild.lib.commit_patches=false
122122

123123
if [ "$PLATFORM" != "windows" ] && [ "$ARCHITECTURE" = "x64" ]; then
124124
echo "Building k-NN library after enabling AVX2"
125125
# Skip applying patches as patches were applied already from previous :buildJniLib task
126126
# If we apply patches again, it fails with conflict
127-
./gradlew :buildJniLib -Dsimd.enabled=true -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false
127+
./gradlew :buildJniLib -Davx2.enabled=true -Davx512.enabled=false -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false
128+
129+
echo "Building k-NN library after enabling AVX512"
130+
./gradlew :buildJniLib -Davx512.enabled=true -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false
128131
fi
129132

130133
./gradlew publishPluginZipPublicationToZipStagingRepository -Dopensearch.version=$VERSION -Dbuild.snapshot=$SNAPSHOT -Dbuild.version_qualifier=$QUALIFIER

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,7 @@ public class KNNConstants {
140140
private static final String JNI_LIBRARY_PREFIX = "opensearchknn_";
141141
public static final String FAISS_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + FAISS_NAME;
142142
public static final String FAISS_AVX2_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + FAISS_NAME + "_avx2";
143+
public static final String FAISS_AVX512_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + FAISS_NAME + "_avx512";
143144
public static final String NMSLIB_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + NMSLIB_NAME;
144145
public static final String COMMON_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + COMMONS_NAME;
145146

src/main/java/org/opensearch/knn/index/KNNSettings.java

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.opensearch.client.Client;
1515
import org.opensearch.cluster.metadata.IndexMetadata;
1616
import org.opensearch.cluster.service.ClusterService;
17+
import org.opensearch.common.Booleans;
1718
import org.opensearch.common.settings.Setting;
1819
import org.opensearch.common.settings.Settings;
1920
import org.opensearch.common.unit.TimeValue;
@@ -86,11 +87,13 @@ public class KNNSettings {
8687
public static final String KNN_FAISS_AVX2_DISABLED = "knn.faiss.avx2.disabled";
8788
public static final String QUANTIZATION_STATE_CACHE_SIZE_LIMIT = "knn.quantization.cache.size.limit";
8889
public static final String QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES = "knn.quantization.cache.expiry.minutes";
90+
public static final String KNN_FAISS_AVX512_DISABLED = "knn.faiss.avx512.disabled";
8991

9092
/**
9193
* Default setting values
9294
*/
9395
public static final boolean KNN_DEFAULT_FAISS_AVX2_DISABLED_VALUE = false;
96+
public static final boolean KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE = false;
9497
public static final String INDEX_KNN_DEFAULT_SPACE_TYPE = "l2";
9598
public static final String INDEX_KNN_DEFAULT_SPACE_TYPE_FOR_BINARY = "hamming";
9699
public static final Integer INDEX_KNN_DEFAULT_ALGO_PARAM_M = 16;
@@ -302,6 +305,12 @@ public class KNNSettings {
302305
Dynamic
303306
);
304307

308+
public static final Setting<Boolean> KNN_FAISS_AVX512_DISABLED_SETTING = Setting.boolSetting(
309+
KNN_FAISS_AVX512_DISABLED,
310+
KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE,
311+
NodeScope
312+
);
313+
305314
/**
306315
* Dynamic settings
307316
*/
@@ -429,6 +438,10 @@ private Setting<?> getSetting(String key) {
429438
return KNN_FAISS_AVX2_DISABLED_SETTING;
430439
}
431440

441+
if (KNN_FAISS_AVX512_DISABLED.equals(key)) {
442+
return KNN_FAISS_AVX512_DISABLED_SETTING;
443+
}
444+
432445
if (KNN_VECTOR_STREAMING_MEMORY_LIMIT_IN_MB.equals(key)) {
433446
return KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING;
434447
}
@@ -460,6 +473,7 @@ public List<Setting<?>> getSettings() {
460473
ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD_SETTING,
461474
KNN_FAISS_AVX2_DISABLED_SETTING,
462475
KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING,
476+
KNN_FAISS_AVX512_DISABLED_SETTING,
463477
QUANTIZATION_STATE_CACHE_SIZE_LIMIT_SETTING,
464478
QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES_SETTING
465479
);
@@ -499,6 +513,13 @@ public static boolean isFaissAVX2Disabled() {
499513
}
500514
}
501515

516+
public static boolean isFaissAVX512Disabled() {
517+
return Booleans.parseBoolean(
518+
KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX512_DISABLED).toString(),
519+
KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE
520+
);
521+
}
522+
502523
public static Integer getFilteredExactSearchThreshold(final String indexName) {
503524
return KNNSettings.state().clusterService.state()
504525
.getMetadata()

0 commit comments

Comments
 (0)