Skip to content

Commit be4b635

Browse files
prwhelanprwhelan
committed
Merge branch 'main' of github.com:prwhelan/elasticsearch into fix/133135-2-1
2 parents c5c1bfa + 3090a4a commit be4b635

File tree

9 files changed

+648
-52
lines changed

9 files changed

+648
-52
lines changed

‎benchmarks/src/main/java/org/elasticsearch/benchmark/vector/OSQScorerBenchmark.java‎

Lines changed: 88 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,10 @@
1414
import org.apache.lucene.store.IndexInput;
1515
import org.apache.lucene.store.IndexOutput;
1616
import org.apache.lucene.store.MMapDirectory;
17-
import org.apache.lucene.util.VectorUtil;
17+
import org.apache.lucene.store.NIOFSDirectory;
1818
import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
1919
import org.elasticsearch.common.logging.LogConfigurator;
20+
import org.elasticsearch.core.IOUtils;
2021
import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
2122
import org.elasticsearch.simdvec.internal.vectorization.ESVectorizationProvider;
2223
import org.openjdk.jmh.annotations.Benchmark;
@@ -29,6 +30,7 @@
2930
import org.openjdk.jmh.annotations.Scope;
3031
import org.openjdk.jmh.annotations.Setup;
3132
import org.openjdk.jmh.annotations.State;
33+
import org.openjdk.jmh.annotations.TearDown;
3234
import org.openjdk.jmh.annotations.Warmup;
3335
import org.openjdk.jmh.infra.Blackhole;
3436

@@ -66,9 +68,14 @@ public class OSQScorerBenchmark {
6668
float centroidDp;
6769

6870
byte[] scratch;
69-
ES91OSQVectorsScorer scorer;
71+
ES91OSQVectorsScorer scorerMmap;
72+
ES91OSQVectorsScorer scorerNfios;
7073

71-
IndexInput in;
74+
Directory dirMmap;
75+
IndexInput inMmap;
76+
77+
Directory dirNiofs;
78+
IndexInput inNiofs;
7279

7380
float[] scratchScores;
7481
float[] corrections;
@@ -84,18 +91,24 @@ public void setup() throws IOException {
8491
random.nextBytes(binaryVector);
8592
}
8693

87-
Directory dir = new MMapDirectory(Files.createTempDirectory("vectorData"));
88-
IndexOutput out = dir.createOutput("vectors", IOContext.DEFAULT);
94+
dirMmap = new MMapDirectory(Files.createTempDirectory("vectorDataMmap"));
95+
dirNiofs = new NIOFSDirectory(Files.createTempDirectory("vectorDataNFIOS"));
96+
IndexOutput outMmap = dirMmap.createOutput("vectors", IOContext.DEFAULT);
97+
IndexOutput outNfios = dirNiofs.createOutput("vectors", IOContext.DEFAULT);
8998
byte[] correctionBytes = new byte[14 * ES91OSQVectorsScorer.BULK_SIZE];
9099
for (int i = 0; i < numVectors; i += ES91OSQVectorsScorer.BULK_SIZE) {
91100
for (int j = 0; j < ES91OSQVectorsScorer.BULK_SIZE; j++) {
92-
out.writeBytes(binaryVectors[i + j], 0, binaryVectors[i + j].length);
101+
outMmap.writeBytes(binaryVectors[i + j], 0, binaryVectors[i + j].length);
102+
outNfios.writeBytes(binaryVectors[i + j], 0, binaryVectors[i + j].length);
93103
}
94104
random.nextBytes(correctionBytes);
95-
out.writeBytes(correctionBytes, 0, correctionBytes.length);
105+
outMmap.writeBytes(correctionBytes, 0, correctionBytes.length);
106+
outNfios.writeBytes(correctionBytes, 0, correctionBytes.length);
96107
}
97-
out.close();
98-
in = dir.openInput("vectors", IOContext.DEFAULT);
108+
outMmap.close();
109+
outNfios.close();
110+
inMmap = dirMmap.openInput("vectors", IOContext.DEFAULT);
111+
inNiofs = dirNiofs.openInput("vectors", IOContext.DEFAULT);
99112

100113
binaryQueries = new byte[numVectors][4 * length];
101114
for (byte[] binaryVector : binaryVectors) {
@@ -110,42 +123,40 @@ public void setup() throws IOException {
110123
centroidDp = random.nextFloat();
111124

112125
scratch = new byte[length];
113-
scorer = ESVectorizationProvider.getInstance().newES91OSQVectorsScorer(in, dims);
126+
scorerMmap = ESVectorizationProvider.getInstance().newES91OSQVectorsScorer(inMmap, dims);
127+
scorerNfios = ESVectorizationProvider.getInstance().newES91OSQVectorsScorer(inNiofs, dims);
114128
scratchScores = new float[16];
115129
corrections = new float[3];
116130
}
117131

132+
@TearDown
133+
public void teardown() throws IOException {
134+
IOUtils.close(dirMmap, inMmap, dirNiofs, inNiofs);
135+
}
136+
137+
@Benchmark
138+
public void scoreFromMemorySegmentOnlyVectorMmapScalar(Blackhole bh) throws IOException {
139+
scoreFromMemorySegmentOnlyVector(bh, inMmap, scorerMmap);
140+
}
141+
118142
@Benchmark
119143
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
120-
public void scoreFromArray(Blackhole bh) throws IOException {
121-
for (int j = 0; j < numQueries; j++) {
122-
in.seek(0);
123-
for (int i = 0; i < numVectors; i++) {
124-
in.readBytes(scratch, 0, length);
125-
float qDist = VectorUtil.int4BitDotProduct(binaryQueries[j], scratch);
126-
in.readFloats(corrections, 0, corrections.length);
127-
int addition = Short.toUnsignedInt(in.readShort());
128-
float score = scorer.score(
129-
result.lowerInterval(),
130-
result.upperInterval(),
131-
result.quantizedComponentSum(),
132-
result.additionalCorrection(),
133-
VectorSimilarityFunction.EUCLIDEAN,
134-
centroidDp,
135-
corrections[0],
136-
corrections[1],
137-
addition,
138-
corrections[2],
139-
qDist
140-
);
141-
bh.consume(score);
142-
}
143-
}
144+
public void scoreFromMemorySegmentOnlyVectorMmapVect(Blackhole bh) throws IOException {
145+
scoreFromMemorySegmentOnlyVector(bh, inMmap, scorerMmap);
146+
}
147+
148+
@Benchmark
149+
public void scoreFromMemorySegmentOnlyVectorNiofsScalar(Blackhole bh) throws IOException {
150+
scoreFromMemorySegmentOnlyVector(bh, inNiofs, scorerNfios);
144151
}
145152

146153
@Benchmark
147154
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
148-
public void scoreFromMemorySegmentOnlyVector(Blackhole bh) throws IOException {
155+
public void scoreFromMemorySegmentOnlyVectorNiofsVect(Blackhole bh) throws IOException {
156+
scoreFromMemorySegmentOnlyVector(bh, inNiofs, scorerNfios);
157+
}
158+
159+
private void scoreFromMemorySegmentOnlyVector(Blackhole bh, IndexInput in, ES91OSQVectorsScorer scorer) throws IOException {
149160
for (int j = 0; j < numQueries; j++) {
150161
in.seek(0);
151162
for (int i = 0; i < numVectors; i++) {
@@ -170,9 +181,29 @@ public void scoreFromMemorySegmentOnlyVector(Blackhole bh) throws IOException {
170181
}
171182
}
172183

184+
@Benchmark
185+
public void scoreFromMemorySegmentOnlyVectorBulkMmapScalar(Blackhole bh) throws IOException {
186+
scoreFromMemorySegmentOnlyVectorBulk(bh, inMmap, scorerMmap);
187+
}
188+
189+
@Benchmark
190+
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
191+
public void scoreFromMemorySegmentOnlyVectorBulkMmapVect(Blackhole bh) throws IOException {
192+
scoreFromMemorySegmentOnlyVectorBulk(bh, inMmap, scorerMmap);
193+
}
194+
195+
@Benchmark
196+
public void scoreFromMemorySegmentOnlyVectorBulkNiofsScalar(Blackhole bh) throws IOException {
197+
scoreFromMemorySegmentOnlyVectorBulk(bh, inNiofs, scorerNfios);
198+
}
199+
173200
@Benchmark
174201
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
175-
public void scoreFromMemorySegmentOnlyVectorBulk(Blackhole bh) throws IOException {
202+
public void scoreFromMemorySegmentOnlyVectorBulkNiofsVect(Blackhole bh) throws IOException {
203+
scoreFromMemorySegmentOnlyVectorBulk(bh, inNiofs, scorerNfios);
204+
}
205+
206+
private void scoreFromMemorySegmentOnlyVectorBulk(Blackhole bh, IndexInput in, ES91OSQVectorsScorer scorer) throws IOException {
176207
for (int j = 0; j < numQueries; j++) {
177208
in.seek(0);
178209
for (int i = 0; i < numVectors; i += 16) {
@@ -199,9 +230,29 @@ public void scoreFromMemorySegmentOnlyVectorBulk(Blackhole bh) throws IOExceptio
199230
}
200231
}
201232

233+
@Benchmark
234+
public void scoreFromMemorySegmentAllBulkMmapScalar(Blackhole bh) throws IOException {
235+
scoreFromMemorySegmentAllBulk(bh, inMmap, scorerMmap);
236+
}
237+
238+
@Benchmark
239+
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
240+
public void scoreFromMemorySegmentAllBulkMmapVect(Blackhole bh) throws IOException {
241+
scoreFromMemorySegmentAllBulk(bh, inMmap, scorerMmap);
242+
}
243+
244+
@Benchmark
245+
public void scoreFromMemorySegmentAllBulkNiofsScalar(Blackhole bh) throws IOException {
246+
scoreFromMemorySegmentAllBulk(bh, inNiofs, scorerNfios);
247+
}
248+
202249
@Benchmark
203250
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
204-
public void scoreFromMemorySegmentAllBulk(Blackhole bh) throws IOException {
251+
public void scoreFromMemorySegmentAllBulkNiofsVect(Blackhole bh) throws IOException {
252+
scoreFromMemorySegmentAllBulk(bh, inNiofs, scorerNfios);
253+
}
254+
255+
private void scoreFromMemorySegmentAllBulk(Blackhole bh, IndexInput in, ES91OSQVectorsScorer scorer) throws IOException {
205256
for (int j = 0; j < numQueries; j++) {
206257
in.seek(0);
207258
for (int i = 0; i < numVectors; i += 16) {

‎libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentES91OSQVectorsScorer.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN;
3030
import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT;
3131

32-
/** Panamized scorer for quantized vectors stored as an {@link IndexInput}. */
32+
/** Panamized scorer for quantized vectors stored as a {@link MemorySegment}. */
3333
public final class MemorySegmentES91OSQVectorsScorer extends ES91OSQVectorsScorer {
3434

3535
private static final VectorSpecies<Integer> INT_SPECIES_128 = IntVector.SPECIES_128;

0 commit comments

Comments
 (0)