vector methods

Author: davidkyle

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

@@ -435,4 +435,15 @@ public static int indexOf(byte[] bytes, int offset, int length, byte marker) {
         Objects.checkFromIndexSize(offset, length, bytes.length);
         return IMPL.indexOf(bytes, offset, length, marker);
     }
+
+    public static void vectorAccumulateAdd(float[] a, float[] b) {
+        if (a.length != b.length) {
+            throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
+        }
+        IMPL.vectorAccumulateAdd(a, b);
+    }
+
+    public static void vectorScalerDivide(float [] a, float b) {
+        IMPL.vectorScalerDivide(a, b);
+    }
 }

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

@@ -13,6 +13,9 @@
 import org.apache.lucene.util.Constants;
 import org.apache.lucene.util.VectorUtil;
 
+import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicLong;
+
 final class DefaultESVectorUtilSupport implements ESVectorUtilSupport {
 
     private static float fma(float a, float b, float c) {
@@ -455,4 +458,16 @@ public static void transposeHalfByteImpl(int[] q, byte[] quantQueryByte) {
     public int indexOf(byte[] bytes, int offset, int length, byte marker) {
         return ByteArrayUtils.indexOf(bytes, offset, length, marker);
     }
+
+    public void vectorScalerDivide(float [] a, float b) {
+        for (int d = 0; d < a.length; d++) {
+            a[d] /= b;
+        }
+    }
+
+    public void vectorAccumulateAdd(float[] a, float[] b) {
+        for (int d = 0; d < a.length; d++) {
+            a[d] += b[d];
+        }
+    }
 }

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

@@ -71,4 +71,8 @@ void soarDistanceBulk(
     void transposeHalfByte(int[] q, byte[] quantQueryByte);
 
     int indexOf(byte[] bytes, int offset, int length, byte marker);
+
+    void vectorScalerDivide(float [] a, float b);
+
+    void vectorAccumulateAdd(float[] a, float[] b);
 }

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

@@ -1140,4 +1140,37 @@ public int indexOf(final byte[] bytes, final int offset, final int length, final
         }
         return -1;
     }
+
+    @Override
+    public void vectorAccumulateAdd(float[] a, float[] b) {
+        final VectorSpecies<Float> SPECIES = FloatVector.SPECIES_PREFERRED;
+
+        int i = 0;
+        for (; i < SPECIES.loopBound(a.length); i += SPECIES.length()) {
+            FloatVector va = FloatVector.fromArray(SPECIES, a, i);
+            FloatVector vb = FloatVector.fromArray(SPECIES, b, i);
+            FloatVector vc = va.add(vb);
+            vc.intoArray(a, i);
+        }
+
+        for (; i < a.length; i++) {
+            a[i] += b[i];
+        }
+    }
+
+    @Override
+    public void vectorScalerDivide(float[] a, float b) {
+        final VectorSpecies<Float> SPECIES = FloatVector.SPECIES_PREFERRED;
+
+        int i = 0;
+        for (; i < SPECIES.loopBound(a.length); i += SPECIES.length()) {
+            FloatVector va = FloatVector.fromArray(SPECIES, a, i);
+            FloatVector vc = va.div(b);
+            vc.intoArray(a, i);
+        }
+
+        for (; i < a.length; i++) {
+            a[i] = a[i] / b;
+        }
+    }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/cluster/KMeansLocal.java

@@ -106,25 +106,19 @@ private static boolean stepLloyd(
             for (int idx = 0; idx < vectors.size(); idx++) {
                 final int assignment = assignments[translateOrd.apply(idx)];
                 if (centroidChanged.get(assignment)) {
-                    float[] centroid = centroids[assignment];
                     if (centroidCounts[assignment]++ == 0) {
-                        Arrays.fill(centroid, 0.0f);
-                    }
-                    float[] vector = vectors.vectorValue(idx);
-                    for (int d = 0; d < dim; d++) {
-                        centroid[d] += vector[d];
+                        centroids[assignment] = vectors.vectorValue(idx);
+                        continue;
                     }
+                    ESVectorUtil.vectorAccumulateAdd(centroids[assignment], vectors.vectorValue(idx));
                 }
             }
 
             for (int clusterIdx = 0; clusterIdx < centroids.length; clusterIdx++) {
                 if (centroidChanged.get(clusterIdx)) {
                     float count = (float) centroidCounts[clusterIdx];
                     if (count > 0) {
-                        float[] centroid = centroids[clusterIdx];
-                        for (int d = 0; d < dim; d++) {
-                            centroid[d] /= count;
-                        }
+                        ESVectorUtil.vectorScalerDivide(centroids[clusterIdx], count);
                     }
                 }
             }