Merge branch 'feat/f64-cpu-precision'

crates/core/rtx-backend-cpu/src/tensor.rs

@@ -4,13 +4,15 @@ use crate::CpuDevice;
 
 /// CPU tensor storage.
 ///
-/// Stores tensor data as a contiguous `Vec<f32>` in CPU memory.
+/// Stores tensor data as a contiguous `Vec<E>` in CPU memory, where the element
-/// All operations are performed directly on CPU with optional
+/// type `E` defaults to `f32`. The element parameter lets the same primitive back
-/// parallelization via Rayon.
+/// both the f32 `CpuBackend` and the f64 `CpuBackendF64`; existing
+/// `CpuTensorPrimitive<D>` references remain `f32` via the default. All operations
+/// are performed directly on CPU with optional parallelization via Rayon.
 #[derive(Clone)]
-pub struct CpuTensorPrimitive<const D: usize> {
+pub struct CpuTensorPrimitive<const D: usize, E = f32> {
     /// Tensor data
-    pub(crate) data: Vec<f32>,
+    pub(crate) data: Vec<E>,
     /// Tensor shape
     pub(crate) shape: [usize; D],
     /// Tensor strides
@@ -19,9 +21,9 @@ pub struct CpuTensorPrimitive<const D: usize> {
     pub(crate) device: CpuDevice,
 }
 
-impl<const D: usize> CpuTensorPrimitive<D> {
+impl<const D: usize, E> CpuTensorPrimitive<D, E> {
     /// Create a new CPU tensor.
-    pub fn new(data: Vec<f32>, shape: [usize; D], device: CpuDevice) -> Self {
+    pub fn new(data: Vec<E>, shape: [usize; D], device: CpuDevice) -> Self {
         let strides = Self::compute_strides(&shape);
         let numel: usize = shape.iter().product();
         assert_eq!(data.len(), numel, "Data length must match shape");
@@ -62,22 +64,24 @@ impl<const D: usize> CpuTensorPrimitive<D> {
     }
 
     /// Get reference to data.
-    pub fn data(&self) -> &[f32] {
+    pub fn data(&self) -> &[E] {
         &self.data
     }
 
     /// Get mutable reference to data.
-    pub fn data_mut(&mut self) -> &mut [f32] {
+    pub fn data_mut(&mut self) -> &mut [E] {
         &mut self.data
     }
+}
 
+impl<const D: usize, E: Clone> CpuTensorPrimitive<D, E> {
     /// Clone data to a vector.
-    pub fn to_vec(&self) -> Vec<f32> {
+    pub fn to_vec(&self) -> Vec<E> {
         self.data.clone()
     }
 }
 
-impl<const D: usize> std::fmt::Debug for CpuTensorPrimitive<D> {
+impl<const D: usize, E> std::fmt::Debug for CpuTensorPrimitive<D, E> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("CpuTensorPrimitive")
             .field("shape", &self.shape)
@@ -88,14 +92,14 @@ impl<const D: usize> std::fmt::Debug for CpuTensorPrimitive<D> {
 }
 
 // SAFETY: CpuTensorPrimitive only contains:
-// - data: Vec<f32> which is Send+Sync when f32 is (and f32 is)
+// - data: Vec<E> which is Send+Sync when E is
 // - shape: [usize; D] which is Copy and trivially thread-safe
 // - strides: [usize; D] which is Copy and trivially thread-safe
 // - device: CpuDevice which derives Clone and has no mutable state
-// All fields are thread-safe, so the type is safe to send between threads
+// All fields are thread-safe (given E is), so the type is safe to send between
-// and share immutably across threads.
+// threads and share immutably across threads.
-unsafe impl<const D: usize> Send for CpuTensorPrimitive<D> {}
+unsafe impl<const D: usize, E: Send> Send for CpuTensorPrimitive<D, E> {}
-unsafe impl<const D: usize> Sync for CpuTensorPrimitive<D> {}
+unsafe impl<const D: usize, E: Sync> Sync for CpuTensorPrimitive<D, E> {}
 
 #[cfg(test)]
 mod tests {
@@ -104,7 +108,7 @@ mod tests {
     #[test]
     fn test_tensor_creation() {
         let device = CpuDevice::new();
-        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
+        let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
         let tensor = CpuTensorPrimitive::new(data.clone(), [2, 3], device);
 
         assert_eq!(tensor.shape(), &[2, 3]);
@@ -115,8 +119,14 @@ mod tests {
 
     #[test]
     fn test_strides_computation() {
-        assert_eq!(CpuTensorPrimitive::compute_strides(&[2, 3, 4]), [12, 4, 1]);
+        assert_eq!(
-        assert_eq!(CpuTensorPrimitive::compute_strides(&[10, 20]), [20, 1]);
+            CpuTensorPrimitive::<3, f32>::compute_strides(&[2, 3, 4]),
-        assert_eq!(CpuTensorPrimitive::compute_strides(&[5]), [1]);
+            [12, 4, 1]
+        );
+        assert_eq!(
+            CpuTensorPrimitive::<2, f32>::compute_strides(&[10, 20]),
+            [20, 1]
+        );
+        assert_eq!(CpuTensorPrimitive::<1, f32>::compute_strides(&[5]), [1]);
     }
 }