JMH benchmarks for transfer vectors change

Signed-off-by: Navneet Verma <[email protected]>

jni/include/org_opensearch_knn_jni_FaissService.h

@@ -122,8 +122,20 @@ JNIEXPORT jbyteArray JNICALL Java_org_opensearch_knn_jni_FaissService_trainIndex
 JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectors
   (JNIEnv *, jclass, jlong, jobjectArray);
 
+/*
+ * Class:     org_opensearch_knn_jni_FaissService
+ * Method:    transferVectorsV2
+ * Signature: (J[[F)J
+ */
 JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectorsV2
         (JNIEnv *, jclass, jlong, jobjectArray);
+/*
+ * Class:     org_opensearch_knn_jni_FaissService
+ * Method:    transferVectorsV3
+ * Signature: (J[[F)J
+ */
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectorsV3(JNIEnv *, jclass ,
+jlong ,jobjectArray , jlong , jlong );
 
 /*
  * Class:     org_opensearch_knn_jni_FaissService

jni/src/org_opensearch_knn_jni_FaissService.cpp

@@ -209,6 +209,32 @@ jlong vectorsPointerJ,
     return (jlong) vect;
 }
 
+JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectorsV3(JNIEnv * env, jclass cls,
+jlong vectorsPointerJ,
+        jobjectArray vectorsJ, jlong startingIndex, jlong initialCapacity)
+{
+    std::vector<float> *vect;
+    if ((long) vectorsPointerJ == 0) {
+        vect = new std::vector<float>((long) initialCapacity);
+    } else {
+        vect = reinterpret_cast<std::vector<float>*>(vectorsPointerJ);
+    }
+
+    int dim = jniUtil.GetInnerDimensionOf2dJavaFloatArray(env, vectorsJ);
+    auto dataset = jniUtil.Convert2dJavaObjectArrayToCppFloatVector(env, vectorsJ, dim);
+
+    for(long i = 0; i < dataset.size(); i++) {
+        vect->at(i + startingIndex) = dataset[i];
+    }
+
+//    if((long)startingIndex == 0) {
+//        vect->insert(vect->end(), dataset.begin(), dataset.end());
+//    } else {
+//        vect->insert(vect->begin() + (long)startingIndex, dataset.begin(), dataset.end());
+//    }
+    return (jlong) vect;
+}
+
 JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_freeVectors(JNIEnv * env, jclass cls,
                                                                             jlong vectorsPointerJ)
 {

micro-benchmarks/README.md

@@ -88,4 +88,8 @@ If you want to disassemble a single method do something like this:
 gradlew -p micro-benchmarks run --args ' MemoryStatsBenchmark -jvmArgs "-XX:+UnlockDiagnosticVMOptions -XX:CompileCommand=print,*.yourMethodName -XX:PrintAssemblyOptions=intel"
 ```
 
+```
+./gradlew -p micro-benchmarks run --args ' -gc true -jvmArgs "-Xms2048m -Xmx2048m"'
+```
+
 If you want `perf` to find the hot methods for you then do add `-prof:perfasm`.

micro-benchmarks/build.gradle

@@ -43,6 +43,8 @@ run {
     // This is required for C++ code
     systemProperty "java.library.path", "$rootDir/jni/release"
     executable = "${BuildParams.runtimeJavaHome}/bin/java"
+    var jvmHeapSize = System.getProperty("jvm.heap.size", "6g")
+    jvmArgs("-Xms" + jvmHeapSize, "-Xmx"  + jvmHeapSize)
 }
 
 

micro-benchmarks/src/main/java/org/opensearch/knn/TransferVectorsBenchmarks.java

@@ -30,21 +30,21 @@
 import java.util.Random;
 import java.util.concurrent.TimeUnit;
 
-@Warmup(iterations = 1, timeUnit = TimeUnit.SECONDS, time = 20)
-@Measurement(iterations = 3, timeUnit = TimeUnit.SECONDS, time = 20)
-@Fork(1)
-@BenchmarkMode(Mode.AverageTime)
+@Warmup(iterations = 1, timeUnit = TimeUnit.SECONDS, time = 300)
+@Measurement(iterations = 1, timeUnit = TimeUnit.SECONDS, time = 300)
+@Fork(3)
+@BenchmarkMode(Mode.SingleShotTime)
 @OutputTimeUnit(TimeUnit.SECONDS)
 @State(Scope.Benchmark)
+
 public class TransferVectorsBenchmarks {
     private static final Random random = new Random(1212121212);
     private static final int TOTAL_NUMBER_OF_VECTOR_TO_BE_TRANSFERRED = 1000000;
 
     @Param({ "128", "256", "384", "512" })
     private int dimension;
 
-    //@Param({ "1000000", "500000", "100000" })
-    @Param({ "500000"})
+    @Param({ "100000", "500000", "1000000" })
     private int vectorsPerTransfer;
 
     private List<float[]> vectorList;
@@ -61,6 +61,7 @@ public void setup() {
     public void transferVectors() {
         long vectorsAddress = 0;
         List<float[]> vectorToTransfer = new ArrayList<>();
+        long startTime = System.currentTimeMillis();
         for (float[] floats : vectorList) {
             if (vectorToTransfer.size() == vectorsPerTransfer) {
                 vectorsAddress = JNIService.transferVectorsV2(vectorsAddress, vectorToTransfer.toArray(new float[][]{}));
@@ -73,6 +74,34 @@ public void transferVectors() {
         }
 
         JNIService.freeVectors(vectorsAddress);
+        long endTime = System.currentTimeMillis();
+        System.out.println("Time taken to transfer vectors " + (endTime - startTime));
+    }
+
+    @Benchmark
+    public void transferVectorsInitialCapacity() {
+        long vectorsAddress = 0;
+        long startingIndex = 0;
+        long startTime = System.currentTimeMillis();
+        List<float[]> vectorToTransfer = new ArrayList<>();
+        for (float[] floats : vectorList) {
+            if (vectorToTransfer.size() == vectorsPerTransfer) {
+                vectorsAddress = JNIService.transferVectorsV3(vectorsAddress,
+                        vectorToTransfer.toArray(new float[][]{}), startingIndex,
+                        TOTAL_NUMBER_OF_VECTOR_TO_BE_TRANSFERRED * dimension);
+                vectorToTransfer = new ArrayList<>();
+                startingIndex += vectorsPerTransfer;
+            }
+            vectorToTransfer.add(floats);
+        }
+        if(!vectorToTransfer.isEmpty()) {
+            vectorsAddress = JNIService.transferVectorsV3(vectorsAddress, vectorToTransfer.toArray(new float[][]{}),
+                    startingIndex, TOTAL_NUMBER_OF_VECTOR_TO_BE_TRANSFERRED * dimension);
+            startingIndex += vectorToTransfer.size();
+        }
+        JNIService.freeVectors(vectorsAddress);
+        long endTime = System.currentTimeMillis();
+        System.out.println("Time taken to transferVectorsInitialCapacity " + (endTime - startTime));
     }
 
     private float[] generateRandomVector(int dimensions) {

src/main/java/org/opensearch/knn/jni/FaissService.java

@@ -11,6 +11,7 @@
 
 package org.opensearch.knn.jni;
 
+import lombok.extern.log4j.Log4j2;
 import org.opensearch.knn.common.KNNConstants;
 import org.opensearch.knn.index.query.KNNQueryResult;
 import org.opensearch.knn.index.util.KNNEngine;
@@ -30,23 +31,26 @@
  *      src/main/java/org/opensearch/knn/index/query/KNNQueryResult.java
  *      src/main/java/org/opensearch/knn/common/KNNConstants.java
  */
+@Log4j2
 class FaissService {
 
     static {
         AccessController.doPrivileged((PrivilegedAction<Void>) () -> {
-            boolean isFaissAVX2Disabled = true;
-            try {
-                isFaissAVX2Disabled = isFaissAVX2Disabled();
-            } catch (Exception e) {
-
-            }
-            // 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()) {
-                System.loadLibrary(KNNConstants.FAISS_AVX2_JNI_LIBRARY_NAME);
-            } else {
+//            boolean isFaissAVX2Disabled = true;
+//            try {
+//                isFaissAVX2Disabled = isFaissAVX2Disabled();
+//            } catch (Exception e) {
+//
+//            }
+//            // 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()) {
+//                log.error("SIMD is enabled");
+//                System.loadLibrary(KNNConstants.FAISS_AVX2_JNI_LIBRARY_NAME);
+//            } else {
+                log.error("SIMD is not enabled");
                 System.loadLibrary(KNNConstants.FAISS_JNI_LIBRARY_NAME);
-            }
+            //}
 
             initLibrary();
             KNNEngine.FAISS.setInitialized(true);
@@ -188,6 +192,8 @@ public static native KNNQueryResult[] queryIndexWithFilter(
 
     public static native long transferVectorsV2(long vectorsPointer, float[][] trainingData);
 
+    public static native long transferVectorsV3(long vectorsPointer, float[][] trainingData, long startingIndex, long initialCapacity);
+
     /**
      * Free vectors from memory
      *

src/main/java/org/opensearch/knn/jni/JNIService.java

@@ -257,4 +257,9 @@ public static void freeVectors(long vectorsPointer) {
     public static long transferVectorsV2(long vectorsPointer, float[][] trainingData) {
         return FaissService.transferVectorsV2(vectorsPointer, trainingData);
     }
+
+    public static long transferVectorsV3(long vectorsPointer, float[][] trainingData, long startingIndex,
+                                                long initialCapacity) {
+        return FaissService.transferVectorsV3(vectorsPointer, trainingData, startingIndex, initialCapacity);
+    }
 }

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

@@ -27,23 +27,23 @@ public class TransferVectorsBenchmarksIT extends OpenSearchTestCase {
     private static final int MAX_DATA_SIZE_TO_BE_TRANSFERRED_IN_MB = 100;
     private static final int CONVERSION_CONSTANT = 1024;
 
-    private static final List<Integer> DIMENSIONS_LIST = List.of(1024, 1536,  968, 768, 512, 256, 128);
-    //List.of(128, 256, 512, 768, 968, 1024, 1536);
+    private static final List<Integer> DIMENSIONS_LIST = List.of(1024, 1536, 968, 768, 512, 256, 128);
+    // List.of(128, 256, 512, 768, 968, 1024, 1536);
 
     public void test_transferVectorV2Speed() {
         final List<Double> elapsedTime = new ArrayList<>();
         System.out.println("Detail for Vector transfer based on specific MB: " + MAX_DATA_SIZE_TO_BE_TRANSFERRED_IN_MB);
         System.out.println("Dimension, Elapsed Time");
         for (int i = 0; i < DIMENSIONS_LIST.size(); i++) {
             System.gc();
-            int vectorsPerTransfer = (MAX_DATA_SIZE_TO_BE_TRANSFERRED_IN_MB * CONVERSION_CONSTANT * CONVERSION_CONSTANT) / (DIMENSIONS_LIST.get(i) * Float.BYTES);
+            int vectorsPerTransfer = (MAX_DATA_SIZE_TO_BE_TRANSFERRED_IN_MB * CONVERSION_CONSTANT * CONVERSION_CONSTANT) / (DIMENSIONS_LIST
+                .get(i) * Float.BYTES);
             elapsedTime.add(transferVectorUtil(DIMENSIONS_LIST.get(i), 1000000, vectorsPerTransfer));
             System.out.println(DIMENSIONS_LIST.get(i) + " , " + elapsedTime.get(i));
         }
         Assert.assertTrue(true);
     }
 
-
     public void test_transferVectorOfSpecificCount() {
         final List<Double> elapsedTime = new ArrayList<>();
         int vectorsPerTransfer = 10 * 1000;
@@ -73,7 +73,6 @@ public void test_transferVectorSendAllAtOnce() {
         Assert.assertTrue(true);
     }
 
-
     private static float[] generateRandomVector(int dimensions) {
         float[] vector = new float[dimensions];
         for (int i = 0; i < dimensions; i++) {
@@ -87,12 +86,11 @@ private double transferVectorUtil(int dimension, int totalNumberOfVectorsToBeTra
         long vectorsAddress = 0;
         long timeInNano = 0;
 
-
         for (int i = 0; i < totalNumberOfVectorsToBeTransferred; i++) {
             vectorList.add(generateRandomVector(dimension));
             if (vectorList.size() == vectorsPerTransfer) {
                 long startTime = System.nanoTime();
-                vectorsAddress = FaissService.transferVectorsV2(vectorsAddress, vectorList.toArray(new float[][]{}));
+                vectorsAddress = FaissService.transferVectorsV2(vectorsAddress, vectorList.toArray(new float[][] {}));
                 long endTime = System.nanoTime();
                 timeInNano = timeInNano + (endTime - startTime);
                 vectorList = new ArrayList<>();
@@ -101,7 +99,7 @@ private double transferVectorUtil(int dimension, int totalNumberOfVectorsToBeTra
 
         if (!CollectionUtils.isEmpty(vectorList)) {
             long startTime = System.nanoTime();
-            vectorsAddress = FaissService.transferVectorsV2(vectorsAddress, vectorList.toArray(new float[][]{}));
+            vectorsAddress = FaissService.transferVectorsV2(vectorsAddress, vectorList.toArray(new float[][] {}));
             long endTime = System.nanoTime();
             timeInNano = timeInNano + (endTime - startTime);
             vectorList = new ArrayList<>();