You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by jp...@apache.org on 2025/08/18 05:44:02 UTC
(lucene) branch main updated: Better vectorize score computations. (#15039)
This is an automated email from the ASF dual-hosted git repository.
jpountz 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 772f63c0a4b Better vectorize score computations. (#15039)
772f63c0a4b is described below
commit 772f63c0a4ba507fc0b26307c0d8cdac22ff3957
Author: Adrien Grand <jp...@gmail.com>
AuthorDate: Mon Aug 18 07:43:56 2025 +0200
Better vectorize score computations. (#15039)
Existing auto-vectorization of scores is a bit fragile since it relies on
`SimScorer#score` being inlined in the for loops where it is called. This is
currently the case in nightly benchmarks, but may not be the case in the real
world where more implementations of `SimScorer` may be used, in particular
those from `FeatureField`.
Furthermore, existing auto-vectorization has some room for improvement as
@gf2121 highlighted at
https://github.com/apache/lucene/pull/14679#discussion_r2093432298.
---
lucene/CHANGES.txt | 3 ++
.../java/org/apache/lucene/search/TermScorer.java | 10 ++---
.../lucene/search/similarities/BM25Similarity.java | 34 ++++++++++++++--
.../lucene/search/similarities/Similarity.java | 45 ++++++++++++++++++++++
.../search/similarities/AssertingSimilarity.java | 20 ++++++++++
.../similarities/BaseSimilarityTestCase.java | 39 +++++++++++++++++++
6 files changed, 143 insertions(+), 8 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 8e8ee1870d0..cde538fd534 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -243,6 +243,9 @@ Optimizations
* GITHUB#15004: Wraps all iterator with likelyImpactsEnum under BlockMaxConjunctionBulkScorer. (Ge Song)
+* GITHUB#15039: Score computations are now more reliably vectorized.
+ (Adrien Grand, Guo Feng)
+
Changes in Runtime Behavior
---------------------
* GITHUB#14823: Decrease TieredMergePolicy's default number of segments per
diff --git a/lucene/core/src/java/org/apache/lucene/search/TermScorer.java b/lucene/core/src/java/org/apache/lucene/search/TermScorer.java
index 278f735838e..46cf1753441 100644
--- a/lucene/core/src/java/org/apache/lucene/search/TermScorer.java
+++ b/lucene/core/src/java/org/apache/lucene/search/TermScorer.java
@@ -22,6 +22,7 @@ import org.apache.lucene.index.ImpactsEnum;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.SlowImpactsEnum;
+import org.apache.lucene.search.similarities.Similarity.BulkSimScorer;
import org.apache.lucene.search.similarities.Similarity.SimScorer;
import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.Bits;
@@ -36,6 +37,7 @@ public final class TermScorer extends Scorer {
private final PostingsEnum postingsEnum;
private final DocIdSetIterator iterator;
private final SimScorer scorer;
+ private final BulkSimScorer bulkScorer;
private final NumericDocValues norms;
private final ImpactsDISI impactsDisi;
private final MaxScoreCache maxScoreCache;
@@ -49,6 +51,7 @@ public final class TermScorer extends Scorer {
impactsDisi = null;
this.scorer = scorer;
this.norms = norms;
+ this.bulkScorer = scorer.asBulkSimScorer();
}
/**
@@ -71,6 +74,7 @@ public final class TermScorer extends Scorer {
}
this.scorer = scorer;
this.norms = norms;
+ this.bulkScorer = scorer.asBulkSimScorer();
}
@Override
@@ -165,10 +169,6 @@ public final class TermScorer extends Scorer {
}
}
- for (int i = 0; i < size; ++i) {
- // Unless SimScorer#score is megamorphic, SimScorer#score should inline and (part of) score
- // computations should auto-vectorize.
- buffer.features[i] = scorer.score(buffer.features[i], normValues[i]);
- }
+ bulkScorer.score(buffer.size, buffer.features, normValues, buffer.features);
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/similarities/BM25Similarity.java b/lucene/core/src/java/org/apache/lucene/search/similarities/BM25Similarity.java
index 274bb475d6c..78e89f51516 100644
--- a/lucene/core/src/java/org/apache/lucene/search/similarities/BM25Similarity.java
+++ b/lucene/core/src/java/org/apache/lucene/search/similarities/BM25Similarity.java
@@ -21,6 +21,7 @@ import java.util.List;
import org.apache.lucene.search.CollectionStatistics;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.TermStatistics;
+import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.SmallFloat;
/**
@@ -217,8 +218,7 @@ public class BM25Similarity extends Similarity {
this.weight = boost * idf.getValue().floatValue();
}
- @Override
- public float score(float freq, long encodedNorm) {
+ private float doScore(float freq, float normInverse) {
// In order to guarantee monotonicity with both freq and norm without
// promoting to doubles, we rewrite freq / (freq + norm) to
// 1 - 1 / (1 + freq * 1/norm).
@@ -228,10 +228,38 @@ public class BM25Similarity extends Similarity {
// x -> 1 + x and x -> 1 - 1/x.
// Finally we expand weight * (1 - 1 / (1 + freq * 1/norm)) to
// weight - weight / (1 + freq * 1/norm), which runs slightly faster.
- float normInverse = cache[((byte) encodedNorm) & 0xFF];
return weight - weight / (1f + freq * normInverse);
}
+ @Override
+ public float score(float freq, long encodedNorm) {
+ float normInverse = cache[((byte) encodedNorm) & 0xFF];
+ return doScore(freq, normInverse);
+ }
+
+ @Override
+ public BulkSimScorer asBulkSimScorer() {
+ return new BulkSimScorer() {
+
+ private float[] normInverses = new float[0];
+
+ @Override
+ public void score(int size, float[] freqs, long[] norms, float[] scores) {
+ if (normInverses.length < size) {
+ normInverses = new float[ArrayUtil.oversize(size, Float.BYTES)];
+ }
+ for (int i = 0; i < size; ++i) {
+ normInverses[i] = cache[((byte) norms[i]) & 0xFF];
+ }
+
+ // This loop auto-vectorizes.
+ for (int i = 0; i < size; ++i) {
+ scores[i] = doScore(freqs[i], normInverses[i]);
+ }
+ }
+ };
+ }
+
@Override
public Explanation explain(Explanation freq, long encodedNorm) {
List<Explanation> subs = new ArrayList<>(explainConstantFactors());
diff --git a/lucene/core/src/java/org/apache/lucene/search/similarities/Similarity.java b/lucene/core/src/java/org/apache/lucene/search/similarities/Similarity.java
index 83582e44e25..721ce2a7311 100644
--- a/lucene/core/src/java/org/apache/lucene/search/similarities/Similarity.java
+++ b/lucene/core/src/java/org/apache/lucene/search/similarities/Similarity.java
@@ -17,6 +17,7 @@
package org.apache.lucene.search.similarities;
import java.util.Collections;
+import java.util.Objects;
import org.apache.lucene.analysis.tokenattributes.PositionIncrementAttribute;
import org.apache.lucene.document.NumericDocValuesField;
import org.apache.lucene.index.FieldInvertState;
@@ -208,6 +209,16 @@ public abstract class Similarity {
*/
public abstract float score(float freq, long norm);
+ /**
+ * Return a {@link BulkSimScorer} that produces the exact same scores as this {@link SimScorer}
+ * but is more efficient at bulk-computing scores.
+ *
+ * <p><b>NOTE</b>: The returned instance is not thread-safe.
+ */
+ public BulkSimScorer asBulkSimScorer() {
+ return new DefaultBulkSimScorer(this);
+ }
+
/**
* Explain the score for a single document
*
@@ -223,4 +234,38 @@ public abstract class Similarity {
Collections.singleton(freq));
}
}
+
+ /** Specialization of {@link SimScorer} for bulk-computation of scores. */
+ public interface BulkSimScorer {
+
+ /**
+ * Bulk computation of scores. For each index {@code i} in [0, size), scores[i] is computed as
+ * score(freqs[i], norms[i]). The default implementation does the following:
+ *
+ * <pre class="prettyprint">
+ * for (int i = 0; i < size; ++i) {
+ * scores[i] = score(freqs[i], norms[i]);
+ * }
+ * </pre>
+ *
+ * <p><b>NOTE</b>: It is legal to pass the same {@code freqs} and {@code scores} arrays.
+ */
+ void score(int size, float[] freqs, long[] norms, float[] scores);
+ }
+
+ private static class DefaultBulkSimScorer implements BulkSimScorer {
+
+ private final SimScorer scorer;
+
+ DefaultBulkSimScorer(SimScorer scorer) {
+ this.scorer = Objects.requireNonNull(scorer);
+ }
+
+ @Override
+ public void score(int size, float[] freqs, long[] norms, float[] scores) {
+ for (int i = 0; i < size; ++i) {
+ scores[i] = scorer.score(freqs[i], norms[i]);
+ }
+ }
+ }
}
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/search/similarities/AssertingSimilarity.java b/lucene/test-framework/src/java/org/apache/lucene/tests/search/similarities/AssertingSimilarity.java
index b4bf4af31ba..8943c8f7a89 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/search/similarities/AssertingSimilarity.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/search/similarities/AssertingSimilarity.java
@@ -99,6 +99,26 @@ public class AssertingSimilarity extends Similarity {
== delegate.score(freq.getValue().floatValue(), norm);
return explanation;
}
+
+ @Override
+ public BulkSimScorer asBulkSimScorer() {
+ BulkSimScorer bulkScorer = delegate.asBulkSimScorer();
+ return new BulkSimScorer() {
+ @Override
+ public void score(int size, float[] freqs, long[] norms, float[] scores) {
+ for (int i = 0; i < size; ++i) {
+ assert freqs[i] > 0;
+ assert norms[i] != 0;
+ }
+ bulkScorer.score(size, freqs, norms, scores);
+ for (int i = 0; i < size; ++i) {
+ float score = scores[i];
+ assert Float.isFinite(score);
+ assert score >= 0;
+ }
+ }
+ };
+ }
}
@Override
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/search/similarities/BaseSimilarityTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/tests/search/similarities/BaseSimilarityTestCase.java
index 763bf16e691..9a8e3802f90 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/search/similarities/BaseSimilarityTestCase.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/search/similarities/BaseSimilarityTestCase.java
@@ -27,12 +27,14 @@ import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.TermStatistics;
import org.apache.lucene.search.similarities.IndriDirichletSimilarity;
import org.apache.lucene.search.similarities.Similarity;
+import org.apache.lucene.search.similarities.Similarity.BulkSimScorer;
import org.apache.lucene.search.similarities.Similarity.SimScorer;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.index.RandomIndexWriter;
import org.apache.lucene.tests.search.CheckHits;
import org.apache.lucene.tests.util.LuceneTestCase;
import org.apache.lucene.tests.util.TestUtil;
+import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.SmallFloat;
@@ -521,4 +523,41 @@ public abstract class BaseSimilarityTestCase extends LuceneTestCase {
}
}
}
+
+ public void testBulkScore() throws IOException {
+ Random random = random();
+ Similarity similarity = getSimilarity(random);
+ CollectionStatistics corpus = newCorpus(random, 1);
+ TermStatistics term = newTerm(random, corpus);
+ SimScorer scorer = similarity.scorer(random().nextFloat(5f), corpus, term);
+ BulkSimScorer bulkScorer = scorer.asBulkSimScorer();
+ int freqUpperBound =
+ Math.toIntExact(Math.min(term.totalTermFreq() - term.docFreq() + 1, Integer.MAX_VALUE));
+
+ float[] freqs = new float[0];
+ long[] norms = new long[0];
+ float[] scores = new float[0];
+
+ int iters = atLeast(3);
+ for (int iter = 0; iter < iters; ++iter) {
+ int size = TestUtil.nextInt(random, 0, 200);
+ if (size > freqs.length) {
+ freqs = new float[ArrayUtil.oversize(size, Float.BYTES)];
+ norms = new long[freqs.length];
+ scores = new float[freqs.length];
+ }
+ for (int i = 0; i < size; ++i) {
+ freqs[i] = TestUtil.nextInt(random, 1, freqUpperBound);
+ norms[i] = TestUtil.nextLong(random, 1, 255);
+ }
+
+ float[] expectedScores = new float[size];
+ for (int i = 0; i < size; ++i) {
+ expectedScores[i] = scorer.score(freqs[i], norms[i]);
+ }
+ bulkScorer.score(size, freqs, norms, scores);
+
+ assertArrayEquals(expectedScores, ArrayUtil.copyOfSubArray(scores, 0, size), 0f);
+ }
+ }
}