[ENH] SIMD vectorization for distance metrics

[sanketkedia]

Description of changes

Adds SIMD vectorization for euclidean, cosine and inner product for x86, x86_64 and arm. Instruction sets whose support has been added are SSE, AVX and NEON.

Test plan

• [+] Tests pass locally with pytest for python, yarn test for js, cargo test for rust

Documentation Changes

No

[github-actions]

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

• Can you think of any use case in which the code does not behave as intended? Have they been tested?
• Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
• If appropriate, are there adequate property based tests?
• If appropriate, are there adequate unit tests?
• Should any logging, debugging, tracing information be added or removed?
• Are error messages user-friendly?
• Have all documentation changes needed been made?
• Have all non-obvious changes been commented?

System Compatibility

• Are there any potential impacts on other parts of the system or backward compatibility?
• Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

• Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

[HammadB]

learning: what does this do (harness)

[sanketkedia]

harness disables libtest benchmarking since I am using the criterion crate

[HammadB]

out of curiosity what alternatives did we evaluate?

[sanketkedia]

Good question. By default, we have libtest using which we can do basic benchmarking. My main motive behind using criterion was to enable sophisticated benchmarking for the future. In particular criterion provides:

1. Statistics: Statistical analysis detects if, and by how much, performance has changed since the last benchmark run
2. Charts: Uses gnuplot to generate detailed graphs of benchmark results

These would be useful in future to run these benchmarks in our CI/CD pipeline or nightly/weekly (some cadence) to detect perf regressions for e.g.

[HammadB]

Awesome thank you

[HammadB]

should we add tests to validate that the simd impls (whatever can run on the target machine) match the base impl?

[sanketkedia]

I think there is a test already in types.rs - test_distance_function_l2sqr() which I was relying on. If that regresses then there is some bug in the SIMD impls. That test validates inner product and l2 norm. Since, cosine is the same as inner product for us, I am guessing that is not needed

[HammadB]

Ok great good point

[codetheweb]

this is super cool, wonder if it'd be worth using an existing vectorized implementation like this one instead of implementing it ourselves?

[sanketkedia]

@codetheweb the main issue with using a third-party is that our distance functions have slightly different definition. For e.g. cosine similarity assumes the vectors are normalized

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
chroma	❌ Failed (Inspect)			May 1, 2024 8:21pm

[sanketkedia]

Verified that there is about 10x perf improvement for both arm and x86_64. Going to merge this unless anyone else has any objections.