Copying the implementations of dotProduct and l2Distance from Lucene's PanamaVectorUtilSupport. Lucene calls it "squaredDistance" instead of "l2Distance", but it's the same thing.
The authors of Lucene have made some really nice optimizations here, which I'm shamelessly copying into Elastiknn. I also considered just using their class and functions directly (things are private, package private, etc), but found it tricky to instantiate, so I just copied over the implementations.
The rough idea in this optimization is: when doing anything with Panama vectors, we have to iterate over the original vector in segments that match the size expected by the processor's SIMD implementation. But we can actually go faster if we unroll the loop by iterating four chunks at a time. Somehow the JVM figures out that we can compute these four chunks in parallel, but does not seem to figure this out if we're going one chunk at a time.
In doing this I went ahead and updated to from JDK 19 to 21 in the .tool-versions file and in the Github workflows. JDK 21 is used by the latest Elasticsearch. I was just behind in this repo.
I updated the JMH micro benchmarks, and it looks like the new implementation is indeed ~4x faster:
Related Issue
611 #617
Changes
Copying the implementations of dotProduct and l2Distance from Lucene's PanamaVectorUtilSupport. Lucene calls it "squaredDistance" instead of "l2Distance", but it's the same thing.
The authors of Lucene have made some really nice optimizations here, which I'm shamelessly copying into Elastiknn. I also considered just using their class and functions directly (things are private, package private, etc), but found it tricky to instantiate, so I just copied over the implementations.
The rough idea in this optimization is: when doing anything with Panama vectors, we have to iterate over the original vector in segments that match the size expected by the processor's SIMD implementation. But we can actually go faster if we unroll the loop by iterating four chunks at a time. Somehow the JVM figures out that we can compute these four chunks in parallel, but does not seem to figure this out if we're going one chunk at a time.
In doing this I went ahead and updated to from JDK 19 to 21 in the
.tool-versionsfile and in the Github workflows. JDK 21 is used by the latest Elasticsearch. I was just behind in this repo.I updated the JMH micro benchmarks, and it looks like the new implementation is indeed ~4x faster:
The ann-benchmarks also got faster:
There was some variability in the results. They went up as high as 96% recall at 211 qps in https://github.com/alexklibisz/elastiknn/pull/620/commits/f7d082fb7e813616f46bab787f8903c49909d204:
Testing and Validation
Standard CI and benchmarking (see above)