[Native][x64] Native implementation of DiskBBQ scoring (distance w adjustment)

libs/native/libraries/build.gradle

@@ -19,7 +19,7 @@ configurations {
 }
 
 var zstdVersion = "1.5.7"
-var vecVersion = "1.0.34"
+var vecVersion = "1.0.36"
 
 repositories {
   exclusiveContent {

libs/native/src/main/java/org/elasticsearch/nativeaccess/VectorSimilarityFunctions.java

@@ -130,4 +130,10 @@ enum Operation {
     MethodHandle getHandle(Function function, DataType dataType, Operation operation);
 
     MethodHandle getHandle(Function function, BBQType bbqType, Operation operation);
+
+    MethodHandle applyCorrectionsEuclideanBulk();
+
+    MethodHandle applyCorrectionsMaxInnerProductBulk();
+
+    MethodHandle applyCorrectionsDotProductBulk();
 }

libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/JdkVectorLibrary.java

@@ -45,6 +45,10 @@ private record OperationSignature<E extends Enum<E>>(Function function, E dataTy
 
     private static final Map<OperationSignature<?>, MethodHandle> HANDLES;
 
+    static final MethodHandle applyCorrectionsEuclideanBulk$mh;
+    static final MethodHandle applyCorrectionsMaxInnerProductBulk$mh;
+    static final MethodHandle applyCorrectionsDotProductBulk$mh;
+
     private static final JdkVectorSimilarityFunctions INSTANCE;
 
     /**
@@ -156,6 +160,26 @@ private static MethodHandle bindFunction(String functionName, int capability, Fu
                 }
 
                 HANDLES = Collections.unmodifiableMap(handles);
+
+                FunctionDescriptor score = FunctionDescriptor.of(
+                    JAVA_FLOAT,
+                    ADDRESS, // corrections
+                    JAVA_INT, // bulkSize,
+                    JAVA_INT, // dimensions,
+                    JAVA_FLOAT, // queryLowerInterval,
+                    JAVA_FLOAT, // queryUpperInterval,
+                    JAVA_INT, // queryComponentSum,
+                    JAVA_FLOAT, // queryAdditionalCorrection,
+                    JAVA_FLOAT, // queryBitScale,
+                    JAVA_FLOAT, // indexBitScale,
+                    JAVA_FLOAT, // centroidDp,
+                    ADDRESS // scores
+                );
+
+                applyCorrectionsEuclideanBulk$mh = bindFunction("diskbbq_apply_corrections_euclidean_bulk", caps, score);
+                applyCorrectionsMaxInnerProductBulk$mh = bindFunction("diskbbq_apply_corrections_maximum_inner_product_bulk", caps, score);
+                applyCorrectionsDotProductBulk$mh = bindFunction("diskbbq_apply_corrections_dot_product_bulk", caps, score);
+
                 INSTANCE = new JdkVectorSimilarityFunctions();
             } else {
                 if (caps < 0) {
@@ -164,6 +188,9 @@ private static MethodHandle bindFunction(String functionName, int capability, Fu
                         enable them in your OS/Hypervisor/VM/container""");
                 }
                 HANDLES = null;
+                applyCorrectionsEuclideanBulk$mh = null;
+                applyCorrectionsMaxInnerProductBulk$mh = null;
+                applyCorrectionsDotProductBulk$mh = null;
                 INSTANCE = null;
             }
         } catch (Throwable t) {
@@ -379,8 +406,108 @@ private static void checkByteSize(MemorySegment a, MemorySegment b) {
             }
         }
 
+        private static float applyCorrectionsEuclideanBulk(
+            MemorySegment corrections,
+            int bulkSize,
+            int dimensions,
+            float queryLowerInterval,
+            float queryUpperInterval,
+            int queryComponentSum,
+            float queryAdditionalCorrection,
+            float queryBitScale,
+            float indexBitScale,
+            float centroidDp,
+            MemorySegment scores
+        ) {
+            try {
+                return (float) applyCorrectionsEuclideanBulk$mh.invokeExact(
+                    corrections,
+                    bulkSize,
+                    dimensions,
+                    queryLowerInterval,
+                    queryUpperInterval,
+                    queryComponentSum,
+                    queryAdditionalCorrection,
+                    queryBitScale,
+                    indexBitScale,
+                    centroidDp,
+                    scores
+                );
+            } catch (Throwable t) {
+                throw new AssertionError(t);
+            }
+        }
+
+        private static float applyCorrectionsMaxInnerProductBulk(
+            MemorySegment corrections,
+            int bulkSize,
+            int dimensions,
+            float queryLowerInterval,
+            float queryUpperInterval,
+            int queryComponentSum,
+            float queryAdditionalCorrection,
+            float queryBitScale,
+            float indexBitScale,
+            float centroidDp,
+            MemorySegment scores
+        ) {
+            try {
+                return (float) applyCorrectionsMaxInnerProductBulk$mh.invokeExact(
+                    corrections,
+                    bulkSize,
+                    dimensions,
+                    queryLowerInterval,
+                    queryUpperInterval,
+                    queryComponentSum,
+                    queryAdditionalCorrection,
+                    queryBitScale,
+                    indexBitScale,
+                    centroidDp,
+                    scores
+                );
+            } catch (Throwable t) {
+                throw new AssertionError(t);
+            }
+        }
+
+        private static float applyCorrectionsDotProductBulk(
+            MemorySegment corrections,
+            int bulkSize,
+            int dimensions,
+            float queryLowerInterval,
+            float queryUpperInterval,
+            int queryComponentSum,
+            float queryAdditionalCorrection,
+            float queryBitScale,
+            float indexBitScale,
+            float centroidDp,
+            MemorySegment scores
+        ) {
+            try {
+                return (float) applyCorrectionsDotProductBulk$mh.invokeExact(
+                    corrections,
+                    bulkSize,
+                    dimensions,
+                    queryLowerInterval,
+                    queryUpperInterval,
+                    queryComponentSum,
+                    queryAdditionalCorrection,
+                    queryBitScale,
+                    indexBitScale,
+                    centroidDp,
+                    scores
+                );
+            } catch (Throwable t) {
+                throw new AssertionError(t);
+            }
+        }
+
         private static final Map<OperationSignature<?>, MethodHandle> HANDLES_WITH_CHECKS;
 
+        static final MethodHandle APPLY_CORRECTIONS_EUCLIDEAN_HANDLE_BULK;
+        static final MethodHandle APPLY_CORRECTIONS_MAX_INNER_PRODUCT_HANDLE_BULK;
+        static final MethodHandle APPLY_CORRECTIONS_DOT_PRODUCT_HANDLE_BULK;
+
         static {
             MethodHandles.Lookup lookup = MethodHandles.lookup();
 
@@ -535,6 +662,37 @@ private static void checkByteSize(MemorySegment a, MemorySegment b) {
                 }
 
                 HANDLES_WITH_CHECKS = Collections.unmodifiableMap(handlesWithChecks);
+
+                MethodType scoringFunction = MethodType.methodType(
+                    float.class,
+                    MemorySegment.class,
+                    int.class,
+                    int.class,
+                    float.class,
+                    float.class,
+                    int.class,
+                    float.class,
+                    float.class,
+                    float.class,
+                    float.class,
+                    MemorySegment.class
+                );
+
+                APPLY_CORRECTIONS_EUCLIDEAN_HANDLE_BULK = lookup.findStatic(
+                    JdkVectorSimilarityFunctions.class,
+                    "applyCorrectionsEuclideanBulk",
+                    scoringFunction
+                );
+                APPLY_CORRECTIONS_MAX_INNER_PRODUCT_HANDLE_BULK = lookup.findStatic(
+                    JdkVectorSimilarityFunctions.class,
+                    "applyCorrectionsMaxInnerProductBulk",
+                    scoringFunction
+                );
+                APPLY_CORRECTIONS_DOT_PRODUCT_HANDLE_BULK = lookup.findStatic(
+                    JdkVectorSimilarityFunctions.class,
+                    "applyCorrectionsDotProductBulk",
+                    scoringFunction
+                );
             } catch (ReflectiveOperationException e) {
                 throw new AssertionError(e);
             }
@@ -555,5 +713,20 @@ public MethodHandle getHandle(Function function, BBQType bbqType, Operation oper
             if (mh == null) throw new IllegalArgumentException("Signature not implemented: " + key);
             return mh;
         }
+
+        @Override
+        public MethodHandle applyCorrectionsEuclideanBulk() {
+            return APPLY_CORRECTIONS_EUCLIDEAN_HANDLE_BULK;
+        }
+
+        @Override
+        public MethodHandle applyCorrectionsMaxInnerProductBulk() {
+            return APPLY_CORRECTIONS_MAX_INNER_PRODUCT_HANDLE_BULK;
+        }
+
+        @Override
+        public MethodHandle applyCorrectionsDotProductBulk() {
+            return APPLY_CORRECTIONS_DOT_PRODUCT_HANDLE_BULK;
+        }
     }
 }

libs/simdvec/native/publish_vec_binaries.sh

@@ -20,7 +20,7 @@ if [ -z "$ARTIFACTORY_API_KEY" ]; then
   exit 1;
 fi
 
-VERSION="1.0.34"
+VERSION="1.0.36"
 ARTIFACTORY_REPOSITORY="${ARTIFACTORY_REPOSITORY:-https://artifactory.elastic.dev/artifactory/elasticsearch-native/}"
 TEMP=$(mktemp -d)
 

libs/simdvec/native/src/vec/c/aarch64/score_1.cpp

@@ -0,0 +1,141 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+#include <stddef.h>
+#include <arm_neon.h>
+#include <math.h>
+#include <limits>
+#include "vec.h"
+#include "vec_common.h"
+#include "aarch64/aarch64_vec_common.h"
+
+#include "score_common.h"
+
+EXPORT f32_t diskbbq_apply_corrections_euclidean_bulk(
+        const int8_t* corrections,
+		const int32_t bulkSize,
+        const int32_t dimensions,
+        const f32_t queryLowerInterval,
+        const f32_t queryUpperInterval,
+        const int32_t queryComponentSum,
+        const f32_t queryAdditionalCorrection,
+        const f32_t queryBitScale,
+        const f32_t indexBitScale,
+        const f32_t centroidDp,
+        f32_t* scores
+) {
+    f32_t maxScore = -std::numeric_limits<f32_t>::infinity();
+
+    const corrections_t c = unpack_corrections(corrections, bulkSize);
+
+    int i = 0;
+    for (; i < bulkSize; ++i) {
+        f32_t score = apply_corrections_euclidean_inner(
+            dimensions,
+            queryLowerInterval,
+            queryUpperInterval,
+            queryComponentSum,
+            queryAdditionalCorrection,
+            queryBitScale,
+            indexBitScale,
+            centroidDp,
+            *(c.lowerIntervals + i),
+            *(c.upperIntervals + i),
+            *(c.targetComponentSums + i),
+            *(c.additionalCorrections + i),
+            *(scores + i)
+        );
+        *(scores + i) = score;
+        maxScore = fmax(maxScore, score);
+    }
+
+    return maxScore;
+}
+
+EXPORT f32_t diskbbq_apply_corrections_maximum_inner_product_bulk(
+        const int8_t* corrections,
+		const int32_t bulkSize,
+        const int32_t dimensions,
+        const f32_t queryLowerInterval,
+        const f32_t queryUpperInterval,
+        const int32_t queryComponentSum,
+        const f32_t queryAdditionalCorrection,
+        const f32_t queryBitScale,
+        const f32_t indexBitScale,
+        const f32_t centroidDp,
+        f32_t* scores
+) {
+    f32_t maxScore = -std::numeric_limits<f32_t>::infinity();
+
+    const corrections_t c = unpack_corrections(corrections, bulkSize);
+
+    int i = 0;
+    for (; i < bulkSize; ++i) {
+        f32_t score = apply_corrections_maximum_inner_product_inner(
+            dimensions,
+            queryLowerInterval,
+            queryUpperInterval,
+            queryComponentSum,
+            queryAdditionalCorrection,
+            queryBitScale,
+            indexBitScale,
+            centroidDp,
+            *(c.lowerIntervals + i),
+            *(c.upperIntervals + i),
+            *(c.targetComponentSums + i),
+            *(c.additionalCorrections + i),
+            *(scores + i)
+        );
+        *(scores + i) = score;
+        maxScore = fmax(maxScore, score);
+    }
+
+    return maxScore;
+}
+
+EXPORT f32_t diskbbq_apply_corrections_dot_product_bulk(
+        const int8_t* corrections,
+		const int32_t bulkSize,
+        const int32_t dimensions,
+        const f32_t queryLowerInterval,
+        const f32_t queryUpperInterval,
+        const int32_t queryComponentSum,
+        const f32_t queryAdditionalCorrection,
+        const f32_t queryBitScale,
+        const f32_t indexBitScale,
+        const f32_t centroidDp,
+        f32_t* scores
+) {
+    f32_t maxScore = -std::numeric_limits<f32_t>::infinity();
+
+    const corrections_t c = unpack_corrections(corrections, bulkSize);
+
+    int i = 0;
+    for (; i < bulkSize; ++i) {
+        f32_t score = apply_corrections_dot_product_inner(
+            dimensions,
+            queryLowerInterval,
+            queryUpperInterval,
+            queryComponentSum,
+            queryAdditionalCorrection,
+            queryBitScale,
+            indexBitScale,
+            centroidDp,
+            *(c.lowerIntervals + i),
+            *(c.upperIntervals + i),
+            *(c.targetComponentSums + i),
+            *(c.additionalCorrections + i),
+            *(scores + i)
+        );
+        *(scores + i) = score;
+        maxScore = fmax(maxScore, score);
+    }
+
+    return maxScore;
+}