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

//! Metal device wrapper for Backend trait.
//!
//! Provides multi-GPU support for Metal devices, including eGPU detection.

use rtx_backend::{DeviceId, DeviceOps};
use rtx_metal::{MetalDevice, MetalDeviceInfo};
use std::sync::Arc;

use crate::{MetalBackend, MetalBackendError, MetalBackendResult};

/// Metal device wrapper implementing DeviceOps.
///
/// Wraps the existing `rtx_metal::MetalDevice` to implement the Backend device trait.
/// Supports multiple GPUs including eGPUs via device index.
#[derive(Clone)]
pub struct MetalDeviceWrapper {
    /// Wrapped Metal device
    pub(crate) device: Arc<MetalDevice>,
    /// Device index
    pub(crate) index: usize,
}

impl MetalDeviceWrapper {
    /// Create a new Metal device wrapper for the default device.
    pub fn new() -> MetalBackendResult<Self> {
        Self::new_with_index(0)
    }

    /// Create a new Metal device wrapper for a specific device index.
    ///
    /// # Arguments
    /// * `index` - The device index (0-based). Use `device_count()` to get available count.
    ///
    /// # Example
    /// ```rust,ignore
    /// // Get the second GPU (e.g., an eGPU)
    /// let device = MetalDeviceWrapper::new_with_index(1)?;
    /// if device.is_egpu() {
    ///     println!("Using external GPU: {}", device.name());
    /// }
    /// ```
    pub fn new_with_index(index: usize) -> MetalBackendResult<Self> {
        let device = MetalDevice::by_index(index).map_err(|e| {
            MetalBackendError::DeviceInit(format!(
                "Failed to initialize Metal device {}: {:?}",
                index, e
            ))
        })?;

        Ok(Self {
            device: Arc::new(device),
            index,
        })
    }

    /// Get the number of available Metal devices.
    ///
    /// Returns the total number of Metal-capable GPUs, including eGPUs.
    pub fn device_count() -> usize {
        MetalDevice::device_count()
    }

    /// Get the underlying Metal device.
    pub fn metal_device(&self) -> &MetalDevice {
        &self.device
    }

    /// Get device info.
    pub fn info(&self) -> MetalDeviceInfo {
        self.device.info().clone()
    }

    /// Get device name.
    pub fn name(&self) -> &str {
        self.device.name()
    }

    /// Get the device index.
    pub fn device_index(&self) -> usize {
        self.index
    }

    /// Check if this device is an eGPU (external GPU).
    ///
    /// Returns `true` for Thunderbolt-connected external GPUs.
    pub fn is_egpu(&self) -> bool {
        self.device.is_egpu()
    }

    /// Synchronize all pending operations.
    pub fn synchronize(&self) {
        let _ = self.device.synchronize();
    }
}

impl Default for MetalDeviceWrapper {
    fn default() -> Self {
        Self::new().expect("No Metal device available")
    }
}

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

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

impl Eq for MetalDeviceWrapper {}

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

impl DeviceOps<MetalBackend> for MetalDeviceWrapper {
    fn id(&self) -> DeviceId {
        DeviceId::Metal(self.index)
    }

    fn memory_capacity(&self) -> usize {
        self.device.info().recommended_max_working_set_size as usize
    }

    fn memory_available(&self) -> usize {
        // Metal doesn't provide real-time memory queries easily
        // Return half of capacity as estimate
        self.memory_capacity() / 2
    }

    fn compute_capability(&self) -> Option<(u32, u32)> {
        // Metal doesn't use compute capability like CUDA
        // Return a version based on Apple GPU family
        Some((3, 0)) // Approximate Metal 3.0 support
    }

    fn synchronize(&self) {
        let _ = self.device.synchronize();
    }

    fn is_available(&self) -> bool {
        rtx_metal::is_available()
    }
}

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

    #[test]
    fn test_device_creation() {
        if rtx_metal::is_available() {
            let device = MetalDeviceWrapper::new();
            assert!(device.is_ok());

            let device = device.unwrap();
            assert!(device.is_available());
            println!("Metal device: {}", device.name());
        }
    }

    #[test]
    fn test_device_count() {
        let count = MetalDeviceWrapper::device_count();
        println!("Metal device count: {}", count);

        if rtx_metal::is_available() {
            assert!(
                count >= 1,
                "Should have at least one device when Metal is available"
            );
        }
    }

    #[test]
    fn test_device_by_index() {
        if rtx_metal::is_available() {
            // Device 0 should always work
            let device0 = MetalDeviceWrapper::new_with_index(0);
            assert!(device0.is_ok());

            let device = device0.unwrap();
            assert_eq!(device.device_index(), 0);
            println!("Device 0: {} (eGPU: {})", device.name(), device.is_egpu());

            // Check for additional devices (e.g., eGPU)
            let count = MetalDeviceWrapper::device_count();
            if count > 1 {
                let device1 = MetalDeviceWrapper::new_with_index(1);
                assert!(device1.is_ok());
                let device = device1.unwrap();
                println!("Device 1: {} (eGPU: {})", device.name(), device.is_egpu());
            }

            // Invalid index should fail
            let invalid = MetalDeviceWrapper::new_with_index(999);
            assert!(invalid.is_err());
        }
    }
}