JDKVectorLibrary: update low-level bounds checks and add benchmark

[ChrisHegarty]

This commit updates the low-level bounds checks in JDKVectorLibrary and add benchmark, so that we can more easily bench the low-level operations.

Note: I added the mr-jar gradle plugin to the benchmarks so that we can compile with preview features in Java 21, namely MemorySegment.

[elasticsearchmachine]

Pinging @elastic/es-search-relevance (Team:Search Relevance)

[benwtrent]

🥹 its beautiful

[benwtrent]

@ChrisHegarty I get error: package org.elasticsearch.nativeaccess does not exist when trying to compile, do we need to expose that module to the benchmarks?

[benwtrent]

@ChrisHegarty well, now I don't get that exception. Maybe I didn't apply your changes correctly locally. Testing some more :)

[ChrisHegarty]

@ChrisHegarty well, now I don't get that exception. Maybe I didn't apply your changes correctly locally. Testing some more :)

check that you have the change from benchmarks/build.gradle. (which contains a dependency on native)

[benwtrent]

OK, yes! that worked, but I am running into other things (likely other weirdness) when trying to access native methods directly.

Adding this to my benchmark (to access native method handle directly):

    final VectorSimilarityFunctions DISTANCE_FUNCS = NativeAccess.instance()
        .getVectorSimilarityFunctions()
        .orElseThrow(AssertionError::new);
    final MethodHandle INT4_BIT_DP = DISTANCE_FUNCS.int4BitDotProductHandle();

    private int int4BitDotProduct(MemorySegment a, MemorySegment b, long offset, MemorySegment scores, int count, int length) {
        try {
            return (int) INT4_BIT_DP.invokeExact(a, b, offset, scores, count, length);
        } catch (Throwable e) {
            if (e instanceof Error err) {
                throw err;
            } else if (e instanceof RuntimeException re) {
                throw re;
            } else {
                throw new RuntimeException(e);
            }
        }
    }

results in

java.lang.AssertionError
        at java.base/java.util.Optional.orElseThrow(Optional.java:403)
        at org.elasticsearch.benchmark.vector.OSQNativeScorerBenchmark.<init>(OSQNativeScorerBenchmark.java:58)
        at org.elasticsearch.benchmark.vector.jmh_generated.OSQNativeScorerBenchmark_jmhType_B1.<init>(OSQNativeScorerBenchmark_jmhType_B1.java:3)
        at org.elasticsearch.benchmark.vector.jmh_generated.OSQNativeScorerBenchmark_jmhType_B2.<init>(OSQNativeScorerBenchmark_jmhType_B2.java:3)
        at org.elasticsearch.benchmark.vector.jmh_generated.OSQNativeScorerBenchmark_jmhType_B3.<init>(OSQNativeScorerBenchmark_jmhType_B3.java:2)

So, likely I need to do some initialization step, or adjust my JVM version?

[benwtrent]

@ChrisHegarty I was able to figure it out. Apparently the native libraries didn't get built right locally, rebuilt and its all good :)