| 
 | 1 | +/*  | 
 | 2 | + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one  | 
 | 3 | + * or more contributor license agreements. Licensed under the "Elastic License  | 
 | 4 | + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side  | 
 | 5 | + * Public License v 1"; you may not use this file except in compliance with, at  | 
 | 6 | + * your election, the "Elastic License 2.0", the "GNU Affero General Public  | 
 | 7 | + * License v3.0 only", or the "Server Side Public License, v 1".  | 
 | 8 | + */  | 
 | 9 | +package org.elasticsearch.benchmark.vector;  | 
 | 10 | + | 
 | 11 | +import org.apache.lucene.util.VectorUtil;  | 
 | 12 | +import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;  | 
 | 13 | +import org.elasticsearch.common.logging.LogConfigurator;  | 
 | 14 | +import org.elasticsearch.simdvec.ESVectorUtil;  | 
 | 15 | +import org.openjdk.jmh.annotations.Benchmark;  | 
 | 16 | +import org.openjdk.jmh.annotations.BenchmarkMode;  | 
 | 17 | +import org.openjdk.jmh.annotations.Fork;  | 
 | 18 | +import org.openjdk.jmh.annotations.Measurement;  | 
 | 19 | +import org.openjdk.jmh.annotations.Mode;  | 
 | 20 | +import org.openjdk.jmh.annotations.OutputTimeUnit;  | 
 | 21 | +import org.openjdk.jmh.annotations.Param;  | 
 | 22 | +import org.openjdk.jmh.annotations.Scope;  | 
 | 23 | +import org.openjdk.jmh.annotations.Setup;  | 
 | 24 | +import org.openjdk.jmh.annotations.State;  | 
 | 25 | +import org.openjdk.jmh.annotations.Warmup;  | 
 | 26 | +import org.openjdk.jmh.infra.Blackhole;  | 
 | 27 | + | 
 | 28 | +import java.io.IOException;  | 
 | 29 | +import java.util.Random;  | 
 | 30 | +import java.util.concurrent.TimeUnit;  | 
 | 31 | + | 
 | 32 | +@BenchmarkMode(Mode.Throughput)  | 
 | 33 | +@OutputTimeUnit(TimeUnit.MILLISECONDS)  | 
 | 34 | +@State(Scope.Benchmark)  | 
 | 35 | +// first iteration is complete garbage, so make sure we really warmup  | 
 | 36 | +@Warmup(iterations = 4, time = 1)  | 
 | 37 | +// real iterations. not useful to spend tons of time here, better to fork more  | 
 | 38 | +@Measurement(iterations = 5, time = 1)  | 
 | 39 | +// engage some noise reduction  | 
 | 40 | +@Fork(value = 1)  | 
 | 41 | +public class DistanceBulkBenchmark {  | 
 | 42 | + | 
 | 43 | +    static {  | 
 | 44 | +        LogConfigurator.configureESLogging(); // native access requires logging to be initialized  | 
 | 45 | +    }  | 
 | 46 | + | 
 | 47 | +    @Param({ "384", "782", "1024" })  | 
 | 48 | +    int dims;  | 
 | 49 | + | 
 | 50 | +    int length;  | 
 | 51 | + | 
 | 52 | +    int numVectors = 4 * 100;  | 
 | 53 | +    int numQueries = 10;  | 
 | 54 | + | 
 | 55 | +    float[][] vectors;  | 
 | 56 | +    float[][] queries;  | 
 | 57 | +    float[] distances = new float[4];  | 
 | 58 | + | 
 | 59 | +    @Setup  | 
 | 60 | +    public void setup() throws IOException {  | 
 | 61 | +        Random random = new Random(123);  | 
 | 62 | + | 
 | 63 | +        this.length = OptimizedScalarQuantizer.discretize(dims, 64) / 8;  | 
 | 64 | + | 
 | 65 | +        vectors = new float[numVectors][dims];  | 
 | 66 | +        for (float[] vector : vectors) {  | 
 | 67 | +            for (int i = 0; i < dims; i++) {  | 
 | 68 | +                vector[i] = random.nextFloat();  | 
 | 69 | +            }  | 
 | 70 | +        }  | 
 | 71 | + | 
 | 72 | +        queries = new float[numQueries][dims];  | 
 | 73 | +        for (float[] query : queries) {  | 
 | 74 | +            for (int i = 0; i < dims; i++) {  | 
 | 75 | +                query[i] = random.nextFloat();  | 
 | 76 | +            }  | 
 | 77 | +        }  | 
 | 78 | +    }  | 
 | 79 | + | 
 | 80 | +    @Benchmark  | 
 | 81 | +    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 82 | +    public void squareDistance(Blackhole bh) {  | 
 | 83 | +        for (int j = 0; j < numQueries; j++) {  | 
 | 84 | +            float[] query = queries[j];  | 
 | 85 | +            for (int i = 0; i < numVectors; i++) {  | 
 | 86 | +                float[] vector = vectors[i];  | 
 | 87 | +                float distance = VectorUtil.squareDistance(query, vector);  | 
 | 88 | +                bh.consume(distance);  | 
 | 89 | +            }  | 
 | 90 | +        }  | 
 | 91 | +    }  | 
 | 92 | + | 
 | 93 | +    @Benchmark  | 
 | 94 | +    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 95 | +    public void soarDistance(Blackhole bh) {  | 
 | 96 | +        for (int j = 0; j < numQueries; j++) {  | 
 | 97 | +            float[] query = queries[j];  | 
 | 98 | +            for (int i = 0; i < numVectors; i++) {  | 
 | 99 | +                float[] vector = vectors[i];  | 
 | 100 | +                float distance = ESVectorUtil.soarDistance(query, vector, vector, 1.0f, 1.0f);  | 
 | 101 | +                bh.consume(distance);  | 
 | 102 | +            }  | 
 | 103 | +        }  | 
 | 104 | +    }  | 
 | 105 | + | 
 | 106 | +    @Benchmark  | 
 | 107 | +    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 108 | +    public void squareDistanceBulk(Blackhole bh) {  | 
 | 109 | +        for (int j = 0; j < numQueries; j++) {  | 
 | 110 | +            float[] query = queries[j];  | 
 | 111 | +            for (int i = 0; i < numVectors; i += 4) {  | 
 | 112 | +                ESVectorUtil.squareDistanceBulk(query, vectors[i], vectors[i + 1], vectors[i + 2], vectors[i + 3], distances);  | 
 | 113 | +                for (float distance : distances) {  | 
 | 114 | +                    bh.consume(distance);  | 
 | 115 | +                }  | 
 | 116 | + | 
 | 117 | +            }  | 
 | 118 | +        }  | 
 | 119 | +    }  | 
 | 120 | + | 
 | 121 | +    @Benchmark  | 
 | 122 | +    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 123 | +    public void soarDistanceBulk(Blackhole bh) {  | 
 | 124 | +        for (int j = 0; j < numQueries; j++) {  | 
 | 125 | +            float[] query = queries[j];  | 
 | 126 | +            for (int i = 0; i < numVectors; i += 4) {  | 
 | 127 | +                ESVectorUtil.soarDistanceBulk(  | 
 | 128 | +                    query,  | 
 | 129 | +                    vectors[i],  | 
 | 130 | +                    vectors[i + 1],  | 
 | 131 | +                    vectors[i + 2],  | 
 | 132 | +                    vectors[i + 3],  | 
 | 133 | +                    vectors[i],  | 
 | 134 | +                    1.0f,  | 
 | 135 | +                    1.0f,  | 
 | 136 | +                    distances  | 
 | 137 | +                );  | 
 | 138 | +                for (float distance : distances) {  | 
 | 139 | +                    bh.consume(distance);  | 
 | 140 | +                }  | 
 | 141 | + | 
 | 142 | +            }  | 
 | 143 | +        }  | 
 | 144 | +    }  | 
 | 145 | +}  | 
0 commit comments