axonml_core/
allocator.rs

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

axonml_core/
allocator.rs
