Adds accelerates optimized scalar quantization with vectorized functions

Author: benwtrent

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

@@ -144,4 +144,71 @@ static int andBitCountLong(byte[] a, byte[] b) {
         }
         return distance;
     }
+
+    /**
+     * Calculate the loss for optimized-scalar quantization for the given parameteres
+     * @param target The vector being quantized, assumed to be centered
+     * @param interval The interval for which to calculate the loss
+     * @param points the quantization points
+     * @param norm2 The norm squared of the target vector
+     * @param lambda The lambda parameter for controlling anisotropic loss calculation
+     * @return The loss for the given parameters
+     */
+    public static float calculateOSQLoss(float[] target, float[] interval, int points, float norm2, float lambda) {
+        assert interval.length == 2;
+        float step = ((interval[1] - interval[0]) / (points - 1.0F));
+        float invStep = 1f / step;
+        return IMPL.calculateOSQLoss(target, interval, step, invStep, norm2, lambda);
+    }
+
+    /**
+     * Calculate the grid points for optimized-scalar quantization
+     * @param target The vector being quantized, assumed to be centered
+     * @param interval The interval for which to calculate the grid points
+     * @param points the quantization points
+     * @param pts The array to store the grid points, must be of length 5
+     */
+    public static void calculateOSQGridPoints(float[] target, float[] interval, int points, float[] pts) {
+        assert interval.length == 2;
+        assert pts.length == 5;
+        float invStep = (points - 1.0F) / (interval[1] - interval[0]);
+        IMPL.calculateOSQGridPoints(target, interval, points, invStep, pts);
+    }
+
+    /**
+     * Center the target vector and calculate the optimized-scalar quantization statistics
+     * @param target The vector being quantized
+     * @param centroid The centroid of the target vector
+     * @param centered The destination of the centered vector, will be overwritten
+     * @param stats The array to store the statistics, must be of length 5
+     */
+    public static void centerAndCalculateOSQStatsEuclidean(float[] target, float[] centroid, float[] centered, float[] stats) {
+        assert target.length == centroid.length;
+        assert stats.length == 5;
+        if (target.length != centroid.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + target.length + "!=" + centroid.length);
+        }
+        if (centered.length != target.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + centered.length + "!=" + target.length);
+        }
+        IMPL.centerAndCalculateOSQStatsEuclidean(target, centroid, centered, stats);
+    }
+
+    /**
+     * Center the target vector and calculate the optimized-scalar quantization statistics
+     * @param target The vector being quantized
+     * @param centroid The centroid of the target vector
+     * @param centered The destination of the centered vector, will be overwritten
+     * @param stats The array to store the statistics, must be of length 6
+     */
+    public static void centerAndCalculateOSQStatsDp(float[] target, float[] centroid, float[] centered, float[] stats) {
+        if (target.length != centroid.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + target.length + "!=" + centroid.length);
+        }
+        if (centered.length != target.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + centered.length + "!=" + target.length);
+        }
+        assert stats.length == 6;
+        IMPL.centerAndCalculateOSQStatsDp(target, centroid, centered, stats);
+    }
 }

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

@@ -44,6 +44,100 @@ public float ipFloatByte(float[] q, byte[] d) {
         return ipFloatByteImpl(q, d);
     }
 
+    @Override
+    public float calculateOSQLoss(float[] target, float[] interval, float step, float invStep, float norm2, float lambda) {
+        float a = interval[0];
+        float b = interval[1];
+        float xe = 0f;
+        float e = 0f;
+        for (float xi : target) {
+            // this is quantizing and then dequantizing the vector
+            float xiq = fma(step, Math.round((Math.min(Math.max(xi, a), b) - a) * invStep), a);
+            // how much does the de-quantized value differ from the original value
+            float xiiq = xi - xiq;
+            e = fma(xiiq, xiiq, e);
+            xe = fma(xi, xiiq, xe);
+        }
+        return (1f - lambda) * xe * xe / norm2 + lambda * e;
+    }
+
+    @Override
+    public void calculateOSQGridPoints(float[] target, float[] interval, int points, float invStep, float[] pts) {
+        float a = interval[0];
+        float b = interval[1];
+        float daa = 0;
+        float dab = 0;
+        float dbb = 0;
+        float dax = 0;
+        float dbx = 0;
+        for (float v : target) {
+            float k = Math.round((Math.min(Math.max(v, a), b) - a) * invStep);
+            float s = k / (points - 1);
+            float ms = 1f - s;
+            daa = fma(ms, ms, daa);
+            dab = fma(ms, s, dab);
+            dbb = fma(s, s, dbb);
+            dax = fma(ms, v, dax);
+            dbx = fma(s, v, dbx);
+        }
+        pts[0] = daa;
+        pts[1] = dab;
+        pts[2] = dbb;
+        pts[3] = dax;
+        pts[4] = dbx;
+    }
+
+    @Override
+    public void centerAndCalculateOSQStatsEuclidean(float[] target, float[] centroid, float[] centered, float[] stats) {
+        float vecMean = 0;
+        float vecVar = 0;
+        float norm2 = 0;
+        float min = Float.MAX_VALUE;
+        float max = -Float.MAX_VALUE;
+        for (int i = 0; i < target.length; i++) {
+            centered[i] = target[i] - centroid[i];
+            min = Math.min(min, centered[i]);
+            max = Math.max(max, centered[i]);
+            norm2 = fma(centered[i], centered[i], norm2);
+            float delta = centered[i] - vecMean;
+            vecMean += delta / (i + 1);
+            float delta2 = centered[i] - vecMean;
+            vecVar = fma(delta, delta2, vecVar);
+        }
+        stats[0] = vecMean;
+        stats[1] = vecVar / target.length;
+        stats[2] = norm2;
+        stats[3] = min;
+        stats[4] = max;
+    }
+
+    @Override
+    public void centerAndCalculateOSQStatsDp(float[] target, float[] centroid, float[] centered, float[] stats) {
+        float vecMean = 0;
+        float vecVar = 0;
+        float norm2 = 0;
+        float centroidDot = 0;
+        float min = Float.MAX_VALUE;
+        float max = -Float.MAX_VALUE;
+        for (int i = 0; i < target.length; i++) {
+            centroidDot = fma(target[i], centroid[i], centroidDot);
+            centered[i] = target[i] - centroid[i];
+            min = Math.min(min, centered[i]);
+            max = Math.max(max, centered[i]);
+            norm2 = fma(centered[i], centered[i], norm2);
+            float delta = centered[i] - vecMean;
+            vecMean += delta / (i + 1);
+            float delta2 = centered[i] - vecMean;
+            vecVar = fma(delta, delta2, vecVar);
+        }
+        stats[0] = vecMean;
+        stats[1] = vecVar / target.length;
+        stats[2] = norm2;
+        stats[3] = min;
+        stats[4] = max;
+        stats[5] = centroidDot;
+    }
+
     public static int ipByteBitImpl(byte[] q, byte[] d) {
         return ipByteBitImpl(q, d, 0);
     }

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

@@ -20,4 +20,12 @@ public interface ESVectorUtilSupport {
     float ipFloatBit(float[] q, byte[] d);
 
     float ipFloatByte(float[] q, byte[] d);
+
+    float calculateOSQLoss(float[] target, float[] interval, float step, float invStep, float norm2, float lambda);
+
+    void calculateOSQGridPoints(float[] target, float[] interval, int points, float invStep, float[] pts);
+
+    void centerAndCalculateOSQStatsEuclidean(float[] target, float[] centroid, float[] centered, float[] stats);
+
+    void centerAndCalculateOSQStatsDp(float[] target, float[] centroid, float[] centered, float[] stats);
 }