diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index 7b44e8368bfd..e309d4e3c9c7 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -235,6 +235,9 @@ Optimizations * GITHUB#14980: Add bulk off-heap scoring for float32 vectors (Chris Hegarty) +* 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 278f735838e7..46cf17534410 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.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 TermScorer(PostingsEnum postingsEnum, SimScorer scorer, NumericDocValues impactsDisi = null; this.scorer = scorer; this.norms = norms; + this.bulkScorer = scorer.asBulkSimScorer(); } /** @@ -71,6 +74,7 @@ public TermScorer( } this.scorer = scorer; this.norms = norms; + this.bulkScorer = scorer.asBulkSimScorer(); } @Override @@ -165,10 +169,6 @@ public void nextDocsAndScores(int upTo, Bits liveDocs, DocAndFloatFeatureBuffer } } - 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 274bb475d6c7..78e89f515160 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 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 @@ private static class BM25Scorer extends SimScorer { 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 float score(float freq, long encodedNorm) { // 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 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 83582e44e25f..721ce2a73113 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 @@ protected SimScorer() {} */ 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. + * + *

NOTE: 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 Explanation explain(Explanation freq, long norm) { 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: + * + * 

+     * for (int i = 0; i < size; ++i) {
+     *   scores[i] = score(freqs[i], norms[i]);
+     * }
+     *
+ * + *

NOTE: 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 b4bf4af31ba5..8943c8f7a893 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 Explanation explain(Explanation freq, long norm) { == 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 763bf16e691f..9a8e3802f902 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.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 @@ private static void doTestScoring( } } } + + 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); + } + } }