From 0fd6f9c7200e72add78758d0041e02fe3064da7a Mon Sep 17 00:00:00 2001 From: CH-Abhinav Date: Sat, 5 Sep 2026 17:18:11 +0000 Subject: [PATCH 1/2] bulk scoring API into VectorUtilSupport --- .../DefaultVectorUtilSupport.java | 7 +++ .../vectorization/VectorUtilSupport.java | 9 ++++ .../org/apache/lucene/util/VectorUtil.java | 28 ++++++++++ .../NativeVectorUtilSupport.java | 5 ++ .../PanamaVectorUtilSupport.java | 51 +++++++++++++++++++ 5 files changed, 100 insertions(+) diff --git a/lucene/core/src/java/org/apache/lucene/internal/vectorization/DefaultVectorUtilSupport.java b/lucene/core/src/java/org/apache/lucene/internal/vectorization/DefaultVectorUtilSupport.java index cc11918baf03..eb79514918f9 100644 --- a/lucene/core/src/java/org/apache/lucene/internal/vectorization/DefaultVectorUtilSupport.java +++ b/lucene/core/src/java/org/apache/lucene/internal/vectorization/DefaultVectorUtilSupport.java @@ -75,6 +75,13 @@ public float dotProduct(short[] a, short[] b) { return sum; } + @Override + public void dotProductBulk(float[] query, float[][] docs, float[] results) { + for (int i = 0; i < docs.length; i++) { + results[i] = dotProduct(query, docs[i]); + } + } + @Override public float cosine(float[] a, float[] b) { float sum = 0.0f; diff --git a/lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorUtilSupport.java b/lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorUtilSupport.java index 5569d77a337d..2acfc322e48f 100644 --- a/lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorUtilSupport.java +++ b/lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorUtilSupport.java @@ -32,6 +32,15 @@ public interface VectorUtilSupport { */ float dotProduct(short[] a, short[] b); + /** + * Calculate the dot product of a query vector against a batch of document vectors + * + * @param query The query vector + * @param docs Array of document vectors to score against + * @param results Result vector which will be overwritten + */ + void dotProductBulk(float[] query, float[][] docs, float[] results); + /** Returns the cosine similarity between the two vectors. */ float cosine(float[] v1, float[] v2); diff --git a/lucene/core/src/java/org/apache/lucene/util/VectorUtil.java b/lucene/core/src/java/org/apache/lucene/util/VectorUtil.java index 641cc420bf67..120839e1d7f1 100644 --- a/lucene/core/src/java/org/apache/lucene/util/VectorUtil.java +++ b/lucene/core/src/java/org/apache/lucene/util/VectorUtil.java @@ -92,6 +92,34 @@ public static float dotProduct(short[] a, short[] b) { return result; } + /** + * Calculates the dot product of a query vector against a batch of document vectors. + * + * @param query The query vector + * @param docs Array of document vectors to score against + * @param results Pre-allocated array where the resulting scores will be written + * @throws IllegalArgumentException if the query dimensions do not match the document dimensions, + * or if the results array is incorrectly sized. + */ + public static void dotProductBulk(float[] query, float[][] docs, float[] results) { + if (docs.length != results.length) { + throw new IllegalArgumentException( + "results array length (" + + results.length + + ") must match docs array length (" + + docs.length + + ")"); + } + if (docs.length == 0) { + return; + } + if (query.length != docs[0].length) { + throw new IllegalArgumentException( + "vector dimensions differ: " + query.length + "!=" + docs[0].length); + } + IMPL.dotProductBulk(query, docs, results); + } + /** * Returns the cosine similarity between the two vectors. * diff --git a/lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java b/lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java index e35b472f71d5..83bec31e0489 100644 --- a/lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java +++ b/lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java @@ -404,6 +404,11 @@ public float squareDistance(float[] a, float[] b) { : delegateVectorUtilSupport.squareDistance(a, b); } + @Override + public void dotProductBulk(float[] query, float[][] docs, float[] results) { + delegateVectorUtilSupport.dotProductBulk(query, docs, results); + } + @Override public int dotProduct(byte[] a, byte[] b) { return (dotProduct$MH != null) diff --git a/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java b/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java index 4ef2dca2512f..74f9c6bb832e 100644 --- a/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java +++ b/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java @@ -169,6 +169,57 @@ private float dotProductBody(float[] a, float[] b, int limit) { return res1.add(res2).reduceLanes(ADD); } + /** */ + @Override + public void dotProductBulk(float[] query, float[][] docs, float[] results) { + int docIdx = 0; + + // Process in batches of 4 documents at a time + for (; docIdx + 3 < docs.length; docIdx += 4) { + float[] d0 = docs[docIdx]; + float[] d1 = docs[docIdx + 1]; + float[] d2 = docs[docIdx + 2]; + float[] d3 = docs[docIdx + 3]; + + FloatVector acc0 = FloatVector.zero(FLOAT_SPECIES); + FloatVector acc1 = FloatVector.zero(FLOAT_SPECIES); + FloatVector acc2 = FloatVector.zero(FLOAT_SPECIES); + FloatVector acc3 = FloatVector.zero(FLOAT_SPECIES); + + int i = 0; + int limit = FLOAT_SPECIES.loopBound(query.length); + + for (; i < limit; i += FLOAT_SPECIES.length()) { + // LOAD QUERY ONCE + FloatVector vq = FloatVector.fromArray(FLOAT_SPECIES, query, i); + + // SCORE AGAINST 4 DOCUMENTS + acc0 = fma(vq, FloatVector.fromArray(FLOAT_SPECIES, d0, i), acc0); + acc1 = fma(vq, FloatVector.fromArray(FLOAT_SPECIES, d1, i), acc1); + acc2 = fma(vq, FloatVector.fromArray(FLOAT_SPECIES, d2, i), acc2); + acc3 = fma(vq, FloatVector.fromArray(FLOAT_SPECIES, d3, i), acc3); + } + + results[docIdx] = acc0.reduceLanes(ADD) + scalarDotProductBulk(query, d0, i); + results[docIdx + 1] = acc1.reduceLanes(ADD) + scalarDotProductBulk(query, d1, i); + results[docIdx + 2] = acc2.reduceLanes(ADD) + scalarDotProductBulk(query, d2, i); + results[docIdx + 3] = acc3.reduceLanes(ADD) + scalarDotProductBulk(query, d3, i); + } + + for (; docIdx < docs.length; docIdx++) { + results[docIdx] = dotProduct(query, docs[docIdx]); + } + } + + /** Helper method to process the remaining un-vectorized tail of the arrays */ + private float scalarDotProductBulk(float[] a, float[] b, int startIndex) { + float res = 0; + for (int i = startIndex; i < a.length; i++) { + res = fma(a[i], b[i], res); + } + return res; + } + @Override public float cosine(float[] a, float[] b) { int i = 0; From 9964fdd6f886d4fc7d24a547528ee3576f2f0bc0 Mon Sep 17 00:00:00 2001 From: CH-Abhinav Date: Sun, 13 Sep 2026 06:44:27 +0000 Subject: [PATCH 2/2] off-heap bulk scoring wired to Panama SIMD --- .../MemorySegmentBulkVectorOps.java | 3 +- .../PanamaVectorUtilSupport.java | 66 +++++++++++++++++++ 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/lucene/core/src/java25/org/apache/lucene/internal/vectorization/MemorySegmentBulkVectorOps.java b/lucene/core/src/java25/org/apache/lucene/internal/vectorization/MemorySegmentBulkVectorOps.java index 0447389d7b81..75c0b63a0d23 100644 --- a/lucene/core/src/java25/org/apache/lucene/internal/vectorization/MemorySegmentBulkVectorOps.java +++ b/lucene/core/src/java25/org/apache/lucene/internal/vectorization/MemorySegmentBulkVectorOps.java @@ -54,7 +54,8 @@ public void dotProductBulk( long d3, long d4, int elementCount) { - dotProductBulkImpl(dataSeg, scores, q, -1L, d1, d2, d3, d4, elementCount); + // Delegate directly to our new Panama SIMD engine + PanamaVectorUtilSupport.dotProductBulk(dataSeg, scores, q, d1, d2, d3, d4, elementCount); } public void dotProductBulk( diff --git a/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java b/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java index 74f9c6bb832e..e11700a63363 100644 --- a/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java +++ b/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java @@ -29,6 +29,7 @@ import static org.apache.lucene.util.VectorUtil.isUnitVector; import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; import jdk.incubator.vector.ByteVector; import jdk.incubator.vector.DoubleVector; import jdk.incubator.vector.FloatVector; @@ -220,6 +221,71 @@ private float scalarDotProductBulk(float[] a, float[] b, int startIndex) { return res; } + /** + * SIMD accelerated bulk dot product for off-heap memory segments. Called directly by + * MemorySegmentBulkVectorOps. + */ + public static void dotProductBulk( + MemorySegment seg, + float[] scores, + float[] query, + long node1Offset, + long node2Offset, + long node3Offset, + long node4Offset, + int dims) { + + FloatVector acc0 = FloatVector.zero(FLOAT_SPECIES); + FloatVector acc1 = FloatVector.zero(FLOAT_SPECIES); + FloatVector acc2 = FloatVector.zero(FLOAT_SPECIES); + FloatVector acc3 = FloatVector.zero(FLOAT_SPECIES); + + int i = 0; + int limit = FLOAT_SPECIES.loopBound(dims); + + for (; i < limit; i += FLOAT_SPECIES.length()) { + FloatVector vq = FloatVector.fromArray(FLOAT_SPECIES, query, i); + + // Load 4 off-heap document vectors directly into registers + FloatVector v0 = + FloatVector.fromMemorySegment( + FLOAT_SPECIES, seg, node1Offset + (long) i * Float.BYTES, LITTLE_ENDIAN); + FloatVector v1 = + FloatVector.fromMemorySegment( + FLOAT_SPECIES, seg, node2Offset + (long) i * Float.BYTES, LITTLE_ENDIAN); + FloatVector v2 = + FloatVector.fromMemorySegment( + FLOAT_SPECIES, seg, node3Offset + (long) i * Float.BYTES, LITTLE_ENDIAN); + FloatVector v3 = + FloatVector.fromMemorySegment( + FLOAT_SPECIES, seg, node4Offset + (long) i * Float.BYTES, LITTLE_ENDIAN); + + acc0 = fma(vq, v0, acc0); + acc1 = fma(vq, v1, acc1); + acc2 = fma(vq, v2, acc2); + acc3 = fma(vq, v3, acc3); + } + + float res0 = acc0.reduceLanes(ADD); + float res1 = acc1.reduceLanes(ADD); + float res2 = acc2.reduceLanes(ADD); + float res3 = acc3.reduceLanes(ADD); + + // Scalar tail for remainder dimensions + for (; i < dims; i++) { + float q = query[i]; + res0 = fma(q, seg.get(ValueLayout.JAVA_FLOAT, node1Offset + (long) i * Float.BYTES), res0); + res1 = fma(q, seg.get(ValueLayout.JAVA_FLOAT, node2Offset + (long) i * Float.BYTES), res1); + res2 = fma(q, seg.get(ValueLayout.JAVA_FLOAT, node3Offset + (long) i * Float.BYTES), res2); + res3 = fma(q, seg.get(ValueLayout.JAVA_FLOAT, node4Offset + (long) i * Float.BYTES), res3); + } + + scores[0] = res0; + scores[1] = res1; + scores[2] = res2; + scores[3] = res3; + } + @Override public float cosine(float[] a, float[] b) { int i = 0;