Coverage Report

Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/home/noah/src/trueno/src/brick/memory.rs
Line
Count
Source
1
//! Memory Management Primitives
2
//!
3
//! Cache line alignment, direct I/O buffers, memory advice, and prefetch utilities.
4
5
use crate::error::TruenoError;
6
7
// ----------------------------------------------------------------------------
8
// LCP-06: Cache Line Padding
9
// ----------------------------------------------------------------------------
10
11
/// Cache line size (64 bytes on most modern CPUs).
12
pub const CACHE_LINE_SIZE: usize = 64;
13
14
/// Number of f32 values per cache line.
15
pub const CACHE_LINE_SIZE_F32: usize = CACHE_LINE_SIZE / std::mem::size_of::<f32>();
16
17
/// Cache-line aligned wrapper to prevent false sharing.
18
///
19
/// # Example
20
/// ```rust
21
/// use trueno::brick::CacheAligned;
22
/// use std::sync::atomic::AtomicU64;
23
///
24
/// let aligned: CacheAligned<AtomicU64> = CacheAligned::new(AtomicU64::new(0));
25
/// assert_eq!(std::mem::align_of_val(&aligned), 64);
26
/// ```
27
#[repr(align(64))]
28
#[derive(Debug)]
29
pub struct CacheAligned<T>(pub T);
30
31
impl<T> CacheAligned<T> {
32
    /// Create a new cache-aligned value.
33
0
    pub const fn new(value: T) -> Self {
34
0
        Self(value)
35
0
    }
36
37
    /// Get a reference to the inner value.
38
0
    pub fn get(&self) -> &T {
39
0
        &self.0
40
0
    }
41
42
    /// Get a mutable reference to the inner value.
43
0
    pub fn get_mut(&mut self) -> &mut T {
44
0
        &mut self.0
45
0
    }
46
47
    /// Consume the wrapper and return the inner value.
48
0
    pub fn into_inner(self) -> T {
49
0
        self.0
50
0
    }
51
}
52
53
impl<T: Default> Default for CacheAligned<T> {
54
0
    fn default() -> Self {
55
0
        Self(T::default())
56
0
    }
57
}
58
59
impl<T: Clone> Clone for CacheAligned<T> {
60
0
    fn clone(&self) -> Self {
61
0
        Self(self.0.clone())
62
0
    }
63
}
64
65
// ----------------------------------------------------------------------------
66
// LCP-02: Direct I/O Alignment
67
// ----------------------------------------------------------------------------
68
69
/// Memory alignment for direct I/O (4KB page aligned).
70
pub const DIRECT_IO_ALIGNMENT: usize = 4096;
71
72
/// Check if a pointer is aligned for direct I/O.
73
#[must_use]
74
0
pub fn is_direct_io_aligned<T>(ptr: *const T) -> bool {
75
0
    (ptr as usize).is_multiple_of(DIRECT_IO_ALIGNMENT)
76
0
}
77
78
/// Aligned buffer for direct I/O operations.
79
#[cfg(not(target_arch = "wasm32"))]
80
pub struct AlignedBuffer {
81
    ptr: *mut u8,
82
    len: usize,
83
    layout: std::alloc::Layout,
84
}
85
86
#[cfg(not(target_arch = "wasm32"))]
87
impl AlignedBuffer {
88
    /// Allocate a new aligned buffer.
89
    ///
90
    /// # Errors
91
    /// Returns an error if allocation fails.
92
0
    pub fn new(size: usize) -> Result<Self, TruenoError> {
93
        use std::alloc::{alloc_zeroed, Layout};
94
95
0
        let layout = Layout::from_size_align(size, DIRECT_IO_ALIGNMENT)
96
0
            .map_err(|_| TruenoError::InvalidInput("invalid alignment".into()))?;
97
98
0
        let ptr = unsafe { alloc_zeroed(layout) };
99
0
        if ptr.is_null() {
100
0
            return Err(TruenoError::InvalidInput("allocation failed".into()));
101
0
        }
102
103
0
        Ok(Self { ptr, len: size, layout })
104
0
    }
105
106
    /// Get the buffer as a slice.
107
0
    pub fn as_slice(&self) -> &[u8] {
108
0
        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
109
0
    }
110
111
    /// Get the buffer as a mutable slice.
112
0
    pub fn as_mut_slice(&mut self) -> &mut [u8] {
113
0
        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
114
0
    }
115
116
    /// Get the raw pointer.
117
0
    pub fn as_ptr(&self) -> *const u8 {
118
0
        self.ptr
119
0
    }
120
121
    /// Get the mutable raw pointer.
122
0
    pub fn as_mut_ptr(&mut self) -> *mut u8 {
123
0
        self.ptr
124
0
    }
125
126
    /// Get the buffer length.
127
0
    pub fn len(&self) -> usize {
128
0
        self.len
129
0
    }
130
131
    /// Check if the buffer is empty.
132
0
    pub fn is_empty(&self) -> bool {
133
0
        self.len == 0
134
0
    }
135
}
136
137
#[cfg(not(target_arch = "wasm32"))]
138
impl Drop for AlignedBuffer {
139
0
    fn drop(&mut self) {
140
0
        unsafe {
141
0
            std::alloc::dealloc(self.ptr, self.layout);
142
0
        }
143
0
    }
144
}
145
146
#[cfg(not(target_arch = "wasm32"))]
147
unsafe impl Send for AlignedBuffer {}
148
149
#[cfg(not(target_arch = "wasm32"))]
150
unsafe impl Sync for AlignedBuffer {}
151
152
// ----------------------------------------------------------------------------
153
// LCP-03: Memory Advice (madvise patterns)
154
// ----------------------------------------------------------------------------
155
156
/// Memory advice for mmap regions.
157
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
158
pub enum MemoryAdvice {
159
    /// Sequential access (enable readahead)
160
    Sequential,
161
    /// Random access (disable readahead)
162
    Random,
163
    /// Will need soon (prefetch)
164
    WillNeed,
165
    /// Don't need (can be paged out)
166
    DontNeed,
167
}
168
169
// Linux madvise constants (from linux/mman.h)
170
#[cfg(target_os = "linux")]
171
const MADV_SEQUENTIAL: i32 = 2;
172
#[cfg(target_os = "linux")]
173
const MADV_RANDOM: i32 = 1;
174
#[cfg(target_os = "linux")]
175
const MADV_WILLNEED: i32 = 3;
176
#[cfg(target_os = "linux")]
177
const MADV_DONTNEED: i32 = 4;
178
179
/// Apply memory advice to a region (Linux only).
180
///
181
/// # Safety
182
/// The pointer must be valid and the length must not exceed the mapped region.
183
#[cfg(target_os = "linux")]
184
0
pub unsafe fn madvise_region(addr: *mut u8, len: usize, advice: MemoryAdvice) -> std::io::Result<()> {
185
    // madvise syscall number is 28 on x86_64
186
    #[cfg(target_arch = "x86_64")]
187
    const SYS_MADVISE: i64 = 28;
188
    #[cfg(target_arch = "aarch64")]
189
    const SYS_MADVISE: i64 = 233;
190
191
0
    let advice_flag: i32 = match advice {
192
0
        MemoryAdvice::Sequential => MADV_SEQUENTIAL,
193
0
        MemoryAdvice::Random => MADV_RANDOM,
194
0
        MemoryAdvice::WillNeed => MADV_WILLNEED,
195
0
        MemoryAdvice::DontNeed => MADV_DONTNEED,
196
    };
197
198
    let ret: i64;
199
    #[cfg(target_arch = "x86_64")]
200
0
    {
201
0
        core::arch::asm!(
202
0
            "syscall",
203
0
            inout("rax") SYS_MADVISE => ret,
204
0
            in("rdi") addr as usize,
205
0
            in("rsi") len,
206
0
            in("rdx") advice_flag as i64,
207
0
            out("rcx") _,
208
0
            out("r11") _,
209
0
            options(nostack)
210
0
        );
211
0
    }
212
    #[cfg(target_arch = "aarch64")]
213
    {
214
        core::arch::asm!(
215
            "svc 0",
216
            inout("x8") SYS_MADVISE => _,
217
            inout("x0") addr as usize => ret,
218
            in("x1") len,
219
            in("x2") advice_flag as i64,
220
            options(nostack)
221
        );
222
    }
223
224
0
    if ret < 0 {
225
0
        return Err(std::io::Error::from_raw_os_error(-ret as i32));
226
0
    }
227
228
0
    Ok(())
229
0
}
230
231
/// Stub for non-Linux platforms.
232
#[cfg(not(target_os = "linux"))]
233
pub unsafe fn madvise_region(_addr: *mut u8, _len: usize, _advice: MemoryAdvice) -> std::io::Result<()> {
234
    Ok(()) // No-op on non-Linux
235
}
236
237
/// Apply dual-level prefetch strategy (WILLNEED + RANDOM).
238
///
239
/// This is the llama.cpp pattern for model loading:
240
/// 1. MADV_WILLNEED: Tell kernel to prefetch the data
241
/// 2. MADV_RANDOM: Disable readahead (model access is random)
242
///
243
/// # Safety
244
/// The pointer must be valid and the length must not exceed the mapped region.
245
#[cfg(target_os = "linux")]
246
0
pub unsafe fn prefetch_for_inference(addr: *mut u8, len: usize) -> std::io::Result<()> {
247
    // First: tell kernel we'll need this data
248
0
    madvise_region(addr, len, MemoryAdvice::WillNeed)?;
249
    // Second: hint random access pattern (disables readahead waste)
250
0
    madvise_region(addr, len, MemoryAdvice::Random)?;
251
0
    Ok(())
252
0
}
253
254
/// Stub for non-Linux platforms.
255
#[cfg(not(target_os = "linux"))]
256
pub unsafe fn prefetch_for_inference(_addr: *mut u8, _len: usize) -> std::io::Result<()> {
257
    Ok(()) // No-op on non-Linux
258
}
259
260
// ----------------------------------------------------------------------------
261
// LCP-11: Prefetch with Locality Hints
262
// ----------------------------------------------------------------------------
263
264
/// Prefetch locality hints.
265
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
266
pub enum PrefetchLocality {
267
    /// No temporal locality (use once, don't pollute cache)
268
    None = 0,
269
    /// Low temporal locality (use a few times)
270
    Low = 1,
271
    /// Moderate temporal locality
272
    Moderate = 2,
273
    /// High temporal locality (keep in all cache levels)
274
    High = 3,
275
}
276
277
/// Prefetch data into cache.
278
///
279
/// # Safety
280
/// The pointer must be valid for reading.
281
#[inline]
282
#[cfg(target_arch = "x86_64")]
283
0
pub unsafe fn prefetch_ptr<T>(ptr: *const T, locality: PrefetchLocality) {
284
    use core::arch::x86_64::*;
285
0
    match locality {
286
0
        PrefetchLocality::None => _mm_prefetch(ptr as *const i8, _MM_HINT_NTA),
287
0
        PrefetchLocality::Low => _mm_prefetch(ptr as *const i8, _MM_HINT_T2),
288
0
        PrefetchLocality::Moderate => _mm_prefetch(ptr as *const i8, _MM_HINT_T1),
289
0
        PrefetchLocality::High => _mm_prefetch(ptr as *const i8, _MM_HINT_T0),
290
    }
291
0
}
292
293
/// Prefetch data into cache (ARM64).
294
#[inline]
295
#[cfg(target_arch = "aarch64")]
296
pub unsafe fn prefetch_ptr<T>(ptr: *const T, _locality: PrefetchLocality) {
297
    // ARM prefetch (PRFM instruction) - locality hints are limited
298
    core::arch::asm!(
299
        "prfm pldl1keep, [{ptr}]",
300
        ptr = in(reg) ptr,
301
        options(nostack, preserves_flags)
302
    );
303
}
304
305
/// Fallback for other architectures.
306
#[inline]
307
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
308
pub unsafe fn prefetch_ptr<T>(_ptr: *const T, _locality: PrefetchLocality) {
309
    // No-op on unsupported architectures
310
}
311
312
/// Prefetch a slice of data.
313
///
314
/// Prefetches each cache line in the slice.
315
#[inline]
316
0
pub fn prefetch_slice<T>(slice: &[T], locality: PrefetchLocality) {
317
0
    let ptr = slice.as_ptr() as *const u8;
318
0
    let len = std::mem::size_of_val(slice);
319
320
0
    for offset in (0..len).step_by(CACHE_LINE_SIZE) {
321
0
        unsafe {
322
0
            prefetch_ptr(ptr.add(offset), locality);
323
0
        }
324
    }
325
0
}
326
327
#[cfg(test)]
328
mod tests {
329
    use super::*;
330
331
    #[test]
332
    fn test_cache_aligned_alignment() {
333
        let aligned: CacheAligned<u64> = CacheAligned::new(42);
334
        assert_eq!(std::mem::align_of_val(&aligned), 64);
335
    }
336
337
    #[test]
338
    fn test_cache_aligned_value() {
339
        let aligned = CacheAligned::new(42u64);
340
        assert_eq!(*aligned.get(), 42);
341
    }
342
343
    #[test]
344
    fn test_cache_aligned_get_mut() {
345
        let mut aligned = CacheAligned::new(42u64);
346
        *aligned.get_mut() = 100;
347
        assert_eq!(*aligned.get(), 100);
348
    }
349
350
    #[test]
351
    fn test_cache_aligned_into_inner() {
352
        let aligned = CacheAligned::new(42u64);
353
        assert_eq!(aligned.into_inner(), 42);
354
    }
355
356
    #[test]
357
    fn test_cache_aligned_default() {
358
        let aligned: CacheAligned<u64> = CacheAligned::default();
359
        assert_eq!(*aligned.get(), 0);
360
    }
361
362
    #[test]
363
    fn test_cache_aligned_clone() {
364
        let aligned = CacheAligned::new(42u64);
365
        let cloned = aligned.clone();
366
        assert_eq!(*cloned.get(), 42);
367
    }
368
369
    #[test]
370
    fn test_cache_line_size_f32() {
371
        assert_eq!(CACHE_LINE_SIZE_F32, 16); // 64 / 4 = 16
372
    }
373
374
    #[test]
375
    fn test_direct_io_alignment() {
376
        assert_eq!(DIRECT_IO_ALIGNMENT, 4096);
377
    }
378
379
    #[test]
380
    fn test_is_direct_io_aligned() {
381
        let aligned_addr: usize = 4096 * 10;
382
        let unaligned_addr: usize = 4096 * 10 + 1;
383
384
        assert!(is_direct_io_aligned(aligned_addr as *const u8));
385
        assert!(!is_direct_io_aligned(unaligned_addr as *const u8));
386
    }
387
388
    #[cfg(not(target_arch = "wasm32"))]
389
    #[test]
390
    fn test_aligned_buffer_creation() {
391
        let buffer = AlignedBuffer::new(4096).unwrap();
392
        assert_eq!(buffer.len(), 4096);
393
        assert!(!buffer.is_empty());
394
    }
395
396
    #[cfg(not(target_arch = "wasm32"))]
397
    #[test]
398
    fn test_aligned_buffer_zeroed() {
399
        let buffer = AlignedBuffer::new(1024).unwrap();
400
        let slice = buffer.as_slice();
401
        assert!(slice.iter().all(|&b| b == 0));
402
    }
403
404
    #[cfg(not(target_arch = "wasm32"))]
405
    #[test]
406
    fn test_aligned_buffer_write() {
407
        let mut buffer = AlignedBuffer::new(1024).unwrap();
408
        buffer.as_mut_slice()[0] = 42;
409
        assert_eq!(buffer.as_slice()[0], 42);
410
    }
411
412
    #[test]
413
    fn test_memory_advice_eq() {
414
        assert_eq!(MemoryAdvice::Sequential, MemoryAdvice::Sequential);
415
        assert_ne!(MemoryAdvice::Sequential, MemoryAdvice::Random);
416
    }
417
418
    #[test]
419
    fn test_prefetch_locality_values() {
420
        assert_eq!(PrefetchLocality::None as u8, 0);
421
        assert_eq!(PrefetchLocality::Low as u8, 1);
422
        assert_eq!(PrefetchLocality::Moderate as u8, 2);
423
        assert_eq!(PrefetchLocality::High as u8, 3);
424
    }
425
426
    #[test]
427
    fn test_prefetch_slice_empty() {
428
        let empty: &[f32] = &[];
429
        prefetch_slice(empty, PrefetchLocality::High);
430
        // Should not panic
431
    }
432
433
    #[test]
434
    fn test_prefetch_slice_small() {
435
        let data = [1.0f32; 8];
436
        prefetch_slice(&data, PrefetchLocality::High);
437
        // Should not panic
438
    }
439
440
    #[test]
441
    fn test_madvise_region_stub() {
442
        // On non-Linux, this is a no-op
443
        unsafe {
444
            let mut data = [0u8; 4096];
445
            let _result = madvise_region(data.as_mut_ptr(), data.len(), MemoryAdvice::WillNeed);
446
            #[cfg(not(target_os = "linux"))]
447
            assert!(_result.is_ok());
448
        }
449
    }
450
451
    #[test]
452
    fn test_prefetch_for_inference_stub() {
453
        unsafe {
454
            let mut data = [0u8; 4096];
455
            let _result = prefetch_for_inference(data.as_mut_ptr(), data.len());
456
            #[cfg(not(target_os = "linux"))]
457
            assert!(_result.is_ok());
458
        }
459
    }
460
}