You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by ch...@apache.org on 2025/08/13 19:52:58 UTC
(lucene) branch main updated: Rework off-heap float32 bulk scoring to reduce garbage creation (#15062)
This is an automated email from the ASF dual-hosted git repository.
chegar pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/main by this push:
new bdd5a87caf3 Rework off-heap float32 bulk scoring to reduce garbage creation (#15062)
bdd5a87caf3 is described below
commit bdd5a87caf3ef4fdeb4032a5d9660a70bfd5b817
Author: Chris Hegarty <62...@users.noreply.github.com>
AuthorDate: Wed Aug 13 20:52:53 2025 +0100
Rework off-heap float32 bulk scoring to reduce garbage creation (#15062)
This commit reworks the off-heap float32 bulk scoring to reduce garbage creation. The JVM fails to optimise things as we'd like, so we can just avoid some of the patterns in the first place!
1. Avoid slicing memory segments: our index input can span over up to 16GB by default, so we can just score a single segment (rather than slicing per scoring operation.
2. The hierarchy to allow scoring against an on-heap and off-heap query does not optimise well, so just flatten this and do the most obvious and straightforward thing - passing and switch on a null array or invalid address offset.
---
.../jmh/VectorScorerFloat32Benchmark.java | 20 +++-
.../jmh/TestVectorScorerFloat32Benchmark.java | 2 +
.../Lucene99MemorySegmentFloatVectorScorer.java | 60 ++++------
...ne99MemorySegmentFloatVectorScorerSupplier.java | 72 +++++------
.../vectorization/MemorySegmentBulkVectorOps.java | 133 ++++++++-------------
.../lucene/codecs/hnsw/TestFlatVectorScorer.java | 3 +-
6 files changed, 128 insertions(+), 162 deletions(-)
diff --git a/lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/VectorScorerFloat32Benchmark.java b/lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/VectorScorerFloat32Benchmark.java
index 4f8223f918f..215f754c3fe 100644
--- a/lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/VectorScorerFloat32Benchmark.java
+++ b/lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/VectorScorerFloat32Benchmark.java
@@ -80,6 +80,9 @@ public class VectorScorerFloat32Benchmark {
@Param({"1024"})
public int size;
+ @Param({"true", "false"})
+ public boolean pollute = false;
+
public int numVectors = 128_000;
public int numVectorsToScore = 20_000;
@@ -105,7 +108,6 @@ public class VectorScorerFloat32Benchmark {
out.writeBytes(ba, 0, ba.length);
}
}
- perIterationInit();
}
@Setup(Level.Iteration)
@@ -129,6 +131,10 @@ public class VectorScorerFloat32Benchmark {
List<Integer> list = IntStream.range(0, numVectors).boxed().collect(Collectors.toList());
Collections.shuffle(list, random);
indices = list.stream().limit(numVectorsToScore).mapToInt(i -> i).toArray();
+
+ if (pollute) {
+ pollute(random);
+ }
}
@TearDown
@@ -139,6 +145,18 @@ public class VectorScorerFloat32Benchmark {
Files.delete(path);
}
+ public void pollute(Random random) throws IOException {
+ float[] vec = randomVector(size, random);
+ var opt = FlatVectorScorerUtil.getLucene99FlatVectorsScorer();
+ var scorer = opt.getRandomVectorScorer(DOT_PRODUCT, values.copy(), vec);
+ for (int i = 0; i < 2; i++) {
+ for (int v = 0; v < numVectorsToScore; v++) {
+ scores[v] = scorer.score(indices[v]);
+ }
+ scorer.bulkScore(indices, scores, indices.length);
+ }
+ }
+
@Benchmark
public float[] dotProductDefault() throws IOException {
for (int v = 0; v < numVectorsToScore; v++) {
diff --git a/lucene/benchmark-jmh/src/test/org/apache/lucene/benchmark/jmh/TestVectorScorerFloat32Benchmark.java b/lucene/benchmark-jmh/src/test/org/apache/lucene/benchmark/jmh/TestVectorScorerFloat32Benchmark.java
index 27ac9486aa2..765ef0455a3 100644
--- a/lucene/benchmark-jmh/src/test/org/apache/lucene/benchmark/jmh/TestVectorScorerFloat32Benchmark.java
+++ b/lucene/benchmark-jmh/src/test/org/apache/lucene/benchmark/jmh/TestVectorScorerFloat32Benchmark.java
@@ -32,10 +32,12 @@ public class TestVectorScorerFloat32Benchmark extends LuceneTestCase {
public void setup() throws IOException {
bench = new VectorScorerFloat32Benchmark();
bench.size = 1024;
+ bench.pollute = true;
bench.numVectors = random().nextInt(1, 256);
bench.numVectorsToScore = random().nextInt(bench.numVectors);
delta = 1e-3f * bench.size;
bench.setup();
+ bench.perIterationInit();
}
@After
diff --git a/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorer.java b/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorer.java
index c1fd989aa64..49e07f36a64 100644
--- a/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorer.java
+++ b/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorer.java
@@ -33,9 +33,8 @@ abstract sealed class Lucene99MemorySegmentFloatVectorScorer
final FloatVectorValues values;
final int vectorByteSize;
- final MemorySegmentAccessInput input;
+ final MemorySegment seg;
final float[] query;
- byte[] scratch;
/**
* Return an optional whose value, if present, is the scorer. Otherwise, an empty optional is
@@ -45,15 +44,18 @@ abstract sealed class Lucene99MemorySegmentFloatVectorScorer
VectorSimilarityFunction type,
IndexInput input,
FloatVectorValues values,
- float[] queryVector) {
+ float[] queryVector)
+ throws IOException {
input = FilterIndexInput.unwrapOnlyTest(input);
- if (!(input instanceof MemorySegmentAccessInput msInput)) {
+ MemorySegment seg;
+ if (!(input instanceof MemorySegmentAccessInput msInput
+ && (seg = msInput.segmentSliceOrNull(0L, msInput.length())) != null)) {
return Optional.empty();
}
checkInvariants(values.size(), values.getVectorByteLength(), input);
return switch (type) {
case COSINE -> Optional.empty(); // of(new CosineScorer(msInput, values, queryVector));
- case DOT_PRODUCT -> Optional.of(new DotProductScorer(msInput, values, queryVector));
+ case DOT_PRODUCT -> Optional.of(new DotProductScorer(seg, values, queryVector));
case EUCLIDEAN -> Optional.empty(); // of(new EuclideanScorer(msInput, values, queryVector));
case MAXIMUM_INNER_PRODUCT ->
Optional.empty(); // of(new MaxInnerProductScorer(msInput, values, queryVector));
@@ -61,28 +63,14 @@ abstract sealed class Lucene99MemorySegmentFloatVectorScorer
}
Lucene99MemorySegmentFloatVectorScorer(
- MemorySegmentAccessInput input, FloatVectorValues values, float[] queryVector) {
+ MemorySegment seg, FloatVectorValues values, float[] queryVector) {
super(values);
this.values = values;
- this.input = input;
+ this.seg = seg;
this.vectorByteSize = values.getVectorByteLength();
this.query = queryVector;
}
- final MemorySegment getSegment(int ord) throws IOException {
- checkOrdinal(ord);
- long byteOffset = (long) ord * vectorByteSize;
- MemorySegment seg = input.segmentSliceOrNull(byteOffset, vectorByteSize);
- if (seg == null) {
- if (scratch == null) {
- scratch = new byte[vectorByteSize];
- }
- input.readBytes(byteOffset, scratch, 0, vectorByteSize);
- seg = MemorySegment.ofArray(scratch);
- }
- return seg;
- }
-
static void checkInvariants(int maxOrd, int vectorByteLength, IndexInput input) {
if (input.length() < (long) vectorByteLength * maxOrd) {
throw new IllegalArgumentException("input length is less than expected vector data");
@@ -97,10 +85,12 @@ abstract sealed class Lucene99MemorySegmentFloatVectorScorer
static final class DotProductScorer extends Lucene99MemorySegmentFloatVectorScorer {
- static final MemorySegmentBulkVectorOps.DotFromQueryArray DOT_OPS =
- new MemorySegmentBulkVectorOps.DotFromQueryArray();
+ static final MemorySegmentBulkVectorOps.DotProduct DOT_OPS =
+ MemorySegmentBulkVectorOps.DOT_INSTANCE;
+
+ final float[] scratchScores = new float[4];
- DotProductScorer(MemorySegmentAccessInput input, FloatVectorValues values, float[] query) {
+ DotProductScorer(MemorySegment input, FloatVectorValues values, float[] query) {
super(input, values, query);
}
@@ -112,16 +102,16 @@ abstract sealed class Lucene99MemorySegmentFloatVectorScorer
}
@Override
- public void bulkScore(int[] nodes, float[] scores, int numNodes) throws IOException {
- float[] scratchScores = new float[4];
+ public void bulkScore(int[] nodes, float[] scores, int numNodes) {
int i = 0;
final int limit = numNodes & ~3;
for (; i < limit; i += 4) {
- MemorySegment ms1 = getSegment(nodes[i]);
- MemorySegment ms2 = getSegment(nodes[i + 1]);
- MemorySegment ms3 = getSegment(nodes[i + 2]);
- MemorySegment ms4 = getSegment(nodes[i + 3]);
- DOT_OPS.dotProductBulk(scratchScores, query, ms1, ms2, ms3, ms4, query.length);
+ long offset1 = (long) nodes[i] * vectorByteSize;
+ long offset2 = (long) nodes[i + 1] * vectorByteSize;
+ long offset3 = (long) nodes[i + 2] * vectorByteSize;
+ long offset4 = (long) nodes[i + 3] * vectorByteSize;
+ DOT_OPS.dotProductBulk(
+ seg, scratchScores, query, offset1, offset2, offset3, offset4, query.length);
scores[i + 0] = normalizeToUnitInterval(scratchScores[0]);
scores[i + 1] = normalizeToUnitInterval(scratchScores[1]);
scores[i + 2] = normalizeToUnitInterval(scratchScores[2]);
@@ -130,10 +120,10 @@ abstract sealed class Lucene99MemorySegmentFloatVectorScorer
// Handle remaining 1–3 nodes in bulk (if any)
int remaining = numNodes - i;
if (remaining > 0) {
- MemorySegment ms1 = getSegment(nodes[i]);
- MemorySegment ms2 = (remaining > 1) ? getSegment(nodes[i + 1]) : ms1;
- MemorySegment ms3 = (remaining > 2) ? getSegment(nodes[i + 2]) : ms1;
- DOT_OPS.dotProductBulk(scratchScores, query, ms1, ms2, ms3, ms1, query.length);
+ long addr1 = (long) nodes[i] * vectorByteSize;
+ long addr2 = (remaining > 1) ? (long) nodes[i + 1] * vectorByteSize : addr1;
+ long addr3 = (remaining > 2) ? (long) nodes[i + 2] * vectorByteSize : addr1;
+ DOT_OPS.dotProductBulk(seg, scratchScores, query, addr1, addr2, addr3, addr1, query.length);
scores[i] = normalizeToUnitInterval(scratchScores[0]);
if (remaining > 1) scores[i + 1] = normalizeToUnitInterval(scratchScores[1]);
if (remaining > 2) scores[i + 2] = normalizeToUnitInterval(scratchScores[2]);
diff --git a/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorerSupplier.java b/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorerSupplier.java
index 33a3a0e33dd..61417452a1e 100644
--- a/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorerSupplier.java
+++ b/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorerSupplier.java
@@ -35,33 +35,34 @@ public abstract sealed class Lucene99MemorySegmentFloatVectorScorerSupplier
final int vectorByteSize;
final int maxOrd;
final int dims;
- final MemorySegmentAccessInput input;
+ final MemorySegment seg;
final FloatVectorValues values; // to support ordToDoc/getAcceptOrds
- byte[] queryScratch, scratch1, scratch2, scratch3, scratch4;
/**
* Return an optional whose value, if present, is the scorer supplier. Otherwise, an empty
* optional is returned.
*/
static Optional<RandomVectorScorerSupplier> create(
- VectorSimilarityFunction type, IndexInput input, FloatVectorValues values) {
+ VectorSimilarityFunction type, IndexInput input, FloatVectorValues values)
+ throws IOException {
input = FilterIndexInput.unwrapOnlyTest(input);
- if (!(input instanceof MemorySegmentAccessInput msInput)) {
+ MemorySegment seg;
+ if (!(input instanceof MemorySegmentAccessInput msInput
+ && (seg = msInput.segmentSliceOrNull(0L, msInput.length())) != null)) {
return Optional.empty();
}
checkInvariants(values.size(), values.getVectorByteLength(), input);
return switch (type) {
case COSINE -> Optional.empty(); // of(new CosineSupplier(msInput, values));
- case DOT_PRODUCT -> Optional.of(new DotProductSupplier(msInput, values));
+ case DOT_PRODUCT -> Optional.of(new DotProductSupplier(seg, values));
case EUCLIDEAN -> Optional.empty(); // of(new EuclideanSupplier(msInput, values));
case MAXIMUM_INNER_PRODUCT ->
Optional.empty(); // of(new MaxInnerProductSupplier(msInput, values));
};
}
- Lucene99MemorySegmentFloatVectorScorerSupplier(
- MemorySegmentAccessInput input, FloatVectorValues values) {
- this.input = input;
+ Lucene99MemorySegmentFloatVectorScorerSupplier(MemorySegment seg, FloatVectorValues values) {
+ this.seg = seg;
this.values = values;
this.vectorByteSize = values.getVectorByteLength();
this.maxOrd = values.size();
@@ -80,26 +81,15 @@ public abstract sealed class Lucene99MemorySegmentFloatVectorScorerSupplier
}
}
- final MemorySegment getSegment(int ord, byte[] scratch) throws IOException {
- long byteOffset = (long) ord * vectorByteSize;
- MemorySegment seg = input.segmentSliceOrNull(byteOffset, vectorByteSize);
- if (seg == null) {
- if (scratch == null) {
- scratch = new byte[vectorByteSize];
- }
- input.readBytes(byteOffset, scratch, 0, vectorByteSize);
- seg = MemorySegment.ofArray(scratch);
- }
- return seg;
- }
-
static final class DotProductSupplier extends Lucene99MemorySegmentFloatVectorScorerSupplier {
- static final MemorySegmentBulkVectorOps.DotFromQuerySegment DOT_OPS =
- new MemorySegmentBulkVectorOps.DotFromQuerySegment();
+ static final MemorySegmentBulkVectorOps.DotProduct DOT_OPS =
+ MemorySegmentBulkVectorOps.DOT_INSTANCE;
+
+ final float[] scratchScores = new float[4];
- DotProductSupplier(MemorySegmentAccessInput input, FloatVectorValues values) {
- super(input, values);
+ DotProductSupplier(MemorySegment seg, FloatVectorValues values) {
+ super(seg, values);
}
@Override
@@ -110,25 +100,25 @@ public abstract sealed class Lucene99MemorySegmentFloatVectorScorerSupplier
@Override
public float score(int node) throws IOException {
checkOrdinal(node);
- MemorySegment query = getSegment(queryOrd, queryScratch);
- MemorySegment ms = getSegment(node, scratch1);
- var raw = DOT_OPS.dotProduct(query, ms, dims);
+ long queryAddr = (long) queryOrd * vectorByteSize;
+ long addr = (long) node * vectorByteSize;
+ var raw = DOT_OPS.dotProduct(seg, queryAddr, addr, dims);
return normalizeToUnitInterval(raw);
}
@Override
- public void bulkScore(int[] nodes, float[] scores, int numNodes) throws IOException {
+ public void bulkScore(int[] nodes, float[] scores, int numNodes) {
// TODO checkOrdinal(node1 ....);
- float[] scratchScores = new float[4];
int i = 0;
- MemorySegment query = getSegment(queryOrd, queryScratch);
+ long queryAddr = (long) queryOrd * vectorByteSize;
final int limit = numNodes & ~3;
for (; i < limit; i += 4) {
- MemorySegment ms1 = getSegment(nodes[i], scratch1);
- MemorySegment ms2 = getSegment(nodes[i + 1], scratch2);
- MemorySegment ms3 = getSegment(nodes[i + 2], scratch3);
- MemorySegment ms4 = getSegment(nodes[i + 3], scratch4);
- DOT_OPS.dotProductBulk(scratchScores, query, ms1, ms2, ms3, ms4, dims);
+ long offset1 = (long) nodes[i + 0] * vectorByteSize;
+ long offset2 = (long) nodes[i + 1] * vectorByteSize;
+ long offset3 = (long) nodes[i + 2] * vectorByteSize;
+ long offset4 = (long) nodes[i + 3] * vectorByteSize;
+ DOT_OPS.dotProductBulk(
+ seg, scratchScores, queryAddr, offset1, offset2, offset3, offset4, dims);
scores[i + 0] = normalizeToUnitInterval(scratchScores[0]);
scores[i + 1] = normalizeToUnitInterval(scratchScores[1]);
scores[i + 2] = normalizeToUnitInterval(scratchScores[2]);
@@ -137,10 +127,10 @@ public abstract sealed class Lucene99MemorySegmentFloatVectorScorerSupplier
// Handle remaining 1–3 nodes in bulk (if any)
int remaining = numNodes - i;
if (remaining > 0) {
- MemorySegment ms1 = getSegment(nodes[i], scratch1);
- MemorySegment ms2 = (remaining > 1) ? getSegment(nodes[i + 1], scratch2) : ms1;
- MemorySegment ms3 = (remaining > 2) ? getSegment(nodes[i + 2], scratch3) : ms1;
- DOT_OPS.dotProductBulk(scratchScores, query, ms1, ms2, ms3, ms1, dims);
+ long addr1 = (long) nodes[i] * vectorByteSize;
+ long addr2 = (remaining > 1) ? (long) nodes[i + 1] * vectorByteSize : addr1;
+ long addr3 = (remaining > 2) ? (long) nodes[i + 2] * vectorByteSize : addr1;
+ DOT_OPS.dotProductBulk(seg, scratchScores, queryAddr, addr1, addr2, addr3, addr3, dims);
scores[i] = normalizeToUnitInterval(scratchScores[0]);
if (remaining > 1) scores[i + 1] = normalizeToUnitInterval(scratchScores[1]);
if (remaining > 2) scores[i + 2] = normalizeToUnitInterval(scratchScores[2]);
@@ -157,7 +147,7 @@ public abstract sealed class Lucene99MemorySegmentFloatVectorScorerSupplier
@Override
public DotProductSupplier copy() throws IOException {
- return new DotProductSupplier(input.clone(), values);
+ return new DotProductSupplier(seg, values);
}
}
}
diff --git a/lucene/core/src/java24/org/apache/lucene/internal/vectorization/MemorySegmentBulkVectorOps.java b/lucene/core/src/java24/org/apache/lucene/internal/vectorization/MemorySegmentBulkVectorOps.java
index 8c6e5b096b7..20aefce4831 100644
--- a/lucene/core/src/java24/org/apache/lucene/internal/vectorization/MemorySegmentBulkVectorOps.java
+++ b/lucene/core/src/java24/org/apache/lucene/internal/vectorization/MemorySegmentBulkVectorOps.java
@@ -35,111 +35,66 @@ public final class MemorySegmentBulkVectorOps {
static final ByteOrder LE = ByteOrder.LITTLE_ENDIAN;
static final ValueLayout.OfFloat LAYOUT_LE_FLOAT = ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(LE);
- private MemorySegmentBulkVectorOps() {}
-
- sealed interface FloatVectorLoader {
- /** Returns the number of float elements. */
- int length();
-
- FloatVector load(VectorSpecies<Float> species, int index);
-
- float tail(int index);
- }
-
- record FloatArrayLoader(float[] arr) implements FloatVectorLoader {
- @Override
- public int length() {
- return arr.length;
- }
-
- @Override
- public FloatVector load(VectorSpecies<Float> species, int index) {
- // assert index + species.length() <= length();
- return FloatVector.fromArray(species, arr, index);
- }
-
- @Override
- public float tail(int index) {
- // assert index <= length();
- return arr[index];
- }
- }
+ public static final DotProduct DOT_INSTANCE = new DotProduct();
- record FloatMemorySegmentLoader(MemorySegment segment) implements FloatVectorLoader {
- @Override
- public int length() {
- return Math.toIntExact(segment.byteSize() / Float.BYTES);
- }
+ private MemorySegmentBulkVectorOps() {}
- @Override
- public FloatVector load(VectorSpecies<Float> species, int index) {
- // assert index + species.length() <= length();
- return FloatVector.fromMemorySegment(species, segment, (long) index * Float.BYTES, LE);
- }
+ public static final class DotProduct {
- @Override
- public float tail(int index) {
- // assert index <= length();
- return segment.get(LAYOUT_LE_FLOAT, (long) index * Float.BYTES);
- }
- }
-
- public static final class DotFromQueryArray extends AbstractDotProduct {
+ private DotProduct() {}
public void dotProductBulk(
+ MemorySegment dataSeg,
float[] scores,
float[] q,
- MemorySegment d1,
- MemorySegment d2,
- MemorySegment d3,
- MemorySegment d4,
+ long d1,
+ long d2,
+ long d3,
+ long d4,
int elementCount) {
- super.dotProductBulk(scores, new FloatArrayLoader(q), d1, d2, d3, d4, elementCount);
+ dotProductBulkImpl(dataSeg, scores, q, -1L, d1, d2, d3, d4, elementCount);
}
- }
-
- public static final class DotFromQuerySegment extends AbstractDotProduct {
public void dotProductBulk(
+ MemorySegment seg,
float[] scores,
- MemorySegment q,
- MemorySegment d1,
- MemorySegment d2,
- MemorySegment d3,
- MemorySegment d4,
+ long q,
+ long d1,
+ long d2,
+ long d3,
+ long d4,
int elementCount) {
- super.dotProductBulk(scores, new FloatMemorySegmentLoader(q), d1, d2, d3, d4, elementCount);
+ dotProductBulkImpl(seg, scores, null, q, d1, d2, d3, d4, elementCount);
}
- public float dotProduct(MemorySegment q, MemorySegment d, int elementCount) {
+ public float dotProduct(MemorySegment seg, long q, long d, int elementCount) {
int i = 0;
FloatVector sv = FloatVector.zero(FLOAT_SPECIES);
final int limit = FLOAT_SPECIES.loopBound(elementCount);
for (; i < limit; i += FLOAT_SPECIES.length()) {
final long offset = (long) i * Float.BYTES;
- FloatVector qv = FloatVector.fromMemorySegment(FLOAT_SPECIES, q, offset, LE);
- FloatVector dv = FloatVector.fromMemorySegment(FLOAT_SPECIES, d, offset, LE);
+ FloatVector qv = FloatVector.fromMemorySegment(FLOAT_SPECIES, seg, q + offset, LE);
+ FloatVector dv = FloatVector.fromMemorySegment(FLOAT_SPECIES, seg, d + offset, LE);
sv = fma(qv, dv, sv);
}
float score = sv.reduceLanes(VectorOperators.ADD);
for (; i < elementCount; i++) {
final long offset = (long) i * Float.BYTES;
- score += q.get(LAYOUT_LE_FLOAT, offset) * d.get(LAYOUT_LE_FLOAT, offset);
+ score += seg.get(LAYOUT_LE_FLOAT, q + offset) * seg.get(LAYOUT_LE_FLOAT, d + offset);
}
return score;
}
- }
-
- abstract static class AbstractDotProduct {
- protected void dotProductBulk(
+ private void dotProductBulkImpl(
+ MemorySegment seg,
float[] scores,
- FloatVectorLoader q,
- MemorySegment d1,
- MemorySegment d2,
- MemorySegment d3,
- MemorySegment d4,
+ float[] qArray,
+ long qOffset,
+ long d1,
+ long d2,
+ long d3,
+ long d4,
int elementCount) {
int i = 0;
FloatVector sv1 = FloatVector.zero(FLOAT_SPECIES);
@@ -150,11 +105,16 @@ public final class MemorySegmentBulkVectorOps {
final int limit = FLOAT_SPECIES.loopBound(elementCount);
for (; i < limit; i += FLOAT_SPECIES.length()) {
final long offset = (long) i * Float.BYTES;
- FloatVector dv1 = FloatVector.fromMemorySegment(FLOAT_SPECIES, d1, offset, LE);
- FloatVector dv2 = FloatVector.fromMemorySegment(FLOAT_SPECIES, d2, offset, LE);
- FloatVector dv3 = FloatVector.fromMemorySegment(FLOAT_SPECIES, d3, offset, LE);
- FloatVector dv4 = FloatVector.fromMemorySegment(FLOAT_SPECIES, d4, offset, LE);
- FloatVector qv = q.load(FLOAT_SPECIES, i);
+ FloatVector dv1 = FloatVector.fromMemorySegment(FLOAT_SPECIES, seg, d1 + offset, LE);
+ FloatVector dv2 = FloatVector.fromMemorySegment(FLOAT_SPECIES, seg, d2 + offset, LE);
+ FloatVector dv3 = FloatVector.fromMemorySegment(FLOAT_SPECIES, seg, d3 + offset, LE);
+ FloatVector dv4 = FloatVector.fromMemorySegment(FLOAT_SPECIES, seg, d4 + offset, LE);
+ FloatVector qv;
+ if (qOffset == -1L) {
+ qv = FloatVector.fromArray(FLOAT_SPECIES, qArray, i);
+ } else {
+ qv = FloatVector.fromMemorySegment(FLOAT_SPECIES, seg, qOffset + offset, LE);
+ }
sv1 = fma(qv, dv1, sv1);
sv2 = fma(qv, dv2, sv2);
sv3 = fma(qv, dv3, sv3);
@@ -167,11 +127,16 @@ public final class MemorySegmentBulkVectorOps {
for (; i < elementCount; i++) {
final long offset = (long) i * Float.BYTES;
- final float qValue = q.tail(i);
- sum1 = fma(qValue, d1.get(LAYOUT_LE_FLOAT, offset), sum1);
- sum2 = fma(qValue, d2.get(LAYOUT_LE_FLOAT, offset), sum2);
- sum3 = fma(qValue, d3.get(LAYOUT_LE_FLOAT, offset), sum3);
- sum4 = fma(qValue, d4.get(LAYOUT_LE_FLOAT, offset), sum4);
+ float qValue;
+ if (qOffset == -1L) {
+ qValue = qArray[i];
+ } else {
+ qValue = seg.get(LAYOUT_LE_FLOAT, qOffset + offset);
+ }
+ sum1 = fma(qValue, seg.get(LAYOUT_LE_FLOAT, d1 + offset), sum1);
+ sum2 = fma(qValue, seg.get(LAYOUT_LE_FLOAT, d2 + offset), sum2);
+ sum3 = fma(qValue, seg.get(LAYOUT_LE_FLOAT, d3 + offset), sum3);
+ sum4 = fma(qValue, seg.get(LAYOUT_LE_FLOAT, d4 + offset), sum4);
}
scores[0] = sum1;
scores[1] = sum2;
diff --git a/lucene/core/src/test/org/apache/lucene/codecs/hnsw/TestFlatVectorScorer.java b/lucene/core/src/test/org/apache/lucene/codecs/hnsw/TestFlatVectorScorer.java
index 54dc8015dde..dada4a7d5a6 100644
--- a/lucene/core/src/test/org/apache/lucene/codecs/hnsw/TestFlatVectorScorer.java
+++ b/lucene/core/src/test/org/apache/lucene/codecs/hnsw/TestFlatVectorScorer.java
@@ -77,7 +77,8 @@ public class TestFlatVectorScorer extends BaseVectorizationTestCase {
var dirs =
List.<IOSupplier<Directory>>of(
TestFlatVectorScorer::newDirectory,
- () -> new MMapDirectory(createTempDir(count.getAndIncrement() + "-")));
+ () -> new MMapDirectory(createTempDir(count.getAndIncrement() + "-")),
+ () -> new MMapDirectory(createTempDir(count.getAndIncrement() + "-"), 1024));
List<Object[]> objs = new ArrayList<>();
for (var scorer : scorers) {