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}