[ENH] Use AVX in distance calculations

[jairad26]

Description of changes

This PR updates the distance crate to prioritize AVX over SSE, and forces builds with AVX flags when ENABLE_AVX512 is set in the Dockerfile

Test plan

How are these changes tested?

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

Migration plan

Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?

Observability plan

What is the plan to instrument and monitor this change?

Documentation Changes

Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the docs section?

[jairad26]

• [ENH] Use AVX in distance calculations #5258 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[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)

[propel-code-bot]

Prioritize AVX Over SSE in Distance Calculations for Rust Distance Crate

This pull request reorders the SIMD feature detection and function selection logic in the Rust distance calculation module to give higher priority to AVX/FMA over SSE for x86/x86_64 architectures. The changes ensure that, when available, AVX and FMA instructions are used for distance computations (euclidean, cosine, and inner product), which can offer higher performance for platforms that support these instruction sets. Additionally, the author notes that builds will be forced to use AVX if ENABLE_AVX512 is set in the Dockerfile, though Dockerfile changes are not present in this diff.

Key Changes

• Reordered feature detection to select AVX/FMA implementations before SSE for relevant architectures.
• All distance calculation branches (euclidean, cosine, inner product) updated to use AVX/FMA when possible, with SSE as fallback.
• Maintained NEON branch for aarch64 unchanged.
• No significant algorithmic or interface changes; the focus is on CPU feature prioritization.

Affected Areas

• rust/distance/src/types.rs (distance calculation logic for all supported DistanceFunctions)

This summary was automatically generated by @propel-code-bot

[propel-code-bot]

[Documentation]

The reordering of instruction set priority from SSE-first to AVX-first looks correct, but consider documenting this prioritization decision. When both SSE and AVX are available at runtime, AVX will now be chosen, which should provide better performance. However, this means the fallback path (SSE → scalar) is now longer when AVX is not available.

Consider adding a comment explaining the prioritization rationale:

// Check instruction sets in order of performance: AVX (best) → SSE → scalar (fallback)

Also verify that the runtime feature detection order aligns with the compile-time feature requirements - all AVX-capable processors should also support SSE, so this ordering should be safe.