From 77bb6fa11467afcfe20d3d194c15a826fead1feb Mon Sep 17 00:00:00 2001 From: Patrick Fox Date: Sun, 25 Oct 2026 10:30:20 +0000 Subject: [PATCH 1/3] [diskann] Add aarch64 NeighborPriortyQueue.get_lower_bound --- diskann/src/neighbor/queue.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/diskann/src/neighbor/queue.rs b/diskann/src/neighbor/queue.rs index f4991574ec..bc84bbf803 100644 --- a/diskann/src/neighbor/queue.rs +++ b/diskann/src/neighbor/queue.rs @@ -8,6 +8,7 @@ use std::{ marker::PhantomData, }; +#[cfg(target_arch = "x86_64")] use diskann_wide::{SIMDMask, SIMDPartialOrd, SIMDVector}; use super::Neighbor; @@ -226,6 +227,7 @@ impl NeighborPriorityQueue { /// Get the lower bound of the neighbor - the index of the first neighbor /// that has a distance greater than or equal to the target neighbor /// PERFORMANCE: This function is performance critical - please benchmark any changes + #[cfg(target_arch = "x86_64")] fn get_lower_bound(&mut self, nbr: &Neighbor) -> usize { diskann_wide::alias!(f32s = f32x8); let target = f32s::splat(diskann_wide::ARCH, nbr.distance); @@ -279,6 +281,22 @@ impl NeighborPriorityQueue { self.size } + /// Get the lower bound of the neighbor - the index of the first neighbor + /// that has a distance greater than or equal to the target neighbor + /// PERFORMANCE: This function is performance critical - please benchmark any changes + #[cfg(target_arch = "aarch64")] + fn get_lower_bound(&mut self, nbr: &Neighbor) -> usize { + let target = nbr.distance; + + for i in 0..self.size { + if self.distances[i] >= target { + return i; + } + } + + self.size + } + /// Get the neighbor at index - SAFETY: index must be less than size fn get_unchecked(&self, index: usize) -> Neighbor { debug_assert!(index < self.size); From fc562377281454b1c50593428fb1757162a879ce Mon Sep 17 00:00:00 2001 From: pfoxARM Date: Mon, 26 Oct 2026 18:04:45 +0000 Subject: [PATCH 2/3] Potential fix for pull request finding 'Restore scalar fallback for get_lower_bound on non-x86_64 targets' Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- diskann/src/neighbor/queue.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann/src/neighbor/queue.rs b/diskann/src/neighbor/queue.rs index bc84bbf803..e21e865bcf 100644 --- a/diskann/src/neighbor/queue.rs +++ b/diskann/src/neighbor/queue.rs @@ -284,7 +284,7 @@ impl NeighborPriorityQueue { /// Get the lower bound of the neighbor - the index of the first neighbor /// that has a distance greater than or equal to the target neighbor /// PERFORMANCE: This function is performance critical - please benchmark any changes - #[cfg(target_arch = "aarch64")] + #[cfg(not(target_arch = "x86_64"))] fn get_lower_bound(&mut self, nbr: &Neighbor) -> usize { let target = nbr.distance; From dbb571a3c673952c49ccf2910132959206ac0a6e Mon Sep 17 00:00:00 2001 From: Patrick Fox Date: Wed, 28 Oct 2026 17:30:40 +0000 Subject: [PATCH 3/3] aarch64 get_lower_bound: Use iterator approach to avoid panic --- diskann/src/neighbor/queue.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diskann/src/neighbor/queue.rs b/diskann/src/neighbor/queue.rs index e21e865bcf..4559fd639e 100644 --- a/diskann/src/neighbor/queue.rs +++ b/diskann/src/neighbor/queue.rs @@ -288,8 +288,8 @@ impl NeighborPriorityQueue { fn get_lower_bound(&mut self, nbr: &Neighbor) -> usize { let target = nbr.distance; - for i in 0..self.size { - if self.distances[i] >= target { + for (i, d) in self.distances.iter().take(self.size).enumerate() { + if *d >= target { return i; } }