diff --git a/prmi/src/index/mod.rs b/prmi/src/index/mod.rs index 1004f80..893594e 100644 --- a/prmi/src/index/mod.rs +++ b/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 diff --git a/prmi/src/index/spectrum.rs b/prmi/src/index/spectrum.rs index ccdb700..2bd3af6 100644 --- a/prmi/src/index/spectrum.rs +++ b/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; diff --git a/prmi/src/sidecar/sa_file.rs b/prmi/src/sidecar/sa_file.rs index f5756bf..8fe924f 100644 --- a/prmi/src/sidecar/sa_file.rs +++ b/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).