feat: add aarch64 neon simd support for distance calculations

timescale / pgvectorscale

A complement to pgvector for high performance, cost efficient vector search on large workloads.

PostgreSQL License

1.37k stars 58 forks source link

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

Closed knudtty closed 21 hours ago

knudtty commented 3 months ago

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 commented 2 months ago

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

Sure

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?

tjgreen42 commented 4 days ago

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