Add support for Faiss SQFP16 Quantization and enable AVX2 Optimization (

opensearch-project#1421) * Add Support for Faiss SQFP16 and enable Faiss AVX2 Optimization Signed-off-by: Naveen Tatikonda <[email protected]> * Add Patch Script to fix build on Linux CI Signed-off-by: Naveen Tatikonda <[email protected]> * Disable AVX2 support on Windows Signed-off-by: Naveen Tatikonda <[email protected]> * Add CHANGELOG Signed-off-by: Naveen Tatikonda <[email protected]> * Update Faiss Submodule Signed-off-by: Naveen Tatikonda <[email protected]> * Address Review Comments Signed-off-by: Naveen Tatikonda <[email protected]> * Update UX Interface Signed-off-by: Naveen Tatikonda <[email protected]> * Add Parameter to enable SIMD Signed-off-by: Naveen Tatikonda <[email protected]> * Update DEVELOPER_GUIDE Signed-off-by: Naveen Tatikonda <[email protected]> * Address Review Comments Signed-off-by: Naveen Tatikonda <[email protected]> --------- Signed-off-by: Naveen Tatikonda <[email protected]>
ryanbogan · Feb 2, 2024 · 5f9511b · 5f9511b
1 parent 4c7a055
commit 5f9511b
Show file tree

Hide file tree

Showing 17 changed files with 1,058 additions and 9 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -45,6 +45,10 @@ jobs:
           cd jni/external/faiss
           git apply --ignore-space-change --ignore-whitespace --3way ../../patches/faiss/0001-Custom-patch-to-support-multi-vector.patch
           rm ../../patches/faiss/0001-Custom-patch-to-support-multi-vector.patch
+          git apply --ignore-space-change --ignore-whitespace --3way ../../patches/faiss/0002-Custom-patch-to-support-sqfp16-neon.patch
+          rm ../../patches/faiss/0002-Custom-patch-to-support-sqfp16-neon.patch
+          git apply --ignore-space-change --ignore-whitespace --3way ../../patches/faiss/0003-Custom-patch-to-support-AVX2-Linux-CI.patch
+          rm ../../patches/faiss/0003-Custom-patch-to-support-AVX2-Linux-CI.patch
         working-directory: ${{ github.workspace }}
 
       - name: Setup Java ${{ matrix.java }}
@@ -56,7 +60,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 build"
+          if lscpu  | grep -i avx2
+          then
+            echo "avx2 available on system"
+            su `id -un 1000` -c "whoami && java -version && ./gradlew build -Dsimd.enabled=true"
+          else
+            echo "avx2 not available on system"
+            su `id -un 1000` -c "whoami && java -version && ./gradlew build"
+          fi  
+          
 
       - name: Upload Coverage Report
         uses: codecov/codecov-action@v1
@@ -88,7 +100,14 @@ jobs:
 
       - name: Run build
         run: |
-          ./gradlew build
+          if sysctl -n machdep.cpu.features machdep.cpu.leaf7_features | grep -i AVX2
+          then
+              echo "avx2 available on system"
+              ./gradlew build -Dsimd.enabled=true
+          else
+              echo "avx2 not available on system"
+              ./gradlew build
+          fi
 
   Build-k-NN-Windows:
     strategy:

diff --git a/.github/workflows/test_security.yml b/.github/workflows/test_security.yml
@@ -45,6 +45,10 @@ jobs:
           cd jni/external/faiss
           git apply --ignore-space-change --ignore-whitespace --3way ../../patches/faiss/0001-Custom-patch-to-support-multi-vector.patch
           rm ../../patches/faiss/0001-Custom-patch-to-support-multi-vector.patch
+          git apply --ignore-space-change --ignore-whitespace --3way ../../patches/faiss/0002-Custom-patch-to-support-sqfp16-neon.patch
+          rm ../../patches/faiss/0002-Custom-patch-to-support-sqfp16-neon.patch
+          git apply --ignore-space-change --ignore-whitespace --3way ../../patches/faiss/0003-Custom-patch-to-support-AVX2-Linux-CI.patch
+          rm ../../patches/faiss/0003-Custom-patch-to-support-AVX2-Linux-CI.patch
         working-directory: ${{ github.workspace }}
 
       - name: Setup Java ${{ matrix.java }}
@@ -56,4 +60,4 @@ 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"
+          su `id -un 1000` -c "whoami && java -version && ./gradlew integTest -Dsecurity.enabled=true -Dsimd.enabled=true"
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Features
 * Add parent join support for lucene knn [#1182](https://github.com/opensearch-project/k-NN/pull/1182)
 * Add parent join support for faiss hnsw [#1398](https://github.com/opensearch-project/k-NN/pull/1398)
+* Add Support for Faiss SQFP16 and enable Faiss AVX2 Optimization [#1421](https://github.com/opensearch-project/k-NN/pull/1421)
 ### Enhancements
 * Increase Lucene max dimension limit to 16,000 [#1346](https://github.com/opensearch-project/k-NN/pull/1346)
 * Tuned default values for ef_search and ef_construction for better indexing and search performance for vector search [#1353](https://github.com/opensearch-project/k-NN/pull/1353)

diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md
@@ -11,6 +11,7 @@
   - [Build](#build)
     - [JNI Library](#jni-library)
     - [JNI Library Artifacts](#jni-library-artifacts)
+    - [Enable SIMD Optimization](#enable-simd-optimization)
   - [Run OpenSearch k-NN](#run-opensearch-k-nn)
     - [Run Single-node Cluster Locally](#run-single-node-cluster-locally)
     - [Run Multi-node Cluster Locally](#run-multi-node-cluster-locally)
@@ -236,6 +237,23 @@ If you want to make a custom patch on JNI library
 3. Place the patch file under `jni/patches`
 4. Make a change in `jni/CmakeLists.txt`, `.github/workflows/CI.yml` to apply the patch during build
 
+### Enable SIMD Optimization
+SIMD(Single Instruction/Multiple Data) Optimization can be enabled by setting this optional parameter `simd.enabled` to `true` 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). So, by default it is set to `false`.
+
+```
+# While building OpenSearch k-NN
+./gradlew build -Dsimd.enabled=true
+
+# While running OpenSearch k-NN
+./gradlew run -Dsimd.enabled=true
+
+# While building the JNI libraries
+cd jni
+cmake . -DSIMD_ENABLED=true
+```
+
 ## Run OpenSearch k-NN
 
 ### Run Single-node Cluster Locally

diff --git a/build.gradle b/build.gradle
@@ -17,6 +17,7 @@ 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", "false")
 
         version_tokens = opensearch_version.tokenize('-')
         opensearch_build = version_tokens[0] + '.0'
@@ -297,10 +298,10 @@ task cmakeJniLib(type:Exec) {
     workingDir 'jni'
     if (Os.isFamily(Os.FAMILY_WINDOWS)) {
         dependsOn windowsPatches
-        commandLine 'cmake', '.', "-G", "Unix Makefiles", "-DKNN_PLUGIN_VERSION=${opensearch_version}", "-DBLAS_LIBRARIES=$rootDir\\src\\main\\resources\\windowsDependencies\\libopenblas.dll", "-DLAPACK_LIBRARIES=$rootDir\\src\\main\\resources\\windowsDependencies\\libopenblas.dll"
+        commandLine 'cmake', '.', "-G", "Unix Makefiles", "-DKNN_PLUGIN_VERSION=${opensearch_version}", "-DBLAS_LIBRARIES=$rootDir\\src\\main\\resources\\windowsDependencies\\libopenblas.dll", "-DLAPACK_LIBRARIES=$rootDir\\src\\main\\resources\\windowsDependencies\\libopenblas.dll", "-DSIMD_ENABLED=${simd_enabled}"
     }
     else {
-        commandLine 'cmake', '.', "-DKNN_PLUGIN_VERSION=${opensearch_version}"
+        commandLine 'cmake', '.', "-DKNN_PLUGIN_VERSION=${opensearch_version}", "-DSIMD_ENABLED=${simd_enabled}"
     }
 }
 

diff --git a/jni/CMakeLists.txt b/jni/CMakeLists.txt
@@ -111,7 +111,13 @@ endif ()
 if (${CONFIG_FAISS} STREQUAL ON OR ${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST} STREQUAL ON)
     set(BUILD_TESTING OFF)          # Avoid building faiss tests
     set(BLA_STATIC ON)              # Statically link BLAS
-    set(FAISS_OPT_LEVEL generic)    # Keep optimization level generic
+    if(${CMAKE_SYSTEM_NAME} STREQUAL Windows OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64" 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)
+    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)
+    endif()
 
     if (${CMAKE_SYSTEM_NAME} STREQUAL Darwin)
         if(CMAKE_C_COMPILER_ID MATCHES "Clang\$")
@@ -143,12 +149,20 @@ if (${CONFIG_FAISS} STREQUAL ON OR ${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST} S
     endif ()
 
     # Check if patch exist, this is to skip git apply during CI build. See CI.yml with ubuntu.
-    find_path(PATCH_FILE NAMES 0001-Custom-patch-to-support-multi-vector.patch PATHS ${CMAKE_CURRENT_SOURCE_DIR}/patches/faiss NO_DEFAULT_PATH)
+    find_path(PATCH_FILE NAMES 0001-Custom-patch-to-support-multi-vector.patch 0002-Custom-patch-to-support-sqfp16-neon.patch PATHS ${CMAKE_CURRENT_SOURCE_DIR}/patches/faiss NO_DEFAULT_PATH)
 
     # If it exists, apply patches
     if (EXISTS ${PATCH_FILE})
         message(STATUS "Applying custom patches.")
         execute_process(COMMAND git apply --ignore-space-change --ignore-whitespace --3way ${CMAKE_CURRENT_SOURCE_DIR}/patches/faiss/0001-Custom-patch-to-support-multi-vector.patch WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss ERROR_VARIABLE ERROR_MSG RESULT_VARIABLE RESULT_CODE)
+
+        # 0002-Custom-patch-to-support-sqfp16-neon.patch is a temporary patch to add NEON support to SQ.
+        # Once the commit conflict issues wrt to Multi vector are resolved, this patch can be removed by updating the faiss submodule with corresponding commit.
+        # Apply the patch if the OS is not Windows and Processor is aarch64.
+        if (NOT ${CMAKE_SYSTEM_NAME} STREQUAL Windows AND ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64" AND ${SIMD_ENABLED})
+            execute_process(COMMAND git apply --ignore-space-change --ignore-whitespace --3way ${CMAKE_CURRENT_SOURCE_DIR}/patches/faiss/0002-Custom-patch-to-support-sqfp16-neon.patch WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss ERROR_VARIABLE ERROR_MSG RESULT_VARIABLE RESULT_CODE)
+        endif()
+
         if(RESULT_CODE)
             message(FATAL_ERROR "Failed to apply patch:\n${ERROR_MSG}")
         endif()
@@ -165,7 +179,7 @@ if (${CONFIG_FAISS} STREQUAL ON OR ${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST} S
         ${CMAKE_CURRENT_SOURCE_DIR}/src/knn_extension/faiss/utils/BitSet.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/src/knn_extension/faiss/MultiVectorResultCollector.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/src/knn_extension/faiss/MultiVectorResultCollectorFactory.cpp)
-    target_link_libraries(${TARGET_LIB_FAISS} faiss ${TARGET_LIB_COMMON} OpenMP::OpenMP_CXX)
+    target_link_libraries(${TARGET_LIB_FAISS} ${TARGET_LINK_FAISS_LIB} ${TARGET_LIB_COMMON} OpenMP::OpenMP_CXX)
     target_include_directories(${TARGET_LIB_FAISS} PRIVATE
         ${CMAKE_CURRENT_SOURCE_DIR}/include
         ${CMAKE_CURRENT_SOURCE_DIR}/include/knn_extension/faiss

diff --git a/jni/external/faiss b/jni/external/faiss