Acceleration of initial centroid selection in kmeans++

Use Java Vector API to accelerate chooseInitialCentroids in KMeansPlusPlusClusterer.

Author: AbhijitKulkarni1

jvector-base/src/main/java/io/github/jbellis/jvector/quantization/KMeansPlusPlusClusterer.java

@@ -181,26 +181,26 @@ private static VectorFloat<?> chooseInitialCentroids(VectorFloat<?>[] points, in
 
         float[] distances = new float[points.length];
         Arrays.fill(distances, Float.MAX_VALUE);
+        VectorFloat<?> distancesVector = vectorTypeSupport.createFloatVector(distances);
+        int distancesLength = points.length;
 
         // Choose the first centroid randomly
         VectorFloat<?> firstCentroid = points[random.nextInt(points.length)];
         centroids.copyFrom(firstCentroid, 0, 0, firstCentroid.length());
+        VectorFloat<?> newDistancesVector = vectorTypeSupport.createFloatVector(points.length);
         for (int i = 0; i < points.length; i++) {
-            float distance1 = squareL2Distance(points[i], firstCentroid);
-            distances[i] = Math.min(distances[i], distance1);
+            newDistancesVector.set(i, squareL2Distance(points[i], firstCentroid));
         }
+        VectorUtil.minInPlace(distancesVector, newDistancesVector);
 
         // For each subsequent centroid
         for (int i = 1; i < k; i++) {
-            float totalDistance = 0;
-            for (float distance : distances) {
-                totalDistance += distance;
-            }
+            float totalDistance = VectorUtil.sum(distancesVector);
 
             float r = random.nextFloat() * totalDistance;
             int selectedIdx = -1;
-            for (int j = 0; j < distances.length; j++) {
-                r -= distances[j];
+            for (int j = 0; j < distancesLength; j++) {
+                r -= distancesVector.get(j);
                 if (r < 1e-6) {
                     selectedIdx = j;
                     break;
@@ -215,10 +215,11 @@ private static VectorFloat<?> chooseInitialCentroids(VectorFloat<?>[] points, in
             centroids.copyFrom(nextCentroid, 0, i * nextCentroid.length(), nextCentroid.length());
 
             // Update distances, but only if the new centroid provides a closer distance
-            for (int j = 0; j < points.length; j++) {
-                float newDistance = squareL2Distance(points[j], nextCentroid);
-                distances[j] = Math.min(distances[j], newDistance);
+            // All entries of newDistancesVector is overwritten with the updated squareL2Distance value
+            for (int j = 0; j < distancesLength; j++) {
+                newDistancesVector.set(j, squareL2Distance(points[j], nextCentroid));
             }
+            VectorUtil.minInPlace(distancesVector, newDistancesVector);
         }
         assertFinite(centroids);
         return centroids;

jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java

@@ -287,6 +287,13 @@ public VectorFloat<?> sub(VectorFloat<?> a, int aOffset, VectorFloat<?> b, int b
     return result;
   }
 
+  @Override
+  public void minInPlace(VectorFloat<?> v1, VectorFloat<?> v2) {
+    for (int i = 0; i < v1.length(); i++) {
+      v1.set(i, Math.min(v1.get(i), v2.get(i)));
+    }
+  }
+
   @Override
   public float assembleAndSum(VectorFloat<?> data, int dataBase, ByteSequence<?> baseOffsets) {
     float sum = 0f;
@@ -541,4 +548,5 @@ public float nvqUniformLoss(VectorFloat<?> vector, float minValue, float maxValu
 
     return squaredSum;
   }
+
 }

jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtil.java

@@ -158,6 +158,10 @@ public static VectorFloat<?> sub(VectorFloat<?> a, int aOffset, VectorFloat<?> b
     return impl.sub(a, aOffset, b, bOffset, length);
   }
 
+  public static void minInPlace(VectorFloat<?> distances1, VectorFloat<?> distances2) {
+    impl.minInPlace(distances1, distances2);
+  }
+
   public static float assembleAndSum(VectorFloat<?> data, int dataBase, ByteSequence<?> dataOffsets) {
     return impl.assembleAndSum(data, dataBase, dataOffsets);
   }
@@ -238,4 +242,5 @@ public static float nvqLoss(VectorFloat<?> vector, float growthRate, float midpo
   public static float nvqUniformLoss(VectorFloat<?> vector, float minValue, float maxValue, int nBits) {
     return impl.nvqUniformLoss(vector, minValue, maxValue, nBits);
   }
+
 }

jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java

@@ -82,6 +82,9 @@ public interface VectorUtilSupport {
   /** @return a - b, element-wise, starting at aOffset and bOffset respectively */
   VectorFloat<?> sub(VectorFloat<?> a, int aOffset, VectorFloat<?> b, int bOffset, int length);
 
+  /** Calculates the minimum value for every corresponding lane values in v1 and v2, in place (v1 will be modified) */
+  void minInPlace(VectorFloat<?> v1, VectorFloat<?> v2);
+
   /**
    * Calculates the sum of sparse points in a vector.
    * <p>
@@ -300,4 +303,5 @@ default float pqDecodedCosineSimilarity(ByteSequence<?> encoded, int clusterCoun
    * @param nBits the number of bits per dimension
    */
   float nvqUniformLoss(VectorFloat<?> vector, float minValue, float maxValue, int nBits);
+
 }

jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java

@@ -110,6 +110,11 @@ public VectorFloat<?> sub(VectorFloat<?> a, int aOffset, VectorFloat<?> b, int b
         return VectorSimdOps.sub((MemorySegmentVectorFloat) a, aOffset, (MemorySegmentVectorFloat) b, bOffset, length);
     }
 
+    @Override
+    public void minInPlace(VectorFloat<?> v1, VectorFloat<?> v2) {
+        VectorSimdOps.minInPlace((MemorySegmentVectorFloat) v1, (MemorySegmentVectorFloat) v2);
+    }
+
     @Override
     public float assembleAndSum(VectorFloat<?> data, int dataBase, ByteSequence<?> baseOffsets) {
         assert baseOffsets.offset() == 0 : "Base offsets are expected to have an offset of 0. Found: " + baseOffsets.offset();
@@ -225,4 +230,5 @@ public float nvqLoss(VectorFloat<?> vector, float growthRate, float midpoint, fl
     public float nvqUniformLoss(VectorFloat<?> vector, float minValue, float maxValue, int nBits) {
         return VectorSimdOps.nvqUniformLoss((MemorySegmentVectorFloat) vector, minValue, maxValue, nBits);
     }
+
 }

jvector-native/src/main/java/io/github/jbellis/jvector/vector/VectorSimdOps.java

@@ -598,6 +598,26 @@ public static int hammingDistance(long[] a, long[] b) {
         return res;
     }
 
+    static void minInPlace(MemorySegmentVectorFloat v1, MemorySegmentVectorFloat v2) {
+        if (v1.length() != v2.length()) {
+            throw new IllegalArgumentException("Vectors must have the same length");
+        }
+
+        int vectorizedLength = FloatVector.SPECIES_PREFERRED.loopBound(v1.length());
+
+        // Process the vectorized part
+        for (int i = 0; i < vectorizedLength; i += FloatVector.SPECIES_PREFERRED.length()) {
+            var a = FloatVector.fromMemorySegment(FloatVector.SPECIES_PREFERRED, v1.get(), v1.offset(i), ByteOrder.LITTLE_ENDIAN);
+            var b = FloatVector.fromMemorySegment(FloatVector.SPECIES_PREFERRED, v2.get(), v2.offset(i), ByteOrder.LITTLE_ENDIAN);
+            a.min(b).intoMemorySegment(v1.get(), v1.offset(i), ByteOrder.LITTLE_ENDIAN);
+        }
+
+        // Process the tail
+        for (int i = vectorizedLength; i < v1.length(); i++) {
+            v1.set(i,  Math.min(v1.get(i), v2.get(i)));
+        }
+    }
+
     public static float max(MemorySegmentVectorFloat vector) {
         var accum = FloatVector.broadcast(FloatVector.SPECIES_PREFERRED, -Float.MAX_VALUE);
         int vectorizedLength = FloatVector.SPECIES_PREFERRED.loopBound(vector.length());

jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java

@@ -105,6 +105,11 @@ public VectorFloat<?> sub(VectorFloat<?> a, int aOffset, VectorFloat<?> b, int b
         return SimdOps.sub((ArrayVectorFloat) a, aOffset, (ArrayVectorFloat) b, bOffset, length);
     }
 
+    @Override
+    public void minInPlace(VectorFloat<?> v1, VectorFloat<?> v2) {
+        SimdOps.minInPlace((ArrayVectorFloat)v1, (ArrayVectorFloat)v2);
+    }
+
     @Override
     public float assembleAndSum(VectorFloat<?> data, int dataBase, ByteSequence<?> baseOffsets) {
         return SimdOps.assembleAndSum(((ArrayVectorFloat) data).get(), dataBase, ((ByteSequence<byte[]>) baseOffsets));

jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/SimdOps.java

@@ -612,6 +612,26 @@ static VectorFloat<?> sub(ArrayVectorFloat a, int aOffset, float value, int leng
         return new ArrayVectorFloat(res);
     }
 
+    static void minInPlace(ArrayVectorFloat v1, ArrayVectorFloat v2) {
+        if (v1.length() != v2.length()) {
+            throw new IllegalArgumentException("Vectors must have the same length");
+        }
+
+        int vectorizedLength = FloatVector.SPECIES_PREFERRED.loopBound(v1.length());
+
+        // Process the vectorized part
+        for (int i = 0; i < vectorizedLength; i += FloatVector.SPECIES_PREFERRED.length()) {
+            var a = FloatVector.fromArray(FloatVector.SPECIES_PREFERRED, v1.get(), i);
+            var b = FloatVector.fromArray(FloatVector.SPECIES_PREFERRED, v2.get(), i);
+            a.min(b).intoArray(v1.get(), i);
+        }
+
+        // Process the tail
+        for (int i = vectorizedLength; i < v1.length(); i++) {
+            v1.set(i,  Math.min(v1.get(i), v2.get(i)));
+        }
+    }
+
     static float assembleAndSum(float[] data, int dataBase, ByteSequence<byte[]> baseOffsets) {
         return switch (PREFERRED_BIT_SIZE)
         {