axonml_core/
allocator.rs

1//! Allocator - Memory Allocation Traits and Implementations
2//!
3//! Provides the allocator abstraction for device-specific memory management.
4//! Each backend implements the Allocator trait for its memory operations.
5//!
6//! # Key Features
7//! - Unified allocator trait for all devices
8//! - Pluggable allocator implementations
9//! - Memory pool support for performance
10//!
11//! @version 0.1.0
12//! @author `AutomataNexus` Development Team
13
14use crate::device::Device;
15use crate::dtype::Scalar;
16use crate::error::Result;
17use sysinfo::System;
18
19// =============================================================================
20// Default Allocator
21// =============================================================================
22
23/// Default CPU allocator using system memory.
24#[derive(Debug, Clone, Copy, Default)]
25pub struct DefaultAllocator;
26
27impl DefaultAllocator {
28    /// Creates a new default allocator.
29    #[must_use]
30    pub const fn new() -> Self {
31        Self
32    }
33
34    /// Returns the device this allocator is for.
35    #[must_use]
36    pub const fn device(&self) -> Device {
37        Device::Cpu
38    }
39
40    /// Allocates memory for `count` elements of type T.
41    pub fn allocate<T: Scalar>(&self, count: usize) -> Result<*mut T> {
42        let mut vec = Vec::<T>::with_capacity(count);
43        let ptr = vec.as_mut_ptr();
44        core::mem::forget(vec);
45        Ok(ptr)
46    }
47
48    /// Deallocates memory previously allocated.
49    ///
50    /// # Safety
51    /// The pointer must have been allocated by this allocator.
52    pub unsafe fn deallocate<T: Scalar>(&self, ptr: *mut T, count: usize) {
53        drop(Vec::from_raw_parts(ptr, 0, count));
54    }
55
56    /// Copies memory from one location to another.
57    ///
58    /// # Safety
59    /// Both pointers must be valid for `count` elements.
60    pub unsafe fn copy<T: Scalar>(&self, dst: *mut T, src: *const T, count: usize) {
61        core::ptr::copy_nonoverlapping(src, dst, count);
62    }
63
64    /// Fills memory with zeros.
65    ///
66    /// # Safety
67    /// The pointer must be valid for `count` elements.
68    pub unsafe fn zero<T: Scalar>(&self, ptr: *mut T, count: usize) {
69        core::ptr::write_bytes(ptr, 0, count);
70    }
71
72    /// Returns the total memory available on the device.
73    #[must_use]
74    pub fn total_memory(&self) -> usize {
75        let sys = System::new_all();
76        sys.total_memory() as usize
77    }
78
79    /// Returns the currently free memory on the device.
80    #[must_use]
81    pub fn free_memory(&self) -> usize {
82        let sys = System::new_all();
83        sys.available_memory() as usize
84    }
85}
86
87// =============================================================================
88// Allocator Trait (for future extensibility)
89// =============================================================================
90
91/// Marker trait for types that can act as allocators.
92///
93/// Note: Due to Rust's object safety rules, we use concrete types
94/// instead of dynamic dispatch for allocators.
95pub trait Allocator {
96    /// Returns the device this allocator is for.
97    fn device(&self) -> Device;
98
99    /// Returns the total memory available.
100    fn total_memory(&self) -> usize;
101
102    /// Returns the free memory available.
103    fn free_memory(&self) -> usize;
104}
105
106impl Allocator for DefaultAllocator {
107    fn device(&self) -> Device {
108        Device::Cpu
109    }
110
111    fn total_memory(&self) -> usize {
112        self.total_memory()
113    }
114
115    fn free_memory(&self) -> usize {
116        self.free_memory()
117    }
118}
119
120// =============================================================================
121// Tests
122// =============================================================================
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn test_default_allocator() {
130        let alloc = DefaultAllocator::new();
131        assert_eq!(alloc.device(), Device::Cpu);
132        assert!(alloc.total_memory() > 0);
133    }
134
135    #[test]
136    fn test_allocate_deallocate() {
137        let alloc = DefaultAllocator::new();
138
139        let ptr = alloc.allocate::<f32>(100).unwrap();
140        assert!(!ptr.is_null());
141
142        unsafe {
143            alloc.zero(ptr, 100);
144            alloc.deallocate(ptr, 100);
145        }
146    }
147
148    #[test]
149    fn test_copy() {
150        let alloc = DefaultAllocator::new();
151
152        let src = alloc.allocate::<f32>(10).unwrap();
153        let dst = alloc.allocate::<f32>(10).unwrap();
154
155        unsafe {
156            for i in 0..10 {
157                *src.add(i) = i as f32;
158            }
159
160            alloc.copy(dst, src, 10);
161
162            for i in 0..10 {
163                assert_eq!(*dst.add(i), i as f32);
164            }
165
166            alloc.deallocate(src, 10);
167            alloc.deallocate(dst, 10);
168        }
169    }
170}
axonml_core/allocator.rs

axonml_core/
allocator.rs
