add benchmark

Author: davidkyle

benchmarks/src/main/java/org/elasticsearch/benchmark/vector/scorer/VectorUtilsBenchmark.java

@@ -0,0 +1,84 @@
+/*
+ * 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.scorer;
+
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
+import org.elasticsearch.common.logging.LogConfigurator;
+import org.elasticsearch.simdvec.ESVectorUtil;
+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.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+import java.io.IOException;
+import java.util.Random;
+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 VectorUtilsBenchmark {
+
+    static {
+        LogConfigurator.configureESLogging(); // native access requires logging to be initialized
+    }
+
+    @Param({ "384", "782", "1024" })
+    int dims;
+
+    int numVectors = 4; // * 100;
+    float[][] vectors;
+
+    @Setup
+    public void setup() throws IOException {
+        Random random = new Random(123);
+
+        vectors = new float[numVectors][dims];
+        for (float[] vector : vectors) {
+            for (int i = 0; i < dims; i++) {
+                vector[i] = random.nextFloat();
+            }
+        }
+    }
+
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void vectorAcc(Blackhole bh) {
+        float[] acc = new float[dims];
+        for (int i = 0; i < numVectors; i++) {
+            ESVectorUtil.vectorAccumulateAdd(acc, vectors[i]);
+        }
+    }
+
+    @Benchmark
+    public void defaultVectorAcc(Blackhole bh) {
+        float[] acc = new float[dims];
+        for (int i = 0; i < numVectors; i++) {
+            float[] vector = vectors[i];
+            for (int j = 0; j < vector.length; j++) {
+                acc[j] += vector[j];
+            }
+        }
+    }
+}

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

@@ -1143,12 +1143,11 @@ public int indexOf(final byte[] bytes, final int offset, final int length, final
 
     @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);
+        int limit = FLOAT_SPECIES.loopBound(a.length);
+        for (; i < limit; i += FLOAT_SPECIES.length()) {
+            FloatVector va = FloatVector.fromArray(FLOAT_SPECIES, a, i);
+            FloatVector vb = FloatVector.fromArray(FLOAT_SPECIES, b, i);
             FloatVector vc = va.add(vb);
             vc.intoArray(a, i);
         }
@@ -1160,11 +1159,9 @@ public void vectorAccumulateAdd(float[] a, float[] b) {
 
     @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);
+        for (; i < FLOAT_SPECIES.loopBound(a.length); i += FLOAT_SPECIES.length()) {
+            FloatVector va = FloatVector.fromArray(FLOAT_SPECIES, a, i);
             FloatVector vc = va.div(b);
             vc.intoArray(a, i);
         }

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

@@ -107,9 +107,10 @@ private static boolean stepLloyd(
                 final int assignment = assignments[translateOrd.apply(idx)];
                 if (centroidChanged.get(assignment)) {
                     if (centroidCounts[assignment]++ == 0) {
-                        centroids[assignment] = vectors.vectorValue(idx);
+                        System.arraycopy(vectors.vectorValue(idx), 0, centroids[assignment], 0, vectors.vectorValue(idx).length);
                         continue;
                     }
+
                     ESVectorUtil.vectorAccumulateAdd(centroids[assignment], vectors.vectorValue(idx));
                 }
             }