-
Notifications
You must be signed in to change notification settings - Fork 147
Optimizes reducelanes in diversityCalculation of PQVectors, for Euclidean function#514
Optimizes reducelanes in diversityCalculation of PQVectors, for Euclidean function #514shyla226 wants to merge 2 commits intodatastax:main from
Conversation
a991b33 to
dd1c332
Compare
MarkWolters
commented
Oct 3, 2025
This PR does successfully remove the calls to reduceLanes from the inside of the loop iterating over the subspaces the Euclidean case. I do have a concern about applying this pattern to one specific case but leaving the handling of other cases as is, leaving the code with 2 different implementations for the same problem. Some questions:
-
Why is the optimization only implemented for the Euclidean case? In PanamaVectorUtilSupport, where the call to reduceLanes is made from the squareDistance methods which are called in a loop for the subspace count, likewise reduceLanes is called from the dotProduct methods. Why only fix one?
-
Why is the optimization only applied to diversityFunctionFor(...) and not scoreFunctionFor(...) which uses the same pattern of calling squareL2Distance(...) inside of a loop for each subspace?
shyla226
commented
Oct 6, 2025
@MarkWolters, thank you very much for the feedback. I can make the changes to dotProduct methods and to function calls in scoreFunctionFor()
MarkWolters
commented
Oct 10, 2025
@shyla226 please be aware that in order to pass the automated GitHub action regression test any pull request must come from a branch inside the datastax/jvector repository and not a fork of the repository. So while I am willing to review the PR from a fork and provide commentary, in order for it to pass the requirements for merging it will have to be a branch and not a fork.
shyla226
commented
Oct 10, 2025
@MarkWolters, Sounds good! I will create a branch
MarkWolters
commented
Oct 10, 2025
@shyla226 I don't think you'll have sufficient permissions to push a new branch to the repository. For now you can keep the changes on a fork and they can be reviewed there while we come up with a solution to this issue.
shyla226
commented
Oct 27, 2025
@MarkWolters , I have added new functions to include dotproduct & cosine for diversityFunction.
I have also added overloaded function for scoreFunction, for all the 3 cases
MarkWolters
commented
Oct 28, 2025
Thanks @shyla226, I will review the updated code asap
shyla226
commented
Oct 28, 2025
Thank you very much @MarkWolters . With the changes, when M=64, dot product shows ~30% reduction in latency & COSINE shows ~43% reduction
shyla226
commented
Oct 28, 2025
@MarkWolters sorry, accidentally closed it! I will update the test setup used in the description above.
shyla226
commented
Oct 30, 2025
@MarkWolters, I have addressed the issue with the 2 failed checks and tested the code changes on Windows server
MarkWolters
commented
Feb 2, 2026
@shyla226 apologies for not addressing this sooner. I have created a branch named reduce_lanes based off of the latest main. Please rebase your change on current main and open a pull request from your fork to the reduce_lanes branch. From the Github UI you should be able to select Contribute --> Open Pull Request and in the resulting screen select the branch you wish to merge to. From there we can run the internal github actions and regression tests and get your changes merged.
shyla226
commented
Feb 3, 2026
@MarkWolters Thank you very much. I will rebase the changes.
c2a8068 to
8e85edb
Compare
Uh oh!
There was an error while loading. Please reload this page.
This pull request introduces improvement to Euclidean similarity function in
PQVectors.diversityFunctionFor. From flamegraph, it is observed that considerable amount of time is spent injdk/incubator/vector/FloatVector.reducelanesTemplate. This is mainly becauseFloatVector.reducelanes()is expensive and it is being called inside aforloop (viaVectorUtil.squareL2Distance). Modification in this pull request moves call toreduceLanes()outside theforloop.Change proposed here was tested with the benchmark,
PQDistanceCalculationBenchmark.diversityCalculationWith this benchmark, ~18% reduction in time was observed when M=64 and ~22% when M=192.
Code modifications:
pqDiversityEuclideaninVectorUtilSupportand its corresponding implementationsforloop inPQVectors.diversityFunctionForand moved it intopqDiversityEuclideanFloatVector.reducelanes()outside theforloopTest setup:
Jvector version : main branch (as of 2025年08月28日)
JDK version : openjdk version "24.0.2" 2025年07月15日
Platform : INTEL(R) XEON(R) PLATINUM 8592+
Benchmark : PQDistanceCalculationBenchmark.diversityCalculation
New changes
I have modified the code to include dot product & cosine functions and implemented similar changes for scoreFunctionFor.
With the changes applied to scoreFunctionFor, when M=64: dot product shows ~30% reduction in latency & cosine shows ~43% reduction.
I can add data points for other subspace counts, if required.
Code modifications:
pqScoreEuclidean, pqScoreDotProduct, pqScoreCosineinVectorUtilSupportand its corresponding implementations fordiversityFunctionForscoreFunctionForforloop inPQVectors.diversityFunctionForandPQVectors.scoreFunctionForand moved them into respective functions inPanamaVectorUtilSupportFloatVector.reducelanes()outside theforloopMutablePQVectorsto test thisTest setup:
Jvector version : main branch (as of 2025年10月22日)
JDK version : openjdk version "25.0.1" 2025年10月21日
Platform : Intel(R) Xeon(R) 6979P
Benchmark :
PQDistanceCalculationMutableVectorBenchmark(Added this by replicatingPQDistanceCalculationBenchmarkto measure performance for MutablePQVectors)