Faiss BBQ Bulk SIMD. (#3170)

Signed-off-by: Dooyong Kim <kdooyong@amazon.com>

Author: 0ctopus13prime

jni/include/org_opensearch_knn_jni_SimdVectorComputeService.h

@@ -19,3 +19,20 @@
     #define LIKELY(x)   (x)
     #define UNLIKELY(x) (x)
 #endif
+
+#if defined(_MSC_VER)
+    // Microsoft Visual C++
+    #define FORCE_INLINE __forceinline
+#elif defined(__GNUC__) || defined(__clang__)
+    // GCC and Clang
+    #define FORCE_INLINE inline __attribute__((always_inline))
+#else
+    // Fallback for other compilers
+    #define FORCE_INLINE inline
+#endif // FORCE_INLINE_H
+
+#if defined(__GNUC__) || defined(__clang__)
+    #define HOT_SPOT [[gnu::hot]]
+#else
+    #define HOT_SPOT
+#endif

jni/include/platform_defs.h

@@ -12,7 +12,9 @@ namespace knn_jni::simd::similarity_function {
         // Max inner product will transform inner product to v < 0 ? 1 / (1 - v) : (1 + v)
         FP16_MAXIMUM_INNER_PRODUCT,
         // L2 for FP16
-        FP16_L2
+        FP16_L2,
+        BBQ_IP,
+        BBQ_L2
     };
 
     struct SimilarityFunction;

jni/include/simd/similarity_function/similarity_function.h

@@ -1,3 +1,6 @@
+#include <cstring>
+#include <limits>
+#include <algorithm>
 #include "org_opensearch_knn_jni_SimdVectorComputeService.h"
 #include "jni_util.h"
 #include "simd/similarity_function/similarity_function.h"
@@ -26,8 +29,11 @@ void JNI_OnUnload(JavaVM *vm, void *reserved) {
     JNI_UTIL.Uninitialize(env);
 }
 
-JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_SimdVectorComputeService_scoreSimilarityInBulk
+JNIEXPORT jfloat JNICALL Java_org_opensearch_knn_jni_SimdVectorComputeService_scoreSimilarityInBulk
   (JNIEnv *env, jclass clazz, jintArray internalVectorIds, jfloatArray jscores, const jint numVectors) {
+    if (numVectors <= 0) {
+      return std::numeric_limits<float>::min();
+    }
 
     try {
       // Get search context
@@ -38,7 +44,14 @@ JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_SimdVectorComputeService_scor
 
       // Get pointers of vectorIds and scores
       jint* vectorIds = static_cast<jint*>(JNI_UTIL.GetPrimitiveArrayCritical(env, internalVectorIds, nullptr));
+      knn_jni::JNIReleaseElements releaseVectorIds {[=]{
+        JNI_UTIL.ReleasePrimitiveArrayCritical(env, internalVectorIds, vectorIds, 0);
+      }};
+
       jfloat* scores = static_cast<jfloat*>(JNI_UTIL.GetPrimitiveArrayCritical(env, jscores, nullptr));
+      knn_jni::JNIReleaseElements releaseScores {[=]{
+        JNI_UTIL.ReleasePrimitiveArrayCritical(env, jscores, scores, 0);
+      }};
 
       // Bulk similarity calculation
       srchContext->similarityFunction->calculateSimilarityInBulk(
@@ -47,11 +60,11 @@ JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_SimdVectorComputeService_scor
           reinterpret_cast<float*>(scores),
           numVectors);
 
-      // Release pinned pointers
-      JNI_UTIL.ReleasePrimitiveArrayCritical(env, internalVectorIds, vectorIds, 0);
-      JNI_UTIL.ReleasePrimitiveArrayCritical(env, jscores, scores, 0);
+      jfloat maxScore = *std::max_element(scores, scores + numVectors);
+      return maxScore;
     } catch (...) {
       JNI_UTIL.CatchCppExceptionAndThrowJava(env);
+      return 0.0f;   // value ignored if exception pending
     }
 }
 
@@ -96,3 +109,46 @@ JNIEXPORT jfloat JNICALL Java_org_opensearch_knn_jni_SimdVectorComputeService_sc
 
     return 0;
 }
+
+JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_SimdVectorComputeService_saveBBQSearchContext
+  (JNIEnv *env, jclass clazz, jbyteArray quantizedQuery,
+   jfloat lowerInterval, jfloat upperInterval, jfloat additionalCorrection,
+   jint quantizedComponentSum, jlongArray addressAndSize,
+   jint functionTypeOrd, jint dimension, jfloat centroidDp) {
+    try {
+      // Get quantized query bytes
+      const jsize queryByteSize = JNI_UTIL.GetJavaBytesArrayLength(env, quantizedQuery);
+      jbyte* queryPtr = static_cast<jbyte*>(JNI_UTIL.GetPrimitiveArrayCritical(env, quantizedQuery, nullptr));
+      knn_jni::JNIReleaseElements queryRelease {[=]{
+        JNI_UTIL.ReleasePrimitiveArrayCritical(env, quantizedQuery, queryPtr, 0);
+      }};
+
+      // Get mmap address and size
+      const jsize mmapAddressAndSizeLength = JNI_UTIL.GetJavaLongArrayLength(env, addressAndSize);
+      jlong* mmapAddressAndSize = static_cast<jlong*>(JNI_UTIL.GetPrimitiveArrayCritical(env, addressAndSize, nullptr));
+      knn_jni::JNIReleaseElements mmapAddressAndSizeRelease {[=]{
+        JNI_UTIL.ReleasePrimitiveArrayCritical(env, addressAndSize, mmapAddressAndSize, 0);
+      }};
+
+      // Store correction factors in tmpBuffer before calling saveSearchContext.
+      // saveSearchContext will reset tmpBuffer at the beginning, so we need to call it first,
+      // then write correction factors after.
+      SimilarityFunction::saveSearchContext(
+          reinterpret_cast<uint8_t*>(queryPtr), queryByteSize,
+          dimension,
+          reinterpret_cast<int64_t*>(mmapAddressAndSize), mmapAddressAndSizeLength,
+          functionTypeOrd);
+
+      // Now store correction factors in tmpBuffer (saveSearchContext clears it, then BBQ_IP branch leaves it empty)
+      SimdVectorSearchContext* ctx = SimilarityFunction::getSearchContext();
+      ctx->tmpBuffer.resize(5 * sizeof(float));
+      auto* correctionPtr = reinterpret_cast<float*>(ctx->tmpBuffer.data());
+      correctionPtr[0] = lowerInterval;
+      correctionPtr[1] = upperInterval;
+      correctionPtr[2] = additionalCorrection;
+      std::memcpy(&correctionPtr[3], &quantizedComponentSum, sizeof(int32_t));
+      correctionPtr[4] = centroidDp;
+    } catch (...) {
+      JNI_UTIL.CatchCppExceptionAndThrowJava(env);
+    }
+}