rustytorch/crates/core/rtx-backend-cpu/src/device.rs

//! CPU device abstraction.

use rtx_backend::{DeviceId, DeviceOps};

use crate::{CpuBackend, CpuBackendF64};

/// CPU device for the CPU backend.
///
/// Represents the system CPU. There is typically only one logical
/// CPU device, though it may have multiple cores.
#[derive(Clone)]
pub struct CpuDevice {
    /// Device index (always 0 for CPU)
    index: usize,
    /// Number of available threads
    num_threads: usize,
}

impl CpuDevice {
    /// Create a new CPU device.
    pub fn new() -> Self {
        let num_threads = rayon::current_num_threads();
        Self {
            index: 0,
            num_threads,
        }
    }

    /// Get the number of available threads.
    pub fn num_threads(&self) -> usize {
        self.num_threads
    }
}

impl Default for CpuDevice {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Debug for CpuDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CpuDevice")
            .field("index", &self.index)
            .field("num_threads", &self.num_threads)
            .finish()
    }
}

impl PartialEq for CpuDevice {
    fn eq(&self, other: &Self) -> bool {
        self.index == other.index
    }
}

impl Eq for CpuDevice {}

impl std::hash::Hash for CpuDevice {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.index.hash(state);
    }
}

impl DeviceOps<CpuBackend> for CpuDevice {
    fn id(&self) -> DeviceId {
        DeviceId::Cpu
    }

    fn memory_capacity(&self) -> usize {
        // Return system memory (rough estimate)
        // In production, would use sys_info crate
        16 * 1024 * 1024 * 1024 // 16GB default
    }

    fn memory_available(&self) -> usize {
        // Estimate available memory
        (16 * 1024 * 1024 * 1024_usize) / 2
    }

    fn compute_capability(&self) -> Option<(u32, u32)> {
        // CPU doesn't have compute capability
        None
    }

    fn synchronize(&self) {
        // CPU operations are synchronous
    }

    fn is_available(&self) -> bool {
        true // CPU is always available
    }
}

/// Same device, viewed through the f64 backend (CPU has one logical device).
impl DeviceOps<CpuBackendF64> for CpuDevice {
    fn id(&self) -> DeviceId {
        DeviceId::Cpu
    }

    fn memory_capacity(&self) -> usize {
        16 * 1024 * 1024 * 1024 // 16GB default
    }

    fn memory_available(&self) -> usize {
        (16 * 1024 * 1024 * 1024_usize) / 2
    }

    fn compute_capability(&self) -> Option<(u32, u32)> {
        None
    }

    fn synchronize(&self) {}

    fn is_available(&self) -> bool {
        true
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_device_creation() {
        let device = CpuDevice::new();
        // `is_available` is shared by both DeviceOps impls (f32/f64); pick one.
        assert!(DeviceOps::<CpuBackend>::is_available(&device));
        assert!(device.num_threads() > 0);
        println!("CPU threads: {}", device.num_threads());
    }
}