somatize_core/
device.rs

1//! Where a node runs.
2//!
3//! An enum and not a validated `String` so that a typo is an error **at
4//! declaration time**: `.on("cude:0")` fails where it was written, instead of
5//! turning into a torch `RuntimeError` halfway through a run.
6//!
7//! The price is that the vocabulary becomes ours rather than torch's, and it can
8//! be paid because **the core does not `match` on a `Device` anywhere else** —
9//! it only carries it to the node. Adding a variant is three lines and nowhere
10//! else stops compiling. Only the ones with a consumer today are here.
11
12use std::fmt;
13use std::str::FromStr;
14
15/// The place where a node executes. It travels **as text**, through
16/// [`Display`](fmt::Display) and [`FromStr`]: a variant number would break
17/// silently the day the enum grows in the middle.
18#[derive(Debug, Clone, PartialEq, Eq, Hash)]
19#[cfg_attr(
20    feature = "serde",
21    derive(serde::Serialize, serde::Deserialize),
22    serde(into = "String", try_from = "String")
23)]
24pub enum Device {
25    /// The processor.
26    Cpu,
27    /// A CUDA GPU, by index. Mandatory: bare `"cuda"` is thread state in torch,
28    /// and to whoever is placing that is not a placement.
29    Cuda(usize),
30    /// Torch's `meta` device: shape and dtype, without memory or compute. The
31    /// only one that proves a placement is obeyed on any machine.
32    Meta,
33}
34
35impl FromStr for Device {
36    type Err = DeviceError;
37
38    /// `cpu`, `cuda:0`, `meta`. Exactly as torch writes them, so what reaches
39    /// the node can be handed to `.to()` without translating anything.
40    fn from_str(s: &str) -> Result<Self, Self::Err> {
41        if s.is_empty() {
42            return Err(DeviceError::Malformed(s.to_string()));
43        }
44        let (kind, index) = match s.split_once(':') {
45            Some((kind, index)) => (kind, Some(index)),
46            None => (s, None),
47        };
48        match (kind, index) {
49            ("cpu", None) => Ok(Self::Cpu),
50            ("meta", None) => Ok(Self::Meta),
51            ("cuda", Some(index)) => index
52                .parse()
53                .map(Self::Cuda)
54                .map_err(|_| DeviceError::Malformed(s.to_string())),
55            ("cuda", None) => Err(DeviceError::NeedsIndex(kind.to_string())),
56            ("cpu" | "meta", Some(_)) => Err(DeviceError::Malformed(s.to_string())),
57            _ => Err(DeviceError::Unknown(kind.to_string())),
58        }
59    }
60}
61
62impl From<Device> for String {
63    fn from(device: Device) -> Self {
64        device.to_string()
65    }
66}
67
68impl TryFrom<String> for Device {
69    type Error = DeviceError;
70
71    fn try_from(s: String) -> Result<Self, Self::Error> {
72        s.parse()
73    }
74}
75
76impl fmt::Display for Device {
77    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
78        match self {
79            Self::Cpu => f.write_str("cpu"),
80            Self::Cuda(index) => write!(f, "cuda:{index}"),
81            Self::Meta => f.write_str("meta"),
82        }
83    }
84}
85
86/// Why that does not name a place to execute.
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub enum DeviceError {
89    /// We do not know that kind of device.
90    Unknown(String),
91    /// The kind is one of ours, but what comes with it is not.
92    Malformed(String),
93    /// It does not say which one.
94    NeedsIndex(String),
95}
96
97impl fmt::Display for DeviceError {
98    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
99        match self {
100            Self::Unknown(kind) => write!(
101                f,
102                "unknown device `{kind}`; today there are `cpu`, `cuda:N` and `meta`"
103            ),
104            Self::Malformed(s) => write!(
105                f,
106                "`{s}` is not shaped like a device; write `cpu`, `cuda:N` or `meta`"
107            ),
108            Self::NeedsIndex(kind) => write!(
109                f,
110                "`{kind}` does not say which one: write `{kind}:0`. \"The current one\" \
111                 is thread state, not a placement"
112            ),
113        }
114    }
115}
116
117impl std::error::Error for DeviceError {}
somatize_core/device.rs

somatize_core/
device.rs
