dynamo_memory/
system.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! System memory storage backed by malloc.
5
6use super::{MemoryDescription, Result, StorageError, StorageKind, actions, nixl::NixlDescriptor};
7use std::any::Any;
8use std::ptr::NonNull;
9
10/// System memory allocated via malloc.
11#[derive(Debug)]
12pub struct SystemStorage {
13    ptr: NonNull<u8>,
14    len: usize,
15}
16
17unsafe impl Send for SystemStorage {}
18unsafe impl Sync for SystemStorage {}
19
20impl SystemStorage {
21    /// Allocate new system memory of the given size.
22    pub fn new(len: usize) -> Result<Self> {
23        if len == 0 {
24            return Err(StorageError::AllocationFailed(
25                "zero-sized allocations are not supported".into(),
26            ));
27        }
28
29        let mut ptr: *mut libc::c_void = std::ptr::null_mut();
30
31        // We need 4KB alignment here for NIXL disk transfers to work.
32        // The O_DIRECT flag is required for GDS.
33        // However, a limitation of this flag is that all operations involving disk
34        // (both read and write) must be page-aligned.
35        // Pinned memory is already page-aligned, so we only need to align system memory.
36        // TODO(jthomson04): Is page size always 4KB?
37
38        // SAFETY: malloc returns suitably aligned memory or null on failure.
39        let result = unsafe { libc::posix_memalign(&mut ptr, 4096, len) };
40        if result != 0 {
41            return Err(StorageError::AllocationFailed(format!(
42                "posix_memalign failed for size {}",
43                len
44            )));
45        }
46        let ptr = NonNull::new(ptr as *mut u8).ok_or_else(|| {
47            StorageError::AllocationFailed(format!("malloc failed for size {}", len))
48        })?;
49
50        // Zero-initialize the memory
51        unsafe {
52            std::ptr::write_bytes(ptr.as_ptr(), 0, len);
53        }
54
55        Ok(Self { ptr, len })
56    }
57
58    /// Get a pointer to the underlying memory.
59    ///
60    /// # Safety
61    /// The caller must ensure the pointer is not used after this storage is dropped.
62    pub unsafe fn as_ptr(&self) -> *const u8 {
63        self.ptr.as_ptr()
64    }
65
66    /// Get a mutable pointer to the underlying memory.
67    ///
68    /// # Safety
69    /// The caller must ensure the pointer is not used after this storage is dropped
70    /// and that there are no other references to this memory.
71    pub unsafe fn as_mut_ptr(&mut self) -> *mut u8 {
72        self.ptr.as_ptr()
73    }
74}
75
76impl Drop for SystemStorage {
77    fn drop(&mut self) {
78        // SAFETY: pointer was allocated by malloc.
79        unsafe {
80            libc::free(self.ptr.as_ptr() as *mut libc::c_void);
81        }
82    }
83}
84
85impl MemoryDescription for SystemStorage {
86    fn addr(&self) -> usize {
87        self.ptr.as_ptr() as usize
88    }
89
90    fn size(&self) -> usize {
91        self.len
92    }
93
94    fn storage_kind(&self) -> StorageKind {
95        StorageKind::System
96    }
97
98    fn as_any(&self) -> &dyn Any {
99        self
100    }
101
102    fn nixl_descriptor(&self) -> Option<NixlDescriptor> {
103        None
104    }
105}
106
107// Support for NIXL registration
108impl super::nixl::NixlCompatible for SystemStorage {
109    fn nixl_params(&self) -> (*const u8, usize, nixl_sys::MemType, u64) {
110        (self.ptr.as_ptr(), self.len, nixl_sys::MemType::Dram, 0)
111    }
112}
113
114impl actions::Memset for SystemStorage {
115    fn memset(&mut self, value: u8, offset: usize, size: usize) -> Result<()> {
116        let end = offset
117            .checked_add(size)
118            .ok_or_else(|| StorageError::OperationFailed("memset: offset overflow".into()))?;
119        if end > self.len {
120            return Err(StorageError::OperationFailed(
121                "memset: offset + size > storage size".into(),
122            ));
123        }
124        unsafe {
125            let ptr = self.ptr.as_ptr().add(offset);
126            std::ptr::write_bytes(ptr, value, size);
127        }
128        Ok(())
129    }
130}
131
132impl actions::Slice for SystemStorage {
133    unsafe fn as_slice(&self) -> Result<&[u8]> {
134        // SAFETY: SystemStorage owns the memory allocated via the global allocator.
135        // The memory remains valid as long as this SystemStorage instance exists.
136        // The ptr is guaranteed to be valid for `self.len` bytes.
137        // Caller must ensure no concurrent mutable access per trait contract.
138        // SAFETY: The pointer is valid, properly aligned, and points to `self.len` bytes.
139        Ok(unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.len) })
140    }
141}
dynamo_memory/system.rs

dynamo_memory/
system.rs
