entrenar/efficiency/device/
compute.rs

1//! Compute device abstraction.
2
3use serde::{Deserialize, Serialize};
4
5use super::apple::AppleSiliconInfo;
6use super::cpu::CpuInfo;
7use super::gpu::GpuInfo;
8use super::tpu::TpuInfo;
9
10/// Compute device abstraction
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12pub enum ComputeDevice {
13    /// CPU device
14    Cpu(CpuInfo),
15    /// GPU device (NVIDIA, AMD, Intel)
16    Gpu(GpuInfo),
17    /// TPU device (Google)
18    Tpu(TpuInfo),
19    /// Apple Silicon with unified memory
20    AppleSilicon(AppleSiliconInfo),
21}
22
23impl ComputeDevice {
24    /// Auto-detect available compute devices
25    pub fn detect() -> Vec<Self> {
26        let mut devices = Vec::new();
27
28        // Always detect CPU
29        devices.push(Self::Cpu(CpuInfo::detect()));
30
31        // Check for Apple Silicon
32        if let Some(apple) = AppleSiliconInfo::detect() {
33            devices.push(Self::AppleSilicon(apple));
34        }
35
36        // GPU detection would require platform-specific APIs (CUDA, ROCm, Metal)
37        // For now, we only auto-detect CPU and Apple Silicon
38
39        devices
40    }
41
42    /// Check if this is a GPU device
43    pub fn is_gpu(&self) -> bool {
44        matches!(self, Self::Gpu(_))
45    }
46
47    /// Check if this is a CPU device
48    pub fn is_cpu(&self) -> bool {
49        matches!(self, Self::Cpu(_))
50    }
51
52    /// Check if this is a TPU device
53    pub fn is_tpu(&self) -> bool {
54        matches!(self, Self::Tpu(_))
55    }
56
57    /// Check if this is Apple Silicon
58    pub fn is_apple_silicon(&self) -> bool {
59        matches!(self, Self::AppleSilicon(_))
60    }
61
62    /// Get available memory in bytes
63    pub fn memory_bytes(&self) -> u64 {
64        match self {
65            Self::Cpu(info) => {
66                // Return system memory (approximation)
67                // In real implementation, would query actual system RAM
68                u64::from(info.cores) * 4 * 1024 * 1024 * 1024 // ~4GB per core estimate
69            }
70            Self::Gpu(info) => info.vram_bytes,
71            Self::Tpu(info) => info.hbm_bytes,
72            Self::AppleSilicon(info) => info.unified_memory_bytes,
73        }
74    }
75
76    /// Get device name
77    pub fn name(&self) -> &str {
78        match self {
79            Self::Cpu(info) => &info.model,
80            Self::Gpu(info) => &info.name,
81            Self::Tpu(info) => &info.version,
82            Self::AppleSilicon(info) => &info.chip,
83        }
84    }
85
86    /// Get compute cores/units
87    pub fn compute_units(&self) -> u32 {
88        match self {
89            Self::Cpu(info) => info.threads,
90            Self::Gpu(_) => 0, // Would need more info
91            Self::Tpu(info) => info.cores,
92            Self::AppleSilicon(info) => info.total_cpu_cores() + info.gpu_cores,
93        }
94    }
95
96    /// Estimate relative compute power (normalized, CPU = 1.0)
97    pub fn relative_compute_power(&self) -> f64 {
98        match self {
99            Self::Cpu(info) => f64::from(info.threads) / 8.0, // Normalize to 8-thread CPU
100            Self::Gpu(info) => {
101                // Rough estimate based on VRAM (proxy for capability)
102                10.0 * (info.vram_gb() / 8.0) // 8GB GPU ~ 10x CPU
103            }
104            Self::Tpu(info) => {
105                // TPUs are highly optimized for matrix ops
106                50.0 * (f64::from(info.cores) / 8.0)
107            }
108            Self::AppleSilicon(info) => {
109                // P-cores are faster, E-cores are efficient
110                (f64::from(info.p_cores) * 1.5 + f64::from(info.e_cores) * 0.5) / 8.0
111                    + f64::from(info.gpu_cores) * 0.5
112            }
113        }
114    }
115}
116
117impl std::fmt::Display for ComputeDevice {
118    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119        match self {
120            Self::Cpu(info) => write!(
121                f,
122                "CPU: {} ({} cores, {} threads, {})",
123                info.model, info.cores, info.threads, info.simd
124            ),
125            Self::Gpu(info) => {
126                write!(f, "GPU: {} ({:.1} GB VRAM", info.name, info.vram_gb())?;
127                if let Some((major, minor)) = info.compute_capability {
128                    write!(f, ", SM {major}.{minor}")?;
129                }
130                write!(f, ")")
131            }
132            Self::Tpu(info) => write!(
133                f,
134                "TPU: {} ({} cores, {:.1} GB HBM)",
135                info.version,
136                info.cores,
137                info.hbm_gb()
138            ),
139            Self::AppleSilicon(info) => write!(
140                f,
141                "Apple Silicon: {} ({}P+{}E cores, {} GPU cores, {:.1} GB)",
142                info.chip,
143                info.p_cores,
144                info.e_cores,
145                info.gpu_cores,
146                info.unified_memory_gb()
147            ),
148        }
149    }
150}
entrenar/efficiency/device/compute.rs

entrenar/efficiency/device/
compute.rs
