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

//! Device abstraction for backends.
//!
//! Each backend defines its own device type that implements [`DeviceOps`].

use crate::Backend;
use serde::{Deserialize, Serialize};
use std::fmt::{Debug, Display};
use std::hash::Hash;

/// Device identifier that can be serialized and compared.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum DeviceId {
    /// CPU device.
    Cpu,
    /// CUDA GPU with index.
    Cuda(usize),
    /// ROCm (AMD) GPU with index.
    Rocm(usize),
    /// Metal (Apple) GPU with index.
    Metal(usize),
    /// SYCL (Intel) GPU with index.
    Sycl(usize),
    /// WebGPU device with index.
    WebGpu(usize),
}

impl DeviceId {
    /// Check if this is a CPU device.
    #[inline]
    pub fn is_cpu(&self) -> bool {
        matches!(self, DeviceId::Cpu)
    }

    /// Check if this is a GPU device.
    #[inline]
    pub fn is_gpu(&self) -> bool {
        !self.is_cpu()
    }

    /// Get the device index (None for CPU).
    #[inline]
    pub fn index(&self) -> Option<usize> {
        match self {
            DeviceId::Cpu => None,
            DeviceId::Cuda(idx)
            | DeviceId::Rocm(idx)
            | DeviceId::Metal(idx)
            | DeviceId::Sycl(idx)
            | DeviceId::WebGpu(idx) => Some(*idx),
        }
    }

    /// Get the device type as a string.
    #[inline]
    pub fn device_type(&self) -> &'static str {
        match self {
            DeviceId::Cpu => "cpu",
            DeviceId::Cuda(_) => "cuda",
            DeviceId::Rocm(_) => "rocm",
            DeviceId::Metal(_) => "metal",
            DeviceId::Sycl(_) => "sycl",
            DeviceId::WebGpu(_) => "webgpu",
        }
    }
}

impl Display for DeviceId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DeviceId::Cpu => write!(f, "cpu"),
            DeviceId::Cuda(idx) => write!(f, "cuda:{}", idx),
            DeviceId::Rocm(idx) => write!(f, "rocm:{}", idx),
            DeviceId::Metal(idx) => write!(f, "metal:{}", idx),
            DeviceId::Sycl(idx) => write!(f, "sycl:{}", idx),
            DeviceId::WebGpu(idx) => write!(f, "webgpu:{}", idx),
        }
    }
}

impl Default for DeviceId {
    fn default() -> Self {
        // Default to first GPU if available, otherwise CPU
        // Actual availability is checked by backend
        DeviceId::Cuda(0)
    }
}

/// Device operations trait for backends.
///
/// Each backend provides its own device type that manages
/// GPU contexts, memory allocators, and compute streams.
///
/// # Example Implementation
///
/// ```rust,ignore
/// #[derive(Clone, Debug)]
/// pub struct CudaDevice {
///     context: Arc<CudaContext>,
///     stream: Arc<CudaStream>,
///     index: usize,
/// }
///
/// impl DeviceOps<CudaBackend> for CudaDevice {
///     fn id(&self) -> DeviceId {
///         DeviceId::Cuda(self.index)
///     }
///
///     fn default() -> Self {
///         CudaDevice::new(0).expect("No CUDA device available")
///     }
/// }
/// ```
pub trait DeviceOps<B: Backend>:
    Clone + Debug + Default + PartialEq + Eq + Hash + Send + Sync + 'static
{
    /// Get the device identifier.
    fn id(&self) -> DeviceId;

    /// Get the device type string.
    #[inline]
    fn device_type(&self) -> &'static str {
        self.id().device_type()
    }

    /// Check if this is a CPU device.
    #[inline]
    fn is_cpu(&self) -> bool {
        self.id().is_cpu()
    }

    /// Check if this is a GPU device.
    #[inline]
    fn is_gpu(&self) -> bool {
        self.id().is_gpu()
    }

    /// Get the device index (None for CPU).
    #[inline]
    fn index(&self) -> Option<usize> {
        self.id().index()
    }

    /// Get total memory capacity in bytes.
    fn memory_capacity(&self) -> usize;

    /// Get currently available memory in bytes.
    fn memory_available(&self) -> usize;

    /// Get compute capability (major, minor) for GPU devices.
    fn compute_capability(&self) -> Option<(u32, u32)>;

    /// Synchronize all pending operations on this device.
    fn synchronize(&self);

    /// Check if this device is available.
    fn is_available(&self) -> bool;
}

/// Helper trait for creating devices from configuration.
///
/// This trait is designed for future use when backends implement device discovery.
#[allow(dead_code)]
pub trait DeviceFactory: DeviceOps<Self::Backend> + Sized {
    /// The backend this device factory is for.
    type Backend: Backend;

    /// Create a device from a device ID.
    fn from_id(id: DeviceId) -> Option<Self>;

    /// Create the default device.
    fn default_device() -> Self;

    /// List all available devices.
    fn available_devices() -> Vec<Self>;
}

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

    #[test]
    fn test_device_id_display() {
        assert_eq!(DeviceId::Cpu.to_string(), "cpu");
        assert_eq!(DeviceId::Cuda(0).to_string(), "cuda:0");
        assert_eq!(DeviceId::Metal(1).to_string(), "metal:1");
    }

    #[test]
    fn test_device_id_is_gpu() {
        assert!(!DeviceId::Cpu.is_gpu());
        assert!(DeviceId::Cuda(0).is_gpu());
        assert!(DeviceId::Metal(0).is_gpu());
        assert!(DeviceId::Sycl(0).is_gpu());
        assert!(DeviceId::WebGpu(0).is_gpu());
    }

    #[test]
    fn test_sycl_device_id() {
        let sycl = DeviceId::Sycl(1);
        assert_eq!(sycl.to_string(), "sycl:1");
        assert_eq!(sycl.device_type(), "sycl");
        assert_eq!(sycl.index(), Some(1));
    }
}