Address Review Comments

Signed-off-by: Naveen Tatikonda <[email protected]>

.github/workflows/CI.yml

@@ -60,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 -Denable_SIMD=true"
+          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
@@ -92,7 +100,20 @@ jobs:
 
       - name: Run build
         run: |
-          ./gradlew build -Denable_SIMD=true
+          git clone https://github.com/NanXiao/lscpu.git
+          cd lscpu
+          make
+          ./lscpu
+          if ./lscpu  | grep -i avx2
+          then
+              echo "avx2 available on system"
+              cd ..
+              ./gradlew build -Dsimd.enabled=true
+          else
+              echo "avx2 not available on system"
+              cd .. 
+              ./gradlew build
+          fi
 
   Build-k-NN-Windows:
     strategy:

.github/workflows/test_security.yml

@@ -60,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 -Denable_SIMD=true"
+          su `id -un 1000` -c "whoami && java -version && ./gradlew integTest -Dsecurity.enabled=true -Dsimd.enabled=true"

DEVELOPER_GUIDE.md

@@ -238,20 +238,20 @@ If you want to make a custom patch on JNI library
 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 `enable_SIMD` to `true` which boosts the performance
+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 -Denable_SIMD=true
+./gradlew build -Dsimd.enabled=true
 
 # While running OpenSearch k-NN
-./gradlew run -Denable_SIMD=true
+./gradlew run -Dsimd.enabled=true
 
 # While building the JNI libraries
 cd jni
-cmake . -DENABLE_SIMD=true
+cmake . -DSIMD_ENABLED=true
 ```
 
 ## Run OpenSearch k-NN

build.gradle

@@ -17,7 +17,7 @@ buildscript {
         version_qualifier = System.getProperty("build.version_qualifier", "")
         opensearch_group = "org.opensearch"
         isSnapshot = "true" == System.getProperty("build.snapshot", "true")
-        enable_SIMD = System.getProperty("enable_SIMD", "false")
+        simd_enabled = System.getProperty("simd.enabled", "false")
 
         version_tokens = opensearch_version.tokenize('-')
         opensearch_build = version_tokens[0] + '.0'
@@ -298,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", "-DENABLE_SIMD=${enable_SIMD}"
+        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}", "-DENABLE_SIMD=${enable_SIMD}"
+        commandLine 'cmake', '.', "-DKNN_PLUGIN_VERSION=${opensearch_version}", "-DSIMD_ENABLED=${simd_enabled}"
     }
 }
 

jni/CMakeLists.txt

@@ -111,10 +111,12 @@ 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
-    if(WIN32 OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64" OR NOT ${ENABLE_SIMD})
+    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)
@@ -157,7 +159,7 @@ if (${CONFIG_FAISS} STREQUAL ON OR ${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST} S
         # 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 WIN32 AND ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64" AND ${ENABLE_SIMD})
+        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()
 
@@ -177,11 +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)
-    if(WIN32 OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64" OR NOT ${ENABLE_SIMD})
-        target_link_libraries(${TARGET_LIB_FAISS} faiss ${TARGET_LIB_COMMON} OpenMP::OpenMP_CXX)
-    else()
-        target_link_libraries(${TARGET_LIB_FAISS} faiss_avx2 ${TARGET_LIB_COMMON} OpenMP::OpenMP_CXX)
-    endif()
+    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

jni/tests/faiss_wrapper_test.cpp

@@ -357,3 +357,50 @@ TEST(FaissTrainIndexTest, BasicAssertions) {
     // Confirm that training succeeded
     ASSERT_TRUE(trainedIndex->is_trained);
 }
+
+TEST(FaissCreateHnswSQfp16IndexTest, BasicAssertions) {
+    // Define the data
+    faiss::idx_t numIds = 200;
+    std::vector<faiss::idx_t> ids;
+    std::vector<std::vector<float>> vectors;
+    int dim = 2;
+    for (int64_t i = 0; i < numIds; ++i) {
+        ids.push_back(i);
+
+        std::vector<float> vect;
+        vect.reserve(dim);
+        for (int j = 0; j < dim; ++j) {
+            vect.push_back(test_util::RandomFloat(-500.0, 500.0));
+        }
+        vectors.push_back(vect);
+    }
+
+    std::string indexPath = test_util::RandomString(10, "tmp/", ".faiss");
+    std::string spaceType = knn_jni::L2;
+    std::string index_description = "HNSW32,SQfp16";
+
+    std::unordered_map<std::string, jobject> parametersMap;
+    parametersMap[knn_jni::SPACE_TYPE] = (jobject)&spaceType;
+    parametersMap[knn_jni::INDEX_DESCRIPTION] = (jobject)&index_description;
+
+    // Set up jni
+    JNIEnv *jniEnv = nullptr;
+    NiceMock<test_util::MockJNIUtil> mockJNIUtil;
+
+    EXPECT_CALL(mockJNIUtil,
+                GetJavaObjectArrayLength(
+                        jniEnv, reinterpret_cast<jobjectArray>(&vectors)))
+            .WillRepeatedly(Return(vectors.size()));
+
+    // Create the index
+    knn_jni::faiss_wrapper::CreateIndex(
+            &mockJNIUtil, jniEnv, reinterpret_cast<jintArray>(&ids),
+            reinterpret_cast<jobjectArray>(&vectors), (jstring)&indexPath,
+            (jobject)&parametersMap);
+
+    // Make sure index can be loaded
+    std::unique_ptr<faiss::Index> index(test_util::FaissLoadIndex(indexPath));
+
+    // Clean up
+    std::remove(indexPath.c_str());
+}

src/test/java/org/opensearch/knn/index/FaissIT.java

@@ -35,6 +35,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.Random;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
 
@@ -278,7 +279,9 @@ public void testHNSWSQFP16_whenIndexedAndQueried_thenSucceed() {
         String fieldName = "test-field-hnsw-sqfp16";
 
         KNNMethod hnswMethod = KNNEngine.FAISS.getMethod(KNNConstants.METHOD_HNSW);
-        SpaceType spaceType = SpaceType.L2;
+        SpaceType[] spaceTypes = { SpaceType.L2, SpaceType.INNER_PRODUCT };
+        Random random = new Random();
+        SpaceType spaceType = spaceTypes[random.nextInt(spaceTypes.length)];
 
         List<Integer> mValues = ImmutableList.of(16, 32, 64, 128);
         List<Integer> efConstructionValues = ImmutableList.of(16, 32, 64, 128);

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

@@ -13,6 +13,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import lombok.SneakyThrows;
 import org.junit.BeforeClass;
 import org.opensearch.Version;
 import org.opensearch.common.xcontent.XContentFactory;
@@ -469,7 +470,8 @@ public void testCreateIndex_faiss_invalid_invalidIndexDescription() throws IOExc
         );
     }
 
-    public void testCreateIndex_faiss_sqfp16_invalidIndexDescription() throws IOException {
+    @SneakyThrows
+    public void testCreateIndex_faiss_sqfp16_invalidIndexDescription() {
 
         int[] docIds = new int[] { 1, 2 };
         float[][] vectors = new float[][] { { 2, 3 }, { 3, 4 } };
@@ -493,7 +495,8 @@ public void testCreateIndex_faiss_sqfp16_invalidIndexDescription() throws IOExce
         );
     }
 
-    public void testLoadIndex_faiss_sqfp16_valid() throws IOException {
+    @SneakyThrows
+    public void testLoadIndex_faiss_sqfp16_valid() {
 
         int[] docIds = new int[] { 1, 2 };
         float[][] vectors = new float[][] { { 2, 3 }, { 3, 4 } };
@@ -513,7 +516,8 @@ public void testLoadIndex_faiss_sqfp16_valid() throws IOException {
         assertNotEquals(0, pointer);
     }
 
-    public void testQueryIndex_faiss_sqfp16_valid() throws IOException {
+    @SneakyThrows
+    public void testQueryIndex_faiss_sqfp16_valid() {
 
         String sqfp16IndexDescription = "HNSW16,SQfp16";
         int k = 10;
@@ -543,7 +547,8 @@ public void testQueryIndex_faiss_sqfp16_valid() throws IOException {
         }
     }
 
-    public void testTrain_whenConfigurationIsIVFSQFP16_thenSucceed() throws IOException {
+    @SneakyThrows
+    public void testTrain_whenConfigurationIsIVFSQFP16_thenSucceed() {
         long trainPointer = transferVectors(10);
         int ivfNlistParam = 16;
         XContentBuilder xContentBuilder = XContentFactory.jsonBuilder()