| 
 | 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.elasticsearch.common.logging.LogConfigurator;  | 
 | 13 | +import org.elasticsearch.common.logging.NodeNamePatternConverter;  | 
 | 14 | +import org.elasticsearch.nativeaccess.NativeAccess;  | 
 | 15 | +import org.elasticsearch.nativeaccess.VectorSimilarityFunctions;  | 
 | 16 | +import org.openjdk.jmh.annotations.Benchmark;  | 
 | 17 | +import org.openjdk.jmh.annotations.BenchmarkMode;  | 
 | 18 | +import org.openjdk.jmh.annotations.Fork;  | 
 | 19 | +import org.openjdk.jmh.annotations.Level;  | 
 | 20 | +import org.openjdk.jmh.annotations.Measurement;  | 
 | 21 | +import org.openjdk.jmh.annotations.Mode;  | 
 | 22 | +import org.openjdk.jmh.annotations.OutputTimeUnit;  | 
 | 23 | +import org.openjdk.jmh.annotations.Param;  | 
 | 24 | +import org.openjdk.jmh.annotations.Scope;  | 
 | 25 | +import org.openjdk.jmh.annotations.Setup;  | 
 | 26 | +import org.openjdk.jmh.annotations.State;  | 
 | 27 | +import org.openjdk.jmh.annotations.TearDown;  | 
 | 28 | +import org.openjdk.jmh.annotations.Warmup;  | 
 | 29 | + | 
 | 30 | +import java.lang.foreign.Arena;  | 
 | 31 | +import java.lang.foreign.MemorySegment;  | 
 | 32 | +import java.lang.foreign.ValueLayout;  | 
 | 33 | +import java.nio.ByteOrder;  | 
 | 34 | +import java.util.concurrent.ThreadLocalRandom;  | 
 | 35 | +import java.util.concurrent.TimeUnit;  | 
 | 36 | + | 
 | 37 | +@BenchmarkMode(Mode.AverageTime)  | 
 | 38 | +@OutputTimeUnit(TimeUnit.NANOSECONDS)  | 
 | 39 | +@State(Scope.Benchmark)  | 
 | 40 | +@Warmup(iterations = 3, time = 1)  | 
 | 41 | +@Measurement(iterations = 5, time = 1)  | 
 | 42 | +public class JDKVectorFloat32Benchmark {  | 
 | 43 | + | 
 | 44 | +    static {  | 
 | 45 | +        NodeNamePatternConverter.setGlobalNodeName("foo");  | 
 | 46 | +        LogConfigurator.loadLog4jPlugins();  | 
 | 47 | +        LogConfigurator.configureESLogging(); // native access requires logging to be initialized  | 
 | 48 | +    }  | 
 | 49 | + | 
 | 50 | +    static final ValueLayout.OfFloat LAYOUT_LE_FLOAT = ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);  | 
 | 51 | + | 
 | 52 | +    float[] floatsA;  | 
 | 53 | +    float[] floatsB;  | 
 | 54 | +    float[] scratch;  | 
 | 55 | +    MemorySegment heapSegA, heapSegB;  | 
 | 56 | +    MemorySegment nativeSegA, nativeSegB;  | 
 | 57 | + | 
 | 58 | +    Arena arena;  | 
 | 59 | + | 
 | 60 | +    @Param({ "1", "128", "207", "256", "300", "512", "702", "1024", "1536", "2048" })  | 
 | 61 | +    public int size;  | 
 | 62 | + | 
 | 63 | +    @Setup(Level.Iteration)  | 
 | 64 | +    public void init() {  | 
 | 65 | +        ThreadLocalRandom random = ThreadLocalRandom.current();  | 
 | 66 | + | 
 | 67 | +        floatsA = new float[size];  | 
 | 68 | +        floatsB = new float[size];  | 
 | 69 | +        scratch = new float[size];  | 
 | 70 | +        for (int i = 0; i < size; ++i) {  | 
 | 71 | +            floatsA[i] = random.nextFloat();  | 
 | 72 | +            floatsB[i] = random.nextFloat();  | 
 | 73 | +        }  | 
 | 74 | +        heapSegA = MemorySegment.ofArray(floatsA);  | 
 | 75 | +        heapSegB = MemorySegment.ofArray(floatsB);  | 
 | 76 | + | 
 | 77 | +        arena = Arena.ofConfined();  | 
 | 78 | +        nativeSegA = arena.allocate((long) floatsA.length * Float.BYTES);  | 
 | 79 | +        MemorySegment.copy(MemorySegment.ofArray(floatsA), LAYOUT_LE_FLOAT, 0L, nativeSegA, LAYOUT_LE_FLOAT, 0L, floatsA.length);  | 
 | 80 | +        nativeSegB = arena.allocate((long) floatsB.length * Float.BYTES);  | 
 | 81 | +        MemorySegment.copy(MemorySegment.ofArray(floatsB), LAYOUT_LE_FLOAT, 0L, nativeSegB, LAYOUT_LE_FLOAT, 0L, floatsB.length);  | 
 | 82 | +    }  | 
 | 83 | + | 
 | 84 | +    @TearDown  | 
 | 85 | +    public void teardown() {  | 
 | 86 | +        arena.close();  | 
 | 87 | +    }  | 
 | 88 | + | 
 | 89 | +    // -- cosine  | 
 | 90 | + | 
 | 91 | +    @Benchmark  | 
 | 92 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 93 | +    public float cosineLucene() {  | 
 | 94 | +        return VectorUtil.cosine(floatsA, floatsB);  | 
 | 95 | +    }  | 
 | 96 | + | 
 | 97 | +    @Benchmark  | 
 | 98 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 99 | +    public float cosineLuceneWithCopy() {  | 
 | 100 | +        // add a copy to better reflect what Lucene has to do to get the target vector on-heap  | 
 | 101 | +        MemorySegment.copy(nativeSegB, LAYOUT_LE_FLOAT, 0L, scratch, 0, scratch.length);  | 
 | 102 | +        return VectorUtil.cosine(floatsA, scratch);  | 
 | 103 | +    }  | 
 | 104 | + | 
 | 105 | +    @Benchmark  | 
 | 106 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 107 | +    public float cosineNativeWithNativeSeg() {  | 
 | 108 | +        return cosineFloat32(nativeSegA, nativeSegB, size);  | 
 | 109 | +    }  | 
 | 110 | + | 
 | 111 | +    @Benchmark  | 
 | 112 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 113 | +    public float cosineNativeWithHeapSeg() {  | 
 | 114 | +        return cosineFloat32(heapSegA, heapSegB, size);  | 
 | 115 | +    }  | 
 | 116 | + | 
 | 117 | +    // -- dot product  | 
 | 118 | + | 
 | 119 | +    @Benchmark  | 
 | 120 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 121 | +    public float dotProductLucene() {  | 
 | 122 | +        return VectorUtil.dotProduct(floatsA, floatsB);  | 
 | 123 | +    }  | 
 | 124 | + | 
 | 125 | +    @Benchmark  | 
 | 126 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 127 | +    public float dotProductLuceneWithCopy() {  | 
 | 128 | +        // add a copy to better reflect what Lucene has to do to get the target vector on-heap  | 
 | 129 | +        MemorySegment.copy(nativeSegB, LAYOUT_LE_FLOAT, 0L, scratch, 0, scratch.length);  | 
 | 130 | +        return VectorUtil.dotProduct(floatsA, scratch);  | 
 | 131 | +    }  | 
 | 132 | + | 
 | 133 | +    @Benchmark  | 
 | 134 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 135 | +    public float dotProductNativeWithNativeSeg() {  | 
 | 136 | +        return dotProductFloat32(nativeSegA, nativeSegB, size);  | 
 | 137 | +    }  | 
 | 138 | + | 
 | 139 | +    @Benchmark  | 
 | 140 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 141 | +    public float dotProductNativeWithHeapSeg() {  | 
 | 142 | +        return dotProductFloat32(heapSegA, heapSegB, size);  | 
 | 143 | +    }  | 
 | 144 | + | 
 | 145 | +    // -- square distance  | 
 | 146 | + | 
 | 147 | +    @Benchmark  | 
 | 148 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 149 | +    public float squareDistanceLucene() {  | 
 | 150 | +        return VectorUtil.squareDistance(floatsA, floatsB);  | 
 | 151 | +    }  | 
 | 152 | + | 
 | 153 | +    @Benchmark  | 
 | 154 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 155 | +    public float squareDistanceLuceneWithCopy() {  | 
 | 156 | +        // add a copy to better reflect what Lucene has to do to get the target vector on-heap  | 
 | 157 | +        MemorySegment.copy(nativeSegB, LAYOUT_LE_FLOAT, 0L, scratch, 0, scratch.length);  | 
 | 158 | +        return VectorUtil.squareDistance(floatsA, scratch);  | 
 | 159 | +    }  | 
 | 160 | + | 
 | 161 | +    @Benchmark  | 
 | 162 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 163 | +    public float squareDistanceNativeWithNativeSeg() {  | 
 | 164 | +        return squareDistanceFloat32(nativeSegA, nativeSegB, size);  | 
 | 165 | +    }  | 
 | 166 | + | 
 | 167 | +    @Benchmark  | 
 | 168 | +    @Fork(value = 3, jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })  | 
 | 169 | +    public float squareDistanceNativeWithHeapSeg() {  | 
 | 170 | +        return squareDistanceFloat32(heapSegA, heapSegB, size);  | 
 | 171 | +    }  | 
 | 172 | + | 
 | 173 | +    static final VectorSimilarityFunctions vectorSimilarityFunctions = vectorSimilarityFunctions();  | 
 | 174 | + | 
 | 175 | +    static VectorSimilarityFunctions vectorSimilarityFunctions() {  | 
 | 176 | +        return NativeAccess.instance().getVectorSimilarityFunctions().get();  | 
 | 177 | +    }  | 
 | 178 | + | 
 | 179 | +    float cosineFloat32(MemorySegment a, MemorySegment b, int length) {  | 
 | 180 | +        try {  | 
 | 181 | +            return (float) vectorSimilarityFunctions.cosineHandleFloat32().invokeExact(a, b, length);  | 
 | 182 | +        } catch (Throwable e) {  | 
 | 183 | +            if (e instanceof Error err) {  | 
 | 184 | +                throw err;  | 
 | 185 | +            } else if (e instanceof RuntimeException re) {  | 
 | 186 | +                throw re;  | 
 | 187 | +            } else {  | 
 | 188 | +                throw new RuntimeException(e);  | 
 | 189 | +            }  | 
 | 190 | +        }  | 
 | 191 | +    }  | 
 | 192 | + | 
 | 193 | +    float dotProductFloat32(MemorySegment a, MemorySegment b, int length) {  | 
 | 194 | +        try {  | 
 | 195 | +            return (float) vectorSimilarityFunctions.dotProductHandleFloat32().invokeExact(a, b, length);  | 
 | 196 | +        } catch (Throwable e) {  | 
 | 197 | +            if (e instanceof Error err) {  | 
 | 198 | +                throw err;  | 
 | 199 | +            } else if (e instanceof RuntimeException re) {  | 
 | 200 | +                throw re;  | 
 | 201 | +            } else {  | 
 | 202 | +                throw new RuntimeException(e);  | 
 | 203 | +            }  | 
 | 204 | +        }  | 
 | 205 | +    }  | 
 | 206 | + | 
 | 207 | +    float squareDistanceFloat32(MemorySegment a, MemorySegment b, int length) {  | 
 | 208 | +        try {  | 
 | 209 | +            return (float) vectorSimilarityFunctions.squareDistanceHandleFloat32().invokeExact(a, b, length);  | 
 | 210 | +        } catch (Throwable e) {  | 
 | 211 | +            if (e instanceof Error err) {  | 
 | 212 | +                throw err;  | 
 | 213 | +            } else if (e instanceof RuntimeException re) {  | 
 | 214 | +                throw re;  | 
 | 215 | +            } else {  | 
 | 216 | +                throw new RuntimeException(e);  | 
 | 217 | +            }  | 
 | 218 | +        }  | 
 | 219 | +    }  | 
 | 220 | +}  | 
0 commit comments