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        unsafe {
16            core::arch::x86_64::_mm_prefetch(ptr as *const i8, core::arch::x86_64::_MM_HINT_T0);
17        }
18    }
19
20    #[cfg(not(target_arch = "x86_64"))]
21    {
22        let _ = ptr; // Suppress unused variable warning
23    }
24}
25
26/// Prefetch memory for writing.
27///
28/// This function hints to the processor that the specified memory location
29/// will be written soon. On most architectures, this is equivalent to
30/// prefetch_read since writing also requires reading the cache line.
31#[inline(always)]
32pub fn prefetch_write<T>(ptr: *const T) {
33    #[cfg(target_arch = "x86_64")]
34    {
35        unsafe {
36            core::arch::x86_64::_mm_prefetch(ptr as *const i8, core::arch::x86_64::_MM_HINT_T0);
37        }
38    }
39
40    #[cfg(not(target_arch = "x86_64"))]
41    {
42        let _ = ptr; // Suppress unused variable warning
43    }
44}
45
46/// Prefetch a range of memory for reading.
47///
48/// This function prefetches multiple cache lines starting from the given
49/// pointer for the specified number of bytes.
50///
51/// # Arguments
52/// * `ptr` - Starting memory address
53/// * `bytes` - Number of bytes to prefetch
54pub fn prefetch_range_read<T>(ptr: *const T, bytes: usize) {
55    const CACHE_LINE_SIZE: usize = 64;
56    let start = ptr as usize;
57    // Saturating/checked address arithmetic: a range ending near `usize::MAX`
58    // must not overflow (which panics under `overflow-checks`). Prefetch hints
59    // are non-faulting, so the only requirement is that the loop terminates.
60    let end = start.saturating_add(bytes);
61
62    let mut addr = start & !(CACHE_LINE_SIZE - 1); // Align to cache line
63    while addr < end {
64        prefetch_read(addr as *const u8);
65        match addr.checked_add(CACHE_LINE_SIZE) {
66            Some(next) => addr = next,
67            None => break,
68        }
69    }
70}
71
72/// Prefetch a slice for reading.
73///
74/// This function prefetches all memory containing the given slice.
75///
76/// # Arguments
77/// * `slice` - The slice to prefetch
78pub fn prefetch_slice_read<T>(slice: &[T]) {
79    if !slice.is_empty() {
80        let bytes = core::mem::size_of_val(slice);
81        prefetch_range_read(slice.as_ptr(), bytes);
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn test_prefetch_slice() {
91        let data = vec![1, 2, 3, 4, 5];
92        // This should not panic
93        prefetch_slice_read(&data);
94
95        // Test empty slice
96        let empty: &[i32] = &[];
97        prefetch_slice_read(empty);
98    }
99
100    #[test]
101    fn test_prefetch_functions() {
102        let data = 42;
103        let ptr = &data as *const i32;
104
105        // These should not panic
106        prefetch_read(ptr);
107        prefetch_write(ptr);
108        prefetch_range_read(ptr, 4);
109    }
110}