From a301a879d965e607d90e63fea891a5c16dc0490e Mon Sep 17 00:00:00 2001 From: Lorenzo Dematte Date: Fri, 27 Mar 2026 12:50:24 +0100 Subject: [PATCH] Adds native VectorScorerInt4OperationBenchmark + tests; fixes VectorScorerInt4 benchmark tests to run only on JDK22+, as the native scorer requires that. --- .../VectorScorerInt4OperationBenchmark.java | 62 ++++++++++++++++++- .../VectorScorerInt4BenchmarkTests.java | 5 +- .../VectorScorerInt4BulkBenchmarkTests.java | 5 +- ...ctorScorerInt4OperationBenchmarkTests.java | 7 ++- 4 files changed, 74 insertions(+), 5 deletions(-) diff --git a/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4OperationBenchmark.java b/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4OperationBenchmark.java index 3f5ee399245fe..5ccd888a45b8f 100644 --- a/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4OperationBenchmark.java +++ b/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4OperationBenchmark.java @@ -10,6 +10,8 @@ import org.apache.lucene.util.VectorUtil; import org.elasticsearch.benchmark.Utils; +import org.elasticsearch.nativeaccess.NativeAccess; +import org.elasticsearch.nativeaccess.VectorSimilarityFunctions; import org.openjdk.jmh.annotations.Benchmark; import org.openjdk.jmh.annotations.BenchmarkMode; import org.openjdk.jmh.annotations.Fork; @@ -21,11 +23,17 @@ import org.openjdk.jmh.annotations.Scope; import org.openjdk.jmh.annotations.Setup; import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; import org.openjdk.jmh.annotations.Warmup; +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import java.lang.invoke.MethodHandle; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; +import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.rethrow; import static org.elasticsearch.nativeaccess.Int4TestUtils.dotProductI4SinglePacked; import static org.elasticsearch.nativeaccess.Int4TestUtils.packNibbles; import static org.elasticsearch.simdvec.internal.vectorization.VectorScorerTestUtils.randomInt4Bytes; @@ -47,8 +55,17 @@ public class VectorScorerInt4OperationBenchmark { Utils.configureBenchmarkLogging(); } - public byte[] unpacked; - public byte[] packed; + private int packedLen; + + private byte[] unpacked; + private byte[] packed; + + MemorySegment unpackedHeapSeg, packedHeapSeg; + MemorySegment unpackedNativeSeg, packedNativeSeg; + + Arena arena; + + private MethodHandle nativeImpl; @Param({ "2", "128", "208", "256", "300", "512", "702", "1024", "1536", "2048" }) public int size; @@ -58,6 +75,27 @@ public void init() { unpacked = new byte[size]; randomInt4Bytes(ThreadLocalRandom.current(), unpacked); packed = packNibbles(unpacked); + packedLen = packed.length; + + unpackedHeapSeg = MemorySegment.ofArray(unpacked); + packedHeapSeg = MemorySegment.ofArray(packed); + + arena = Arena.ofConfined(); + unpackedNativeSeg = arena.allocate(unpacked.length); + MemorySegment.copy(unpacked, 0, unpackedNativeSeg, ValueLayout.JAVA_BYTE, 0L, unpacked.length); + packedNativeSeg = arena.allocate(packed.length); + MemorySegment.copy(packed, 0, packedNativeSeg, ValueLayout.JAVA_BYTE, 0L, packed.length); + + nativeImpl = vectorSimilarityFunctions.getHandle( + VectorSimilarityFunctions.Function.DOT_PRODUCT, + VectorSimilarityFunctions.DataType.INT4, + VectorSimilarityFunctions.Operation.SINGLE + ); + } + + @TearDown + public void teardown() { + arena.close(); } @Benchmark @@ -69,4 +107,24 @@ public int scalar() { public int lucene() { return VectorUtil.int4DotProductSinglePacked(unpacked, packed); } + + @Benchmark + public int nativeWithNativeSeg() { + try { + return (int) nativeImpl.invokeExact(unpackedNativeSeg, packedNativeSeg, packedLen); + } catch (Throwable t) { + throw rethrow(t); + } + } + + @Benchmark + public int nativeWithHeapSeg() { + try { + return (int) nativeImpl.invokeExact(unpackedHeapSeg, packedHeapSeg, packedLen); + } catch (Throwable t) { + throw rethrow(t); + } + } + + static final VectorSimilarityFunctions vectorSimilarityFunctions = NativeAccess.instance().getVectorSimilarityFunctions().orElseThrow(); } diff --git a/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4BenchmarkTests.java b/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4BenchmarkTests.java index e0d3689984462..7cbe142a69b82 100644 --- a/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4BenchmarkTests.java +++ b/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4BenchmarkTests.java @@ -19,6 +19,8 @@ import java.util.Arrays; +import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.supportsHeapSegments; + public class VectorScorerInt4BenchmarkTests extends ESTestCase { private final double delta = 1e-3; @@ -40,8 +42,9 @@ private VectorScorerInt4Benchmark createBench(VectorImplementation impl, VectorS } @BeforeClass - public static void skipWindows() { + public static void skipUnsupported() { assumeFalse("doesn't work on windows yet", Constants.WINDOWS); + assumeTrue("native requires JDK22+", supportsHeapSegments()); } public void testScores() throws Exception { diff --git a/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4BulkBenchmarkTests.java b/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4BulkBenchmarkTests.java index 1900e507796d2..a83eafc58b8df 100644 --- a/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4BulkBenchmarkTests.java +++ b/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4BulkBenchmarkTests.java @@ -19,6 +19,8 @@ import java.util.Arrays; +import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.supportsHeapSegments; + public class VectorScorerInt4BulkBenchmarkTests extends ESTestCase { private final VectorSimilarityType function; @@ -44,8 +46,9 @@ private VectorScorerInt4BulkBenchmark createBench(VectorImplementation impl, Vec } @BeforeClass - public static void skipWindows() { + public static void skipUnsupported() { assumeFalse("doesn't work on windows yet", Constants.WINDOWS); + assumeTrue("native requires JDK22+", supportsHeapSegments()); } public void testSequential() throws Exception { diff --git a/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4OperationBenchmarkTests.java b/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4OperationBenchmarkTests.java index f5a1280ded1f9..d17090d4a9bd4 100644 --- a/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4OperationBenchmarkTests.java +++ b/benchmarks/src/test/java/org/elasticsearch/benchmark/vector/scorer/VectorScorerInt4OperationBenchmarkTests.java @@ -18,6 +18,8 @@ import java.util.Arrays; +import static org.elasticsearch.benchmark.vector.scorer.BenchmarkUtils.supportsHeapSegments; + public class VectorScorerInt4OperationBenchmarkTests extends ESTestCase { private final int size; @@ -27,8 +29,9 @@ public VectorScorerInt4OperationBenchmarkTests(int size) { } @BeforeClass - public static void skipWindows() { + public static void skipUnsupported() { assumeFalse("doesn't work on windows yet", Constants.WINDOWS); + assumeTrue("native requires JDK22+", supportsHeapSegments()); } public void test() { @@ -39,6 +42,8 @@ public void test() { int expected = bench.scalar(); assertEquals(expected, bench.lucene()); + assertEquals(expected, bench.nativeWithNativeSeg()); + assertEquals(expected, bench.nativeWithHeapSeg()); } }