perf(spectrum): software-prefetch SA entries in the boundary-search and forward/tabled probe loops

prmi/src/index/mod.rs

@@ -214,6 +214,15 @@ impl LearnedIndex {
         self.sa.position(i)
     }
 
+    /// Software-prefetch SA entry `i` (position + key, one cache line) into L1.
+    /// Advisory hint used by the boundary-search probe loops to overlap the next
+    /// cold read with the current compare; never affects results. Out-of-range
+    /// `i` is ignored.
+    #[inline(always)]
+    pub(crate) fn prefetch_sa(&self, i: u64) {
+        self.sa.prefetch(i);
+    }
+
     /// Global maximum prediction error bound recorded in `.meta`.
     pub fn max_error_bound(&self) -> u64 {
         self.meta.rmi.max_error_bound

prmi/src/index/spectrum.rs

@@ -965,6 +965,7 @@ impl LearnedIndex {
             let mut a = lo;
             let mut b = hi;
             while a < b {
+                self.prefetch_bsearch(a, b);
                 let mid = a + (b - a) / 2;
                 let pos = self.sa_position_for(mid);
                 bump_probe();
@@ -983,6 +984,7 @@ impl LearnedIndex {
             let mut c = k;
             let mut d = hi;
             while c < d {
+                self.prefetch_bsearch(c, d);
                 let mid = c + (d - c) / 2;
                 let pos = self.sa_position_for(mid);
                 bump_probe();
@@ -1201,6 +1203,7 @@ impl LearnedIndex {
             let mut a = lo;
             let mut b = hi;
             while a < b {
+                self.prefetch_bsearch(a, b);
                 let mid = a + (b - a) / 2;
                 let pos = self.sa_position_for(mid);
                 bump_probe();
@@ -1218,6 +1221,7 @@ impl LearnedIndex {
             let mut c = kk;
             let mut d = hi;
             while c < d {
+                self.prefetch_bsearch(c, d);
                 let mid = c + (d - c) / 2;
                 let pos = self.sa_position_for(mid);
                 bump_probe();
@@ -2203,6 +2207,21 @@ impl LearnedIndex {
         }
     }
 
+    /// Prefetch the two SA entries a binary search over `[lo, hi)` would probe
+    /// next (the children of the current midpoint), so the next cold DRAM read is
+    /// already in flight while the current keyed compare runs. Advisory — it
+    /// never changes a search result, only its latency; out-of-range indices are
+    /// ignored by `prefetch_sa`. Called at the top of every SA binary-search loop.
+    #[inline(always)]
+    fn prefetch_bsearch(&self, lo: u64, hi: u64) {
+        if lo >= hi {
+            return;
+        }
+        let mid = lo + (hi - lo) / 2;
+        self.prefetch_sa(lo + (mid - lo) / 2); // next mid if we branch toward lo
+        self.prefetch_sa(mid + 1 + (hi - mid - 1) / 2); // next mid if we branch toward hi
+    }
+
     /// Find the boundary index in `[domain_lo, domain_hi)` — the first index `i`
     /// where the monotone predicate `go_right(i)` is `false` — seeded from the model
     /// window `[seed_lo, seed_hi)` with exponential expand-on-miss.
@@ -2274,8 +2293,15 @@ impl LearnedIndex {
             }
             break;
         }
-        // `[lo, hi)` now straddles the boundary; standard binary search.
+        // `[lo, hi)` now straddles the boundary; standard binary search. Every
+        // `go_right(mid)` probes SA entry `mid` (a cold DRAM read); prefetch the
+        // two possible next-probe entries so that read is already in flight while
+        // this iteration's keyed compare runs. Advisory only — it never changes
+        // the boundary returned, just the probe latency. (Every caller's
+        // predicate probes SA entry `mid`, so the prefetch targets are SA
+        // indices; out-of-range indices are ignored by `prefetch_sa`.)
         while lo < hi {
+            self.prefetch_bsearch(lo, hi);
             let mid = lo + (hi - lo) / 2;
             if go_right(mid) {
                 lo = mid + 1;

prmi/src/sidecar/sa_file.rs

@@ -359,6 +359,44 @@ impl SaFileReader {
         unpack_position(bytes)
     }
 
+    /// Issue a software prefetch hint for SA entry `i` into L1. Advisory and
+    /// side-effect-free: an out-of-range `i` is ignored, and on architectures
+    /// without a prefetch intrinsic it compiles to nothing. The 5-byte position
+    /// and (mode 2/3) 8-byte key share one entry — a single cache line — so one
+    /// prefetch warms both. Used by the binary-search probe loops to overlap the
+    /// next cold DRAM read with the current keyed compare; it never changes a
+    /// search result, only its latency.
+    #[inline(always)]
+    pub fn prefetch(&self, i: u64) {
+        if i >= self.num_entries {
+            return;
+        }
+        let off = SA_FILE_HEADER_BYTES + (i as usize) * self.bytes_per_entry;
+        // SAFETY: `i < num_entries` and the header validation checked the exact
+        // mapped length, so `off` (the entry start) is within `[0, data_len)`.
+        let addr = unsafe { self.data_ptr.add(off) };
+        #[cfg(target_arch = "x86_64")]
+        // SAFETY: `_mm_prefetch` is a pure hint — it touches no memory
+        // architecturally, never faults, and has no memory-safety effects.
+        unsafe {
+            core::arch::x86_64::_mm_prefetch::<{ core::arch::x86_64::_MM_HINT_T0 }>(
+                addr as *const i8,
+            );
+        }
+        #[cfg(target_arch = "aarch64")]
+        // SAFETY: `prfm` is a hint instruction; it never faults and has no
+        // memory-safety effects.
+        unsafe {
+            core::arch::asm!(
+                "prfm pldl1keep, [{0}]",
+                in(reg) addr,
+                options(nostack, preserves_flags),
+            );
+        }
+        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
+        let _ = addr;
+    }
+
     /// Return the stored 32-mer key at SA index `i`, if this file was built in
     /// mode 2 or mode 3 (both of which store the key). Returns `None` for mode 1
     /// and suffix-key-cache mode (where keys are absent or stored separately).