ferrum_engine/parallel/
device.rs

1//! Device Management
2//!
3//! Device discovery, capability detection, and resource monitoring
4//! for multi-GPU environments.
5
6use ferrum_types::{Device, FerrumError, Result};
7use parking_lot::RwLock;
8use std::collections::HashMap;
9use tracing::{info, warn};
10
11/// Device capability information
12#[derive(Debug, Clone)]
13pub struct DeviceCapability {
14    /// Compute capability (e.g., 8.0 for A100)
15    pub compute_capability: (u32, u32),
16    /// Total memory in bytes
17    pub total_memory: usize,
18    /// Number of SMs (for CUDA) or compute units
19    pub num_compute_units: u32,
20    /// Maximum threads per block
21    pub max_threads_per_block: u32,
22    /// Warp size (32 for CUDA, varies for others)
23    pub warp_size: u32,
24    /// Whether tensor cores are available
25    pub has_tensor_cores: bool,
26    /// Whether unified memory is supported
27    pub unified_memory: bool,
28    /// Maximum shared memory per block
29    pub max_shared_memory: usize,
30    /// Memory bandwidth (GB/s)
31    pub memory_bandwidth: f32,
32    /// Peak FLOPS (TF32)
33    pub peak_tflops: f32,
34}
35
36impl Default for DeviceCapability {
37    fn default() -> Self {
38        Self {
39            compute_capability: (0, 0),
40            total_memory: 0,
41            num_compute_units: 0,
42            max_threads_per_block: 1024,
43            warp_size: 32,
44            has_tensor_cores: false,
45            unified_memory: false,
46            max_shared_memory: 48 * 1024, // 48KB default
47            memory_bandwidth: 0.0,
48            peak_tflops: 0.0,
49        }
50    }
51}
52
53impl DeviceCapability {
54    /// Create capability info for CPU
55    pub fn cpu() -> Self {
56        Self {
57            compute_capability: (0, 0),
58            total_memory: 0,      // To be filled from system info
59            num_compute_units: 0, // CPU cores
60            max_threads_per_block: 1,
61            warp_size: 1,
62            has_tensor_cores: false,
63            unified_memory: true,
64            max_shared_memory: 0,
65            memory_bandwidth: 0.0,
66            peak_tflops: 0.0,
67        }
68    }
69
70    /// Create capability info for Apple Silicon (Metal)
71    pub fn apple_silicon(total_memory: usize, gpu_cores: u32) -> Self {
72        Self {
73            compute_capability: (1, 0), // Metal version indicator
74            total_memory,
75            num_compute_units: gpu_cores,
76            max_threads_per_block: 1024,
77            warp_size: 32,           // SIMD width
78            has_tensor_cores: false, // No tensor cores, but has AMX
79            unified_memory: true,
80            max_shared_memory: 32 * 1024,
81            memory_bandwidth: 200.0, // Varies by chip
82            peak_tflops: 10.0,       // Varies by chip
83        }
84    }
85
86    /// Check if device can run a model of given size
87    pub fn can_fit_model(&self, model_size_bytes: usize) -> bool {
88        // Leave some memory for KV cache and intermediate activations
89        let usable_memory = self.total_memory * 70 / 100; // 70% for model
90        model_size_bytes <= usable_memory
91    }
92}
93
94/// Device information with runtime state
95#[derive(Debug, Clone)]
96pub struct DeviceInfo {
97    /// Device identifier
98    pub device: Device,
99    /// Device name
100    pub name: String,
101    /// Device capabilities
102    pub capability: DeviceCapability,
103    /// Current memory usage in bytes
104    pub used_memory: usize,
105    /// Whether device is available
106    pub is_available: bool,
107    /// Current utilization (0.0 - 1.0)
108    pub utilization: f32,
109}
110
111impl DeviceInfo {
112    /// Create info for CPU
113    pub fn cpu() -> Self {
114        Self {
115            device: Device::CPU,
116            name: "CPU".to_string(),
117            capability: DeviceCapability::cpu(),
118            used_memory: 0,
119            is_available: true,
120            utilization: 0.0,
121        }
122    }
123
124    /// Get available memory
125    pub fn available_memory(&self) -> usize {
126        self.capability
127            .total_memory
128            .saturating_sub(self.used_memory)
129    }
130}
131
132/// Device manager for multi-GPU coordination
133pub struct DeviceManager {
134    /// All discovered devices
135    devices: RwLock<HashMap<Device, DeviceInfo>>,
136    /// Primary device (first available GPU or CPU)
137    primary_device: RwLock<Device>,
138}
139
140impl DeviceManager {
141    /// Create a new device manager
142    pub fn new() -> Self {
143        Self {
144            devices: RwLock::new(HashMap::new()),
145            primary_device: RwLock::new(Device::CPU),
146        }
147    }
148
149    /// Discover available devices
150    pub fn discover_devices(&self) -> Result<()> {
151        info!("Discovering available devices...");
152        let mut devices = self.devices.write();
153        devices.clear();
154
155        // Always add CPU
156        let cpu_info = DeviceInfo::cpu();
157        devices.insert(Device::CPU, cpu_info);
158
159        // Try to detect GPUs
160        #[cfg(all(feature = "metal", any(target_os = "macos", target_os = "ios")))]
161        {
162            if let Ok(metal_info) = self.detect_metal_devices() {
163                for info in metal_info {
164                    info!("Found Metal device: {}", info.name);
165                    *self.primary_device.write() = info.device.clone();
166                    devices.insert(info.device.clone(), info);
167                }
168            }
169        }
170
171        // CUDA device detection would go here
172        #[cfg(feature = "cuda")]
173        {
174            if let Ok(cuda_info) = self.detect_cuda_devices() {
175                for (idx, info) in cuda_info.into_iter().enumerate() {
176                    info!("Found CUDA device {}: {}", idx, info.name);
177                    if idx == 0 {
178                        *self.primary_device.write() = info.device.clone();
179                    }
180                    devices.insert(info.device.clone(), info);
181                }
182            }
183        }
184
185        info!("Discovered {} device(s)", devices.len());
186        Ok(())
187    }
188
189    /// Get all available devices
190    pub fn get_devices(&self) -> Vec<DeviceInfo> {
191        self.devices.read().values().cloned().collect()
192    }
193
194    /// Get device info by ID
195    pub fn get_device(&self, device: &Device) -> Option<DeviceInfo> {
196        self.devices.read().get(device).cloned()
197    }
198
199    /// Get primary device
200    pub fn primary_device(&self) -> Device {
201        self.primary_device.read().clone()
202    }
203
204    /// Set primary device
205    pub fn set_primary_device(&self, device: Device) -> Result<()> {
206        if self.devices.read().contains_key(&device) {
207            *self.primary_device.write() = device;
208            Ok(())
209        } else {
210            Err(FerrumError::not_found(format!(
211                "Device {:?} not found",
212                device
213            )))
214        }
215    }
216
217    /// Get available GPU devices
218    pub fn get_gpu_devices(&self) -> Vec<DeviceInfo> {
219        self.devices
220            .read()
221            .values()
222            .filter(|info| Self::is_gpu_device(&info.device))
223            .cloned()
224            .collect()
225    }
226
227    /// Get total available GPU memory across all devices
228    pub fn total_gpu_memory(&self) -> usize {
229        self.devices
230            .read()
231            .values()
232            .filter(|info| Self::is_gpu_device(&info.device))
233            .map(|info| info.capability.total_memory)
234            .sum()
235    }
236
237    /// Check if a device is a GPU (CUDA, ROCm, or Metal)
238    fn is_gpu_device(device: &Device) -> bool {
239        match device {
240            Device::CPU => false,
241            Device::CUDA(_) => true,
242            Device::ROCm(_) => true,
243            #[cfg(any(target_os = "macos", target_os = "ios"))]
244            Device::Metal => true,
245        }
246    }
247
248    /// Update memory usage for a device
249    pub fn update_memory_usage(&self, device: &Device, used_bytes: usize) {
250        if let Some(info) = self.devices.write().get_mut(device) {
251            info.used_memory = used_bytes;
252        }
253    }
254
255    /// Find devices that can fit a model of given size
256    pub fn find_devices_for_model(&self, model_size_bytes: usize) -> Vec<DeviceInfo> {
257        self.devices
258            .read()
259            .values()
260            .filter(|info| info.capability.can_fit_model(model_size_bytes))
261            .cloned()
262            .collect()
263    }
264
265    /// Select optimal devices for parallel execution
266    pub fn select_devices_for_parallelism(
267        &self,
268        model_size_bytes: usize,
269        min_devices: usize,
270        max_devices: usize,
271    ) -> Vec<DeviceInfo> {
272        let gpu_devices: Vec<_> = self.get_gpu_devices();
273
274        // If we have enough GPUs that can individually fit the model, use those
275        let fitting_devices: Vec<_> = gpu_devices
276            .iter()
277            .filter(|d| d.capability.can_fit_model(model_size_bytes))
278            .cloned()
279            .collect();
280
281        if fitting_devices.len() >= min_devices {
282            return fitting_devices.into_iter().take(max_devices).collect();
283        }
284
285        // Otherwise, select devices for model parallelism based on total memory
286        let mut sorted_devices = gpu_devices;
287        sorted_devices.sort_by(|a, b| b.capability.total_memory.cmp(&a.capability.total_memory));
288
289        let mut selected = Vec::new();
290        let mut total_memory = 0usize;
291
292        for device in sorted_devices {
293            if selected.len() >= max_devices {
294                break;
295            }
296            total_memory += device.capability.total_memory;
297            selected.push(device);
298
299            // Check if we have enough memory for model + overhead
300            let required = model_size_bytes + model_size_bytes / 4; // 25% overhead
301            if total_memory >= required && selected.len() >= min_devices {
302                break;
303            }
304        }
305
306        // Fall back to CPU if no GPUs available
307        if selected.is_empty() {
308            warn!("No suitable GPUs found, falling back to CPU");
309            if let Some(cpu) = self.get_device(&Device::CPU) {
310                selected.push(cpu);
311            }
312        }
313
314        selected
315    }
316
317    #[cfg(all(feature = "metal", any(target_os = "macos", target_os = "ios")))]
318    fn detect_metal_devices(&self) -> Result<Vec<DeviceInfo>> {
319        // In actual implementation, this would use metal-rs to detect devices
320        // For now, return a placeholder for the default GPU
321        Ok(vec![DeviceInfo {
322            device: Device::Metal,
323            name: "Apple GPU".to_string(),
324            capability: DeviceCapability::apple_silicon(16 * 1024 * 1024 * 1024, 76), // Placeholder
325            used_memory: 0,
326            is_available: true,
327            utilization: 0.0,
328        }])
329    }
330
331    #[cfg(feature = "cuda")]
332    fn detect_cuda_devices(&self) -> Result<Vec<DeviceInfo>> {
333        let count = ferrum_kernels::cuda_device_count()
334            .map_err(|error| FerrumError::device(format!("CUDA probe failed: {error}")))?;
335        if count == 0 {
336            Err(FerrumError::unsupported("No CUDA devices found"))
337        } else {
338            Ok((0..count)
339                .map(|idx| DeviceInfo {
340                    device: Device::CUDA(idx),
341                    name: ferrum_kernels::cuda_device_name(idx)
342                        .unwrap_or_else(|_| format!("CUDA Device {idx}")),
343                    capability: DeviceCapability {
344                        compute_capability: (0, 0),
345                        total_memory: 0,
346                        unified_memory: false,
347                        ..DeviceCapability::default()
348                    },
349                    used_memory: 0,
350                    is_available: true,
351                    utilization: 0.0,
352                })
353                .collect())
354        }
355    }
356}
357
358impl Default for DeviceManager {
359    fn default() -> Self {
360        let manager = Self::new();
361        let _ = manager.discover_devices();
362        manager
363    }
364}
365
366/// Global device manager
367static GLOBAL_DEVICE_MANAGER: std::sync::OnceLock<DeviceManager> = std::sync::OnceLock::new();
368
369/// Get the global device manager
370pub fn global_device_manager() -> &'static DeviceManager {
371    GLOBAL_DEVICE_MANAGER.get_or_init(DeviceManager::default)
372}
373
374// ============================================================================
375// Tests
376// ============================================================================
377
378#[cfg(test)]
379mod tests {
380    use super::*;
381
382    #[test]
383    fn test_device_capability_default() {
384        let cap = DeviceCapability::default();
385        assert_eq!(cap.compute_capability, (0, 0));
386        assert!(!cap.has_tensor_cores);
387    }
388
389    #[test]
390    fn test_device_info_cpu() {
391        let info = DeviceInfo::cpu();
392        assert_eq!(info.device, Device::CPU);
393        assert!(info.is_available);
394    }
395
396    #[test]
397    fn test_device_manager_discover() {
398        let manager = DeviceManager::new();
399        manager.discover_devices().unwrap();
400
401        let devices = manager.get_devices();
402        assert!(!devices.is_empty()); // At least CPU should be there
403    }
404
405    #[test]
406    fn test_can_fit_model() {
407        let mut cap = DeviceCapability::default();
408        cap.total_memory = 16 * 1024 * 1024 * 1024; // 16GB
409
410        assert!(cap.can_fit_model(10 * 1024 * 1024 * 1024)); // 10GB model
411        assert!(!cap.can_fit_model(15 * 1024 * 1024 * 1024)); // 15GB model (exceeds 70%)
412    }
413}
ferrum_engine/parallel/device.rs

ferrum_engine/parallel/
device.rs
