nuvec_core/
device.rs

1//! Devices and streams.
2
3use std::cell::Cell;
4use std::fmt;
5
6/// The kind of processor a device is.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
8pub enum DeviceKind {
9    /// The CPU.
10    Cpu,
11    /// A GPU.
12    Gpu,
13}
14
15/// A processor that executes array operations.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
17pub struct Device {
18    /// The kind of processor.
19    pub kind: DeviceKind,
20    /// Index among devices of the same kind.
21    pub index: u32,
22}
23
24impl Device {
25    /// The CPU.
26    pub const fn cpu() -> Self {
27        Self {
28            kind: DeviceKind::Cpu,
29            index: 0,
30        }
31    }
32
33    /// The GPU with the given index.
34    pub const fn gpu(index: u32) -> Self {
35        Self {
36            kind: DeviceKind::Gpu,
37            index,
38        }
39    }
40}
41
42impl fmt::Display for Device {
43    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44        match self.kind {
45            DeviceKind::Cpu => write!(f, "cpu:{}", self.index),
46            DeviceKind::Gpu => write!(f, "gpu:{}", self.index),
47        }
48    }
49}
50
51/// An ordered queue of operations on a device.
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
53pub struct Stream {
54    /// Index among the device's streams.
55    pub index: u32,
56    /// The device the stream executes on.
57    pub device: Device,
58}
59
60impl Stream {
61    /// The first stream of `device`.
62    pub const fn default_for(device: Device) -> Self {
63        Self { index: 0, device }
64    }
65}
66
67thread_local! {
68    static DEFAULT_STREAM: Cell<Option<Stream>> = const { Cell::new(None) };
69}
70
71/// The default device: the GPU when Metal is available, else the CPU (as in MLX).
72pub fn default_device() -> Device {
73    if crate::gpu::available() {
74        Device::gpu(0)
75    } else {
76        Device::cpu()
77    }
78}
79
80/// The stream new operations on this thread are recorded on: the default device's first
81/// stream until [`set_default_stream`] is called on this thread.
82pub fn default_stream() -> Stream {
83    DEFAULT_STREAM.get().unwrap_or_else(|| {
84        let stream = Stream::default_for(default_device());
85        DEFAULT_STREAM.set(Some(stream));
86        stream
87    })
88}
89
90/// Sets the stream new operations on this thread are recorded on.
91pub fn set_default_stream(stream: Stream) {
92    DEFAULT_STREAM.set(Some(stream));
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn default_is_the_default_device_and_settable_per_thread() {
101        assert_eq!(default_stream().device, default_device());
102        let expected = if cfg!(all(target_os = "macos", not(miri))) && crate::gpu::available() {
103            Device::gpu(0)
104        } else {
105            Device::cpu()
106        };
107        assert_eq!(default_device(), expected);
108        set_default_stream(Stream::default_for(Device::cpu()));
109        assert_eq!(default_stream().device, Device::cpu());
110        std::thread::spawn(|| assert_eq!(default_stream().device, default_device()))
111            .join()
112            .unwrap();
113    }
114
115    #[test]
116    fn display() {
117        assert_eq!(Device::gpu(1).to_string(), "gpu:1");
118    }
119}
nuvec_core/device.rs

nuvec_core/
device.rs
