etensor_core/backends/cpu/
alloc.rs

1//! High-performance CPU memory allocators and tensor factories.
2
3use crate::tensor::Tensor;
4use crate::shape::Shape;
5use crate::buffer::Buffer;
6use crate::dtypes::DType;
7use crate::device::Device;
8use crate::errors::{EtensorError, EtensorResult};
9
10/// Physical memory allocator for Host (CPU) RAM.
11pub struct CpuAllocator;
12
13impl CpuAllocator {
14    /// Allocates a contiguous CPU tensor filled with zeros.
15    pub fn zeros(shape: Shape, dtype: DType, requires_grad: bool) -> EtensorResult<Tensor> {
16        if dtype != DType::F32 {
17            return Err(EtensorError::InternalError(
18                "CpuAllocator currently only supports F32 initialization.".to_string(),
19            ));
20        }
21
22        let num_elements = shape.num_elements();
23        let buffer = Buffer::from_f32_vec(vec![0.0; num_elements]);
24
25        Ok(Tensor::new(buffer, shape, Device::Cpu, dtype, requires_grad))
26    }
27
28    /// Allocates a contiguous CPU tensor filled with ones.
29    pub fn ones(shape: Shape, dtype: DType, requires_grad: bool) -> EtensorResult<Tensor> {
30        if dtype != DType::F32 {
31            return Err(EtensorError::InternalError(
32                "CpuAllocator currently only supports F32 initialization.".to_string(),
33            ));
34        }
35
36        let num_elements = shape.num_elements();
37        let buffer = Buffer::from_f32_vec(vec![1.0; num_elements]);
38
39        Ok(Tensor::new(buffer, shape, Device::Cpu, dtype, requires_grad))
40    }
41
42    /// Allocates a contiguous CPU tensor filled with a specific scalar value.
43    pub fn full(shape: Shape, value: f32, dtype: DType, requires_grad: bool) -> EtensorResult<Tensor> {
44        if dtype != DType::F32 {
45            return Err(EtensorError::InternalError(
46                "CpuAllocator currently only supports F32 initialization.".to_string(),
47            ));
48        }
49
50        let num_elements = shape.num_elements();
51        let buffer = Buffer::from_f32_vec(vec![value; num_elements]);
52
53        Ok(Tensor::new(buffer, shape, Device::Cpu, dtype, requires_grad))
54    }
55
56    /// Creates a 1D tensor with evenly spaced values within a given interval.
57    pub fn arange(start: f32, end: f32, step: f32, dtype: DType, requires_grad: bool) -> EtensorResult<Tensor> {
58        if dtype != DType::F32 {
59            return Err(EtensorError::InternalError(
60                "CpuAllocator currently only supports F32 initialization.".to_string(),
61            ));
62        }
63        
64        if step == 0.0 {
65            return Err(EtensorError::InternalError("Arange step size cannot be zero.".to_string()));
66        }
67
68        // Calculate total elements required
69        let size = ((end - start) / step).ceil() as usize;
70        let mut data = Vec::with_capacity(size);
71        
72        let mut current = start;
73        for _ in 0..size {
74            data.push(current);
75            current += step;
76        }
77
78        let shape = Shape::new(vec![size]);
79        let buffer = Buffer::from_f32_vec(data);
80
81        Ok(Tensor::new(buffer, shape, Device::Cpu, dtype, requires_grad))
82    }
83}
84
85// =====================================================================
86// UNIT TESTS
87// =====================================================================
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn test_cpu_alloc_zeros() {
94        let shape = Shape::new(vec![2, 3]);
95        let t = CpuAllocator::zeros(shape, DType::F32, false).unwrap();
96        
97        let slice = t.data.as_f32_slice().unwrap();
98        assert_eq!(slice.len(), 6);
99        assert_eq!(slice, &[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
100    }
101
102    #[test]
103    fn test_cpu_alloc_ones() {
104        let shape = Shape::new(vec![4]);
105        let t = CpuAllocator::ones(shape, DType::F32, true).unwrap();
106        
107        let slice = t.data.as_f32_slice().unwrap();
108        assert_eq!(slice, &[1.0, 1.0, 1.0, 1.0]);
109        assert!(t.requires_grad);
110    }
111
112    #[test]
113    fn test_cpu_alloc_full() {
114        let shape = Shape::new(vec![2, 2]);
115        let t = CpuAllocator::full(shape, 42.0, DType::F32, false).unwrap();
116        
117        let slice = t.data.as_f32_slice().unwrap();
118        assert_eq!(slice, &[42.0, 42.0, 42.0, 42.0]);
119    }
120
121    #[test]
122    fn test_cpu_alloc_arange() {
123        // Range from 0.0 to 5.0 with step 1.0 -> [0.0, 1.0, 2.0, 3.0, 4.0]
124        let t = CpuAllocator::arange(0.0, 5.0, 1.0, DType::F32, false).unwrap();
125        
126        let slice = t.data.as_f32_slice().unwrap();
127        assert_eq!(slice, &[0.0, 1.0, 2.0, 3.0, 4.0]);
128        assert_eq!(t.shape.dims, vec![5]);
129    }
130}
etensor_core/backends/cpu/alloc.rs

etensor_core/backends/cpu/
alloc.rs
