optimize OptimizedScalarQuantizer#scalarQuantize when destination can be an integer array

Author: iverase

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

@@ -44,6 +44,7 @@ public class OptimizedScalarQuantizerBenchmark {
     float[] vector;
     float[] centroid;
     byte[] destination;
+    int[] intDestination;
 
     @Param({ "1", "4", "7" })
     byte bits;
@@ -55,6 +56,7 @@ public void init() {
         ThreadLocalRandom random = ThreadLocalRandom.current();
         // random byte arrays for binary methods
         destination = new byte[dims];
+        intDestination = new int[dims];
         vector = new float[dims];
         centroid = new float[dims];
         for (int i = 0; i < dims; ++i) {
@@ -75,4 +77,11 @@ public byte[] vector() {
         osq.scalarQuantize(vector, destination, bits, centroid);
         return destination;
     }
+
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public byte[] vectorToInt() {
+        osq.scalarQuantizeToInts(vector, intDestination, bits, centroid);
+        return destination;
+    }
 }

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

@@ -258,4 +258,25 @@ public static float soarDistance(float[] v1, float[] centroid, float[] originalR
         }
         return IMPL.soarDistance(v1, centroid, originalResidual, soarLambda, rnorm);
     }
+
+    /**
+     * Optimized-scalar quantization of the provided vector to the provided destination array.
+     *
+     * @param vector the vector to quantize
+     * @param destination the array to store the result
+     * @param lowInterval the minimum value, lower values in the original array will be replaced by this value
+     * @param upperInterval the maximum value, bigger values in the original array will be replaced by this value
+     * @param bit the number of bits to use for quantization, must be between 1 and 8
+     *
+     * @return return the sum of all the elements of the resulting quantized vector.
+     */
+    public static int quantizeVectorWithIntervals(float[] vector, int[] destination, float lowInterval, float upperInterval, byte bit) {
+        if (vector.length != destination.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + vector.length + "!=" + destination.length);
+        }
+        if (bit <= 0 || bit > Byte.SIZE) {
+            throw new IllegalArgumentException("bit must be between 1 and 8, but was: " + bit);
+        }
+        return IMPL.quantizeVectorWithIntervals(vector, destination, lowInterval, upperInterval, bit);
+    }
 }

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

@@ -269,4 +269,18 @@ public static float ipFloatByteImpl(float[] q, byte[] d) {
         }
         return ret;
     }
+
+    @Override
+    public int quantizeVectorWithIntervals(float[] vector, int[] destination, float lowInterval, float upperInterval, byte bits) {
+        float nSteps = ((1 << bits) - 1);
+        float step = (upperInterval - lowInterval) / nSteps;
+        int sumQuery = 0;
+        for (int h = 0; h < vector.length; h++) {
+            float xi = Math.min(Math.max(vector[h], lowInterval), upperInterval);
+            int assignment = Math.round((xi - lowInterval) / step);
+            sumQuery += assignment;
+            destination[h] = assignment;
+        }
+        return sumQuery;
+    }
 }

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

@@ -39,4 +39,6 @@ public interface ESVectorUtilSupport {
 
     float soarDistance(float[] v1, float[] centroid, float[] originalResidual, float soarLambda, float rnorm);
 
+    int quantizeVectorWithIntervals(float[] vector, int[] quantize, float lowInterval, float upperInterval, byte bit);
+
 }

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

@@ -791,4 +791,32 @@ public static float ipFloatByteImpl(float[] q, byte[] d) {
 
         return sum;
     }
+
+    @Override
+    public int quantizeVectorWithIntervals(float[] vector, int[] destination, float lowInterval, float upperInterval, byte bits) {
+        float nSteps = ((1 << bits) - 1);
+        float step = (upperInterval - lowInterval) / nSteps;
+        int sumQuery = 0;
+        int i = 0;
+        if (vector.length > 2 * FLOAT_SPECIES.length()) {
+            int limit = FLOAT_SPECIES.loopBound(vector.length);
+            FloatVector lowVec = FloatVector.broadcast(FLOAT_SPECIES, lowInterval);
+            FloatVector upperVec = FloatVector.broadcast(FLOAT_SPECIES, upperInterval);
+            FloatVector stepVec = FloatVector.broadcast(FLOAT_SPECIES, step);
+            for (; i < limit; i += FLOAT_SPECIES.length()) {
+                FloatVector v = FloatVector.fromArray(FLOAT_SPECIES, vector, i);
+                FloatVector xi = v.max(lowVec).min(upperVec); // clamp
+                IntVector assignment = xi.sub(lowVec).div(stepVec).add(0.5f).convert(VectorOperators.F2I, 0).reinterpretAsInts(); // round
+                sumQuery += assignment.reduceLanes(ADD);
+                assignment.intoArray(destination, i);
+            }
+        }
+        for (; i < vector.length; i++) {
+            float xi = Math.min(Math.max(vector[i], lowInterval), upperInterval);
+            int assignment = Math.round((xi - lowInterval) / step);
+            sumQuery += assignment;
+            destination[i] = assignment;
+        }
+        return sumQuery;
+    }
 }

libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java

@@ -286,6 +286,29 @@ public void testSoarDistance() {
         assertEquals(expected, result, deltaEps);
     }
 
+    public void testQuantizeVectorWithIntervals() {
+        int vectorSize = randomIntBetween(1, 2048);
+        float[] vector = new float[vectorSize];
+
+        byte bits = (byte) randomIntBetween(1, 8);
+        for (int i = 0; i < vectorSize; ++i) {
+            vector[i] = random().nextFloat();
+        }
+        float low = random().nextFloat();
+        float high = random().nextFloat();
+        if (low > high) {
+            float tmp = low;
+            low = high;
+            high = tmp;
+        }
+        byte[] quantizeExpected = new byte[vectorSize];
+        byte[] quantizeResult = new byte[vectorSize];
+        var expected = defaultedProvider.getVectorUtilSupport().quantizeVectorWithIntervals(vector, low, high, quantizeExpected, bits);
+        var result = defOrPanamaProvider.getVectorUtilSupport().quantizeVectorWithIntervals(vector, low, high, quantizeResult, bits);
+        assertArrayEquals(quantizeExpected, quantizeResult);
+        assertEquals(expected, result, 0f);
+    }
+
     void testIpByteBinImpl(ToLongBiFunction<byte[], byte[]> ipByteBinFunc) {
         int iterations = atLeast(50);
         for (int i = 0; i < iterations; i++) {

server/src/main/java/org/elasticsearch/index/codec/vectors/BQSpaceUtils.java

@@ -57,4 +57,34 @@ public static void transposeHalfByte(byte[] q, byte[] quantQueryByte) {
             quantQueryByte[index + 3 * quantQueryByte.length / 4] = (byte) upperByte;
         }
     }
+
+    /**
+     * Same as {@link #transposeHalfByte(byte[], byte[])} but the input vector is provided as
+     * an array of integers.This is useful when using {@link OptimizedScalarQuantizer#scalarQuantizeToInts(float[], int[], byte, float[])}
+     * which performs better than the byte version.
+     *
+     * @param q the query vector, assumed to be half-byte quantized with values between 0 and 15
+     * @param quantQueryByte the byte array to store the transposed query vector
+     * */
+    public static void transposeHalfByte(int[] q, byte[] quantQueryByte) {
+        for (int i = 0; i < q.length;) {
+            assert q[i] >= 0 && q[i] <= 15;
+            int lowerByte = 0;
+            int lowerMiddleByte = 0;
+            int upperMiddleByte = 0;
+            int upperByte = 0;
+            for (int j = 7; j >= 0 && i < q.length; j--) {
+                lowerByte |= (q[i] & 1) << j;
+                lowerMiddleByte |= ((q[i] >> 1) & 1) << j;
+                upperMiddleByte |= ((q[i] >> 2) & 1) << j;
+                upperByte |= ((q[i] >> 3) & 1) << j;
+                i++;
+            }
+            int index = ((i + 7) / 8) - 1;
+            quantQueryByte[index] = (byte) lowerByte;
+            quantQueryByte[index + quantQueryByte.length / 4] = (byte) lowerMiddleByte;
+            quantQueryByte[index + quantQueryByte.length / 2] = (byte) upperMiddleByte;
+            quantQueryByte[index + 3 * quantQueryByte.length / 4] = (byte) upperByte;
+        }
+    }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java

@@ -54,6 +54,20 @@ public static void packAsBinary(byte[] vector, byte[] packed) {
         }
     }
 
+    public static void packAsBinary(int[] vector, byte[] packed) {
+        for (int i = 0; i < vector.length;) {
+            byte result = 0;
+            for (int j = 7; j >= 0 && i < vector.length; j--) {
+                assert vector[i] == 0 || vector[i] == 1;
+                result |= (byte) ((vector[i] & 1) << j);
+                ++i;
+            }
+            int index = ((i + 7) / 8) - 1;
+            assert index < packed.length;
+            packed[index] = result;
+        }
+    }
+
     public static int discretize(int value, int bucket) {
         return ((value + (bucket - 1)) / bucket) * bucket;
     }

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

@@ -182,7 +182,7 @@ private static class MemorySegmentPostingsVisitor implements PostingVisitor {
         DocIdsWriter docIdsWriter = new DocIdsWriter();
 
         final float[] scratch;
-        final byte[] quantizationScratch;
+        final int[] quantizationScratch;
         final byte[] quantizedQueryScratch;
         final OptimizedScalarQuantizer quantizer;
         final float[] correctiveValues = new float[3];
@@ -202,7 +202,7 @@ private static class MemorySegmentPostingsVisitor implements PostingVisitor {
             this.needsScoring = needsScoring;
 
             scratch = new float[target.length];
-            quantizationScratch = new byte[target.length];
+            quantizationScratch = new int[target.length];
             final int discretizedDimensions = discretize(fieldInfo.getVectorDimension(), 64);
             quantizedQueryScratch = new byte[QUERY_BITS * discretizedDimensions / 8];
             quantizedByteLength = discretizedDimensions / 8 + (Float.BYTES * 3) + Short.BYTES;
@@ -344,7 +344,7 @@ private void quantizeQueryIfNecessary() {
                 if (fieldInfo.getVectorSimilarityFunction() == COSINE) {
                     VectorUtil.l2normalize(scratch);
                 }
-                queryCorrections = quantizer.scalarQuantize(scratch, quantizationScratch, (byte) 4, centroid);
+                queryCorrections = quantizer.scalarQuantizeToInts(scratch, quantizationScratch, (byte) 4, centroid);
                 transposeHalfByte(quantizationScratch, quantizedQueryScratch);
                 quantized = true;
             }

server/src/main/java/org/elasticsearch/index/codec/vectors/DefaultIVFVectorsWriter.java

@@ -291,7 +291,7 @@ CentroidAssignments calculateAndWriteCentroids(
     static class BinarizedFloatVectorValues {
         private OptimizedScalarQuantizer.QuantizationResult corrections;
         private final byte[] binarized;
-        private final byte[] initQuantized;
+        private final int[] initQuantized;
         private float[] centroid;
         private final FloatVectorValues values;
         private final OptimizedScalarQuantizer quantizer;
@@ -302,7 +302,7 @@ static class BinarizedFloatVectorValues {
             this.values = delegate;
             this.quantizer = quantizer;
             this.binarized = new byte[discretize(delegate.dimension(), 64) / 8];
-            this.initQuantized = new byte[delegate.dimension()];
+            this.initQuantized = new int[delegate.dimension()];
         }
 
         public OptimizedScalarQuantizer.QuantizationResult getCorrectiveTerms(int ord) {
@@ -323,7 +323,7 @@ public byte[] vectorValue(int ord) throws IOException {
         }
 
         private void binarize(int ord) throws IOException {
-            corrections = quantizer.scalarQuantize(values.vectorValue(ord), initQuantized, INDEX_BITS, centroid);
+            corrections = quantizer.scalarQuantizeToInts(values.vectorValue(ord), initQuantized, INDEX_BITS, centroid);
             packAsBinary(initQuantized, binarized);
         }
     }