[DiskBBQ] Compute only higher three bits for skipping

benchmarks/src/main/java/org/elasticsearch/benchmark/vector/OSQScorerBenchmark.java

@@ -170,6 +170,33 @@ public void scoreFromMemorySegmentOnlyVector(Blackhole bh) throws IOException {
         }
     }
 
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void scoreThreeUpperBitsFromMemorySegmentOnlyVector(Blackhole bh) throws IOException {
+        for (int j = 0; j < numQueries; j++) {
+            in.seek(0);
+            for (int i = 0; i < numVectors; i++) {
+                float qDist = scorer.quantizeScoreThreeUpperBit(binaryQueries[j]);
+                in.readFloats(corrections, 0, corrections.length);
+                int addition = Short.toUnsignedInt(in.readShort());
+                float score = scorer.score(
+                    result.lowerInterval(),
+                    result.upperInterval(),
+                    result.quantizedComponentSum(),
+                    result.additionalCorrection(),
+                    VectorSimilarityFunction.EUCLIDEAN,
+                    centroidDp,
+                    corrections[0],
+                    corrections[1],
+                    addition,
+                    corrections[2],
+                    qDist
+                );
+                bh.consume(score);
+            }
+        }
+    }
+
     @Benchmark
     @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
     public void scoreFromMemorySegmentOnlyVectorBulk(Blackhole bh) throws IOException {
@@ -199,6 +226,35 @@ public void scoreFromMemorySegmentOnlyVectorBulk(Blackhole bh) throws IOExceptio
         }
     }
 
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void scoreThreeUpperBitsFromMemorySegmentOnlyVectorBulk(Blackhole bh) throws IOException {
+        for (int j = 0; j < numQueries; j++) {
+            in.seek(0);
+            for (int i = 0; i < numVectors; i += 16) {
+                scorer.quantizeScoreThreeUpperBitBulk(binaryQueries[j], ES91OSQVectorsScorer.BULK_SIZE, scratchScores);
+                for (int k = 0; k < ES91OSQVectorsScorer.BULK_SIZE; k++) {
+                    in.readFloats(corrections, 0, corrections.length);
+                    int addition = Short.toUnsignedInt(in.readShort());
+                    float score = scorer.score(
+                        result.lowerInterval(),
+                        result.upperInterval(),
+                        result.quantizedComponentSum(),
+                        result.additionalCorrection(),
+                        VectorSimilarityFunction.EUCLIDEAN,
+                        centroidDp,
+                        corrections[0],
+                        corrections[1],
+                        addition,
+                        corrections[2],
+                        scratchScores[k]
+                    );
+                    bh.consume(score);
+                }
+            }
+        }
+    }
+
     @Benchmark
     @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
     public void scoreFromMemorySegmentAllBulk(Blackhole bh) throws IOException {

libs/simdvec/src/main/java/org/elasticsearch/simdvec/ES91OSQVectorsScorer.java

@@ -80,6 +80,54 @@ public long quantizeScore(byte[] q) throws IOException {
         return subRet0 + (subRet1 << 1) + (subRet2 << 2) + (subRet3 << 3);
     }
 
+    public long quantizeScoreThreeUpperBit(byte[] q) throws IOException {
+        assert q.length == length * 4;
+        final int size = length;
+        long subRet1 = 0;
+        long subRet2 = 0;
+        long subRet3 = 0;
+        int r = 0;
+        for (final int upperBound = size & -Long.BYTES; r < upperBound; r += Long.BYTES) {
+            final long value = in.readLong();
+            subRet1 += Long.bitCount((long) BitUtil.VH_LE_LONG.get(q, r + size) & value);
+            subRet2 += Long.bitCount((long) BitUtil.VH_LE_LONG.get(q, r + 2 * size) & value);
+            subRet3 += Long.bitCount((long) BitUtil.VH_LE_LONG.get(q, r + 3 * size) & value);
+        }
+        for (final int upperBound = size & -Integer.BYTES; r < upperBound; r += Integer.BYTES) {
+            final int value = in.readInt();
+            subRet1 += Integer.bitCount((int) BitUtil.VH_LE_INT.get(q, r + size) & value);
+            subRet2 += Integer.bitCount((int) BitUtil.VH_LE_INT.get(q, r + 2 * size) & value);
+            subRet3 += Integer.bitCount((int) BitUtil.VH_LE_INT.get(q, r + 3 * size) & value);
+        }
+        for (; r < size; r++) {
+            final byte value = in.readByte();
+            subRet1 += Integer.bitCount((q[r + size] & value) & 0xFF);
+            subRet2 += Integer.bitCount((q[r + 2 * size] & value) & 0xFF);
+            subRet3 += Integer.bitCount((q[r + 3 * size] & value) & 0xFF);
+        }
+        return (subRet1 << 1) + (subRet2 << 2) + (subRet3 << 3);
+    }
+
+    public long quantizeScoreLowerBit(byte[] q) throws IOException {
+        assert q.length == length * 4;
+        final int size = length;
+        long subRet = 0;
+        int r = 0;
+        for (final int upperBound = size & -Long.BYTES; r < upperBound; r += Long.BYTES) {
+            final long value = in.readLong();
+            subRet += Long.bitCount((long) BitUtil.VH_LE_LONG.get(q, r) & value);
+        }
+        for (final int upperBound = size & -Integer.BYTES; r < upperBound; r += Integer.BYTES) {
+            final int value = in.readInt();
+            subRet += Integer.bitCount((int) BitUtil.VH_LE_INT.get(q, r) & value);
+        }
+        for (; r < size; r++) {
+            final byte value = in.readByte();
+            subRet += Integer.bitCount((q[r] & value) & 0xFF);
+        }
+        return subRet;
+    }
+
     /**
      * compute the quantize distance between the provided quantized query and the quantized vectors
      * that are read from the wrapped {@link IndexInput}. The number of quantized vectors to read is
@@ -91,6 +139,12 @@ public void quantizeScoreBulk(byte[] q, int count, float[] scores) throws IOExce
         }
     }
 
+    public void quantizeScoreThreeUpperBitBulk(byte[] q, int count, float[] scores) throws IOException {
+        for (int i = 0; i < count; i++) {
+            scores[i] = quantizeScoreThreeUpperBit(q);
+        }
+    }
+
     /**
      * Computes the score by applying the necessary corrections to the provided quantized distance.
      */