moirai_utils/memory.rs
1//! Memory utilities for cache optimization and prefetching.
2//!
3//! This module provides utilities for optimizing memory access patterns,
4//! including cache prefetching and aligned memory allocation.
5
6/// Prefetch memory for reading.
7///
8/// This function hints to the processor that the specified memory location
9/// will be read soon, allowing it to preload the data into cache.
10/// On architectures that don't support prefetching, this is a no-op.
11#[inline(always)]
12pub fn prefetch_read<T>(ptr: *const T) {
13 #[cfg(target_arch = "x86_64")]
14 {
15 // SAFETY: prefetch hints are non-faulting and advisory, so `ptr`
16 // need not be dereferenceable; SSE is baseline on x86_64 and the
17 // cast carries only the address.
18 unsafe {
19 core::arch::x86_64::_mm_prefetch(ptr as *const i8, core::arch::x86_64::_MM_HINT_T0);
20 }
21 }
22
23 #[cfg(not(target_arch = "x86_64"))]
24 {
25 let _ = ptr; // Suppress unused variable warning
26 }
27}
28
29/// Prefetch memory for writing.
30///
31/// This function hints to the processor that the specified memory location
32/// will be written soon. On most architectures, this is equivalent to
33/// prefetch_read since writing also requires reading the cache line.
34#[inline(always)]
35pub fn prefetch_write<T>(ptr: *const T) {
36 #[cfg(target_arch = "x86_64")]
37 {
38 // SAFETY: prefetch hints are non-faulting and advisory, so `ptr`
39 // need not be dereferenceable; SSE is baseline on x86_64 and the
40 // cast carries only the address.
41 unsafe {
42 core::arch::x86_64::_mm_prefetch(ptr as *const i8, core::arch::x86_64::_MM_HINT_T0);
43 }
44 }
45
46 #[cfg(not(target_arch = "x86_64"))]
47 {
48 let _ = ptr; // Suppress unused variable warning
49 }
50}
51
52/// Prefetch a range of memory for reading.
53///
54/// This function prefetches multiple cache lines starting from the given
55/// pointer for the specified number of bytes.
56///
57/// # Arguments
58/// * `ptr` - Starting memory address
59/// * `bytes` - Number of bytes to prefetch
60pub fn prefetch_range_read<T>(ptr: *const T, bytes: usize) {
61 // Transfer granularity, not the false-sharing separation: one prefetch
62 // hint covers one line, so striding by the (larger) interference size
63 // would skip every second line in the range.
64 use crate::cache::CACHE_LINE_SIZE;
65
66 let start = ptr as usize;
67 // Saturating/checked address arithmetic: a range ending near `usize::MAX`
68 // must not overflow (which panics under `overflow-checks`). Prefetch hints
69 // are non-faulting, so the only requirement is that the loop terminates.
70 let end = start.saturating_add(bytes);
71
72 let mut addr = start & !(CACHE_LINE_SIZE - 1); // Align to cache line
73 while addr < end {
74 prefetch_read(addr as *const u8);
75 match addr.checked_add(CACHE_LINE_SIZE) {
76 Some(next) => addr = next,
77 None => break,
78 }
79 }
80}
81
82/// Prefetch a slice for reading.
83///
84/// This function prefetches all memory containing the given slice.
85///
86/// # Arguments
87/// * `slice` - The slice to prefetch
88pub fn prefetch_slice_read<T>(slice: &[T]) {
89 if !slice.is_empty() {
90 let bytes = core::mem::size_of_val(slice);
91 prefetch_range_read(slice.as_ptr(), bytes);
92 }
93}
94
95#[cfg(test)]
96mod tests {
97 use super::*;
98
99 #[test]
100 fn test_prefetch_slice() {
101 let data = vec![1, 2, 3, 4, 5];
102 // This should not panic
103 prefetch_slice_read(&data);
104
105 // Test empty slice
106 let empty: &[i32] = &[];
107 prefetch_slice_read(empty);
108 }
109
110 #[test]
111 fn test_prefetch_functions() {
112 let data = 42;
113 let ptr = &data as *const i32;
114
115 // These should not panic
116 prefetch_read(ptr);
117 prefetch_write(ptr);
118 prefetch_range_read(ptr, 4);
119 }
120}