Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add aarch64 neon simd support for distance calculations #127

Open
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

knudtty
Copy link

@knudtty knudtty commented Aug 23, 2024

Closes #126

Added simd distance (cosine and l2) calculations for aarch64 neon intrinsics. Tested on Macbook M1 Max, AWS EC2 t4g.medium, and even a Raspberry Pi 4. Benchmarks showed anywhere from 3-10x performance improvements depending on the machine. Seeing that many cloud managed db's use ARM (AWS RDS is all Gravitron3 chips IIRC), Neon SIMD would go a long ways.

Added a test between this new function and the unoptimized version to ensure distance calculations are equivalent. All tests pass.

@knudtty knudtty requested a review from a team as a code owner August 23, 2024 15:53
Copy link
Collaborator

@cevian cevian left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the great PR a few asks:

  1. can you please add the equivalent of https://github.com/timescale/pgvectorscale/blob/main/pgvectorscale/src/access_method/distance.rs#L12 for neon

  2. optional: can you add something to https://github.com/timescale/pgvectorscale/blob/main/pgvectorscale/benches/distance.rs (or maybe a different bench)

  3. please address the inline comments.

Thanks again

pgvectorscale/src/access_method/distance.rs Outdated Show resolved Hide resolved
pgvectorscale/src/access_method/distance_aarch64.rs Outdated Show resolved Hide resolved
@knudtty
Copy link
Author

knudtty commented Aug 28, 2024

  1. can you please add the equivalent of https://github.com/timescale/pgvectorscale/blob/main/pgvectorscale/src/access_method/distance.rs#L12 for neon

Sure

  1. optional: can you add something to https://github.com/timescale/pgvectorscale/blob/main/pgvectorscale/benches/distance.rs (or maybe a different bench)

The neon distance calcs are being run as part of benchmark_distance since their inclusion in distance_l2 and distance_cosine, were you thinking to also add the unoptimized benchmarks for comparison next to it like this?
Screenshot 2024-08-28 at 6 18 57 PM

@tjgreen42
Copy link
Contributor

Code looks great (thanks @knudtty!). @cevian, can we get this in? Looks like all changes you requested were applied.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Support neon simd for distance calculations in aarch64
4 participants