Skip to main content

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}