Introduce an int4 off-heap vector scorer (#129824)

* Introduce an int4 off-heap vector scorer

* iter

* Update server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsReader.java

Co-authored-by: Benjamin Trent <[email protected]>

---------

Co-authored-by: Benjamin Trent <[email protected]>

Author: iverase

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

@@ -0,0 +1,123 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+package org.elasticsearch.benchmark.vector;
+
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.store.MMapDirectory;
+import org.apache.lucene.util.VectorUtil;
+import org.elasticsearch.common.logging.LogConfigurator;
+import org.elasticsearch.core.IOUtils;
+import org.elasticsearch.simdvec.ES91Int4VectorsScorer;
+import org.elasticsearch.simdvec.internal.vectorization.ESVectorizationProvider;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@State(Scope.Benchmark)
+// first iteration is complete garbage, so make sure we really warmup
+@Warmup(iterations = 4, time = 1)
+// real iterations. not useful to spend tons of time here, better to fork more
+@Measurement(iterations = 5, time = 1)
+// engage some noise reduction
+@Fork(value = 1)
+public class Int4ScorerBenchmark {
+
+    static {
+        LogConfigurator.configureESLogging(); // native access requires logging to be initialized
+    }
+
+    @Param({ "384", "702", "1024" })
+    int dims;
+
+    int numVectors = 200;
+    int numQueries = 10;
+
+    byte[] scratch;
+    byte[][] binaryVectors;
+    byte[][] binaryQueries;
+
+    ES91Int4VectorsScorer scorer;
+    Directory dir;
+    IndexInput in;
+
+    @Setup
+    public void setup() throws IOException {
+        binaryVectors = new byte[numVectors][dims];
+        dir = new MMapDirectory(Files.createTempDirectory("vectorData"));
+        try (IndexOutput out = dir.createOutput("vectors", IOContext.DEFAULT)) {
+            for (byte[] binaryVector : binaryVectors) {
+                for (int i = 0; i < dims; i++) {
+                    // 4-bit quantization
+                    binaryVector[i] = (byte) ThreadLocalRandom.current().nextInt(16);
+                }
+                out.writeBytes(binaryVector, 0, binaryVector.length);
+            }
+        }
+
+        in = dir.openInput("vectors", IOContext.DEFAULT);
+        binaryQueries = new byte[numVectors][dims];
+        for (byte[] binaryVector : binaryVectors) {
+            for (int i = 0; i < dims; i++) {
+                // 4-bit quantization
+                binaryVector[i] = (byte) ThreadLocalRandom.current().nextInt(16);
+            }
+        }
+
+        scratch = new byte[dims];
+        scorer = ESVectorizationProvider.getInstance().newES91Int4VectorsScorer(in, dims);
+    }
+
+    @TearDown
+    public void teardown() throws IOException {
+        IOUtils.close(dir, in);
+    }
+
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void scoreFromArray(Blackhole bh) throws IOException {
+        for (int j = 0; j < numQueries; j++) {
+            in.seek(0);
+            for (int i = 0; i < numVectors; i++) {
+                in.readBytes(scratch, 0, dims);
+                bh.consume(VectorUtil.int4DotProduct(binaryQueries[j], scratch));
+            }
+        }
+    }
+
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void scoreFromMemorySegmentOnlyVector(Blackhole bh) throws IOException {
+        for (int j = 0; j < numQueries; j++) {
+            in.seek(0);
+            for (int i = 0; i < numVectors; i++) {
+                bh.consume(scorer.int4DotProduct(binaryQueries[j]));
+            }
+        }
+    }
+}

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

@@ -0,0 +1,43 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+package org.elasticsearch.simdvec;
+
+import org.apache.lucene.store.IndexInput;
+
+import java.io.IOException;
+
+/** Scorer for quantized vectors stored as an {@link IndexInput}.
+ * <p>
+ * Similar to {@link org.apache.lucene.util.VectorUtil#int4DotProduct(byte[], byte[])} but
+ * one value is read directly from an {@link IndexInput}.
+ *
+ * */
+public class ES91Int4VectorsScorer {
+
+    /** The wrapper {@link IndexInput}. */
+    protected final IndexInput in;
+    protected final int dimensions;
+    protected byte[] scratch;
+
+    /** Sole constructor, called by sub-classes. */
+    public ES91Int4VectorsScorer(IndexInput in, int dimensions) {
+        this.in = in;
+        this.dimensions = dimensions;
+        scratch = new byte[dimensions];
+    }
+
+    public long int4DotProduct(byte[] b) throws IOException {
+        in.readBytes(scratch, 0, dimensions);
+        int total = 0;
+        for (int i = 0; i < dimensions; i++) {
+            total += scratch[i] * b[i];
+        }
+        return total;
+    }
+}

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

@@ -47,6 +47,10 @@ public static ES91OSQVectorsScorer getES91OSQVectorsScorer(IndexInput input, int
         return ESVectorizationProvider.getInstance().newES91OSQVectorsScorer(input, dimension);
     }
 
+    public static ES91Int4VectorsScorer getES91Int4VectorsScorer(IndexInput input, int dimension) throws IOException {
+        return ESVectorizationProvider.getInstance().newES91Int4VectorsScorer(input, dimension);
+    }
+
     public static long ipByteBinByte(byte[] q, byte[] d) {
         if (q.length != d.length * B_QUERY) {
             throw new IllegalArgumentException("vector dimensions incompatible: " + q.length + "!= " + B_QUERY + " x " + d.length);

libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorizationProvider.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.simdvec.internal.vectorization;
 
 import org.apache.lucene.store.IndexInput;
+import org.elasticsearch.simdvec.ES91Int4VectorsScorer;
 import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
 
 import java.io.IOException;
@@ -30,4 +31,9 @@ public ESVectorUtilSupport getVectorUtilSupport() {
     public ES91OSQVectorsScorer newES91OSQVectorsScorer(IndexInput input, int dimension) throws IOException {
         return new ES91OSQVectorsScorer(input, dimension);
     }
+
+    @Override
+    public ES91Int4VectorsScorer newES91Int4VectorsScorer(IndexInput input, int dimension) throws IOException {
+        return new ES91Int4VectorsScorer(input, dimension);
+    }
 }

libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorizationProvider.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.simdvec.internal.vectorization;
 
 import org.apache.lucene.store.IndexInput;
+import org.elasticsearch.simdvec.ES91Int4VectorsScorer;
 import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
 
 import java.io.IOException;
@@ -31,6 +32,9 @@ public static ESVectorizationProvider getInstance() {
     /** Create a new {@link ES91OSQVectorsScorer} for the given {@link IndexInput}. */
     public abstract ES91OSQVectorsScorer newES91OSQVectorsScorer(IndexInput input, int dimension) throws IOException;
 
+    /** Create a new {@link ES91Int4VectorsScorer} for the given {@link IndexInput}. */
+    public abstract ES91Int4VectorsScorer newES91Int4VectorsScorer(IndexInput input, int dimension) throws IOException;
+
     // visible for tests
     static ESVectorizationProvider lookup(boolean testMode) {
         return new DefaultESVectorizationProvider();

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorizationProvider.java

@@ -13,6 +13,7 @@
 import org.apache.lucene.util.Constants;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
+import org.elasticsearch.simdvec.ES91Int4VectorsScorer;
 import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
 
 import java.io.IOException;
@@ -38,6 +39,9 @@ public static ESVectorizationProvider getInstance() {
     /** Create a new {@link ES91OSQVectorsScorer} for the given {@link IndexInput}. */
     public abstract ES91OSQVectorsScorer newES91OSQVectorsScorer(IndexInput input, int dimension) throws IOException;
 
+    /** Create a new {@link ES91Int4VectorsScorer} for the given {@link IndexInput}. */
+    public abstract ES91Int4VectorsScorer newES91Int4VectorsScorer(IndexInput input, int dimension) throws IOException;
+
     // visible for tests
     static ESVectorizationProvider lookup(boolean testMode) {
         final int runtimeVersion = Runtime.version().feature();