Skip to main content

ferrotorch_core/
device.rs

1//! ## REQ status (per `.design/ferrotorch-core/device.md`)
2//!
3//! Tensor location enum mirroring `c10::Device` (`c10/core/Device.h:31`).
4//!
5//! | REQ | Status | Evidence |
6//! |---|---|---|
7//! | REQ-1 (Cpu variant) | SHIPPED | variant `Device::Cpu` at `device.rs:15` with `#[default]`; consumer `storage.rs` `TensorStorage::cpu(...).device() == Device::Cpu`; also `bool_tensor.rs:152` returns `Cpu` for any `TensorStorage::cpu`-backed `BoolTensor` |
8//! | REQ-2 (Cuda variant) | SHIPPED | variant `Device::Cuda(usize)` at `device.rs:18`; consumer `int_tensor.rs:268-323` `IntTensor::to` matches `(Cpu, Cuda(_))` / `(Cuda(_), Cpu)` arms for H2D / D2H transfer |
9//! | REQ-3 (Xpu variant) | SHIPPED | variant `Device::Xpu(usize)` at `device.rs:22`; consumer `error.rs:259` `FerrotorchError::DeviceMismatch { expected, got }` carries Xpu values; `int_tensor.rs:336` rejects `Xpu` destination via structured error |
10//! | REQ-4 (Mps variant) | SHIPPED | variant `Device::Mps(usize)` at `device.rs:26`; consumer `bool_tensor.rs:261-266` `(from, to) => Err(InvalidArgument)` arm pattern-matches on `Mps(_)` |
11//! | REQ-5 (Meta variant) | SHIPPED | variant `Device::Meta` at `device.rs:31`; consumer `storage.rs` `TensorStorage::Meta` arm — `try_as_slice` returns `GpuTensorNotAccessible` for Meta variant |
12//! | REQ-6 (predicates) | SHIPPED | `is_cpu` / `is_cuda` / `is_xpu` / `is_mps` / `is_meta` at `device.rs:36-64`; consumer `bool_tensor.rs:158`, `int_tensor.rs:205`, every `if a.device().is_cuda()` branch across `grad_fns/*.rs` |
13//! | REQ-7 (Display) | SHIPPED | `Display` impl at `device.rs:66-76` matching `c10::Device::str()` (`c10/core/Device.h:167`); consumer `error.rs:11` `#[error("device mismatch: expected {expected}, got {got}")]` |
14//! | REQ-8 (Copy/Hash derives) | SHIPPED | `#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]` at `device.rs:12`; consumer `gpu_dispatch.rs` registry + `Tensor<T>::device() == other.device()` PartialEq compares in `bool_tensor.rs:333`, `int_tensor.rs:436` |
15
16/// Device on which a tensor's data resides.
17///
18/// `Meta` is a special device that does not allocate any backing memory:
19/// meta tensors carry shape, dtype, and device information but no data.
20/// They are useful for shape inference, dry-run model construction, and
21/// inspecting parameter counts of huge models without actually allocating
22/// the weights. Mirrors `torch.device("meta")`.
23///
24/// `Xpu` mirrors PyTorch's `torch.device("xpu")` and addresses Intel
25/// GPUs (Arc series, Data Center GPU Max) via the portable CubeCL
26/// wgpu runtime that the `ferrotorch-xpu` crate wraps. CL-452.
27#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
28pub enum Device {
29    /// CPU main memory.
30    #[default]
31    Cpu,
32    /// CUDA GPU with the given device index.
33    Cuda(usize),
34    /// Intel XPU (Arc / Data Center GPU Max) with the given device index.
35    /// Accessed via `ferrotorch-xpu` which wraps a CubeCL wgpu runtime.
36    /// CL-452.
37    Xpu(usize),
38    /// Apple Silicon Metal Performance Shaders. The `usize` is the Metal
39    /// device index (`0` is the system default GPU). Mirrors
40    /// `torch.device("mps")`. Implemented via `ferrotorch-mps`. (#451)
41    Mps(usize),
42    /// Meta device — shape-only, no backing storage. Operations that need
43    /// data return an error; operations that only manipulate metadata
44    /// (reshape, view, permute, narrow, transpose, …) work normally and
45    /// produce meta tensors as output. CL-395.
46    Meta,
47}
48
49impl Device {
50    /// Returns `true` if this is a CPU device.
51    #[inline]
52    pub fn is_cpu(&self) -> bool {
53        matches!(self, Device::Cpu)
54    }
55
56    /// Returns `true` if this is a CUDA device.
57    #[inline]
58    pub fn is_cuda(&self) -> bool {
59        matches!(self, Device::Cuda(_))
60    }
61
62    /// Returns `true` if this is an Intel XPU device. CL-452.
63    #[inline]
64    pub fn is_xpu(&self) -> bool {
65        matches!(self, Device::Xpu(_))
66    }
67
68    /// Returns `true` if this is an Apple MPS device. (#451)
69    #[inline]
70    pub fn is_mps(&self) -> bool {
71        matches!(self, Device::Mps(_))
72    }
73
74    /// Returns `true` if this is the meta device (shape-only, no data).
75    #[inline]
76    pub fn is_meta(&self) -> bool {
77        matches!(self, Device::Meta)
78    }
79}
80
81impl core::fmt::Display for Device {
82    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
83        match self {
84            Device::Cpu => write!(f, "cpu"),
85            Device::Cuda(id) => write!(f, "cuda:{id}"),
86            Device::Xpu(id) => write!(f, "xpu:{id}"),
87            Device::Mps(id) => write!(f, "mps:{id}"),
88            Device::Meta => write!(f, "meta"),
89        }
90    }
91}