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/// Runs `f` with `stream` as this thread's default stream, and restores the previous
96/// default afterwards, also when `f` panics. A device stands for its first stream.
97///
98/// This replaces per-operation stream arguments: every operation `f` records goes on
99/// `stream`.
100///
101/// ```
102/// use nuvec_core::{Device, array, default_stream, with_stream};
103///
104/// let x = array![1.0, 2.0];
105/// let y = with_stream(Device::cpu(), || x.exp().sum(..));
106/// assert_eq!(y.device(), Device::cpu());
107/// assert_eq!(default_stream(), nuvec_core::default_stream()); // restored
108/// ```
109pub fn with_stream<R>(stream: impl Into<Stream>, f: impl FnOnce() -> R) -> R {
110    let _guard = StreamGuard::set(stream.into());
111    f()
112}
113
114/// Restores the thread's default stream when dropped, also when unwinding.
115pub(crate) struct StreamGuard(Option<Stream>);
116
117impl StreamGuard {
118    /// Makes `stream` the default until the guard is dropped.
119    pub(crate) fn set(stream: Stream) -> Self {
120        let guard = Self(DEFAULT_STREAM.get());
121        DEFAULT_STREAM.set(Some(stream));
122        guard
123    }
124}
125
126impl Drop for StreamGuard {
127    fn drop(&mut self) {
128        DEFAULT_STREAM.set(self.0);
129    }
130}
131
132/// A device's first stream.
133impl From<Device> for Stream {
134    fn from(device: Device) -> Self {
135        Self::default_for(device)
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn default_is_the_default_device_and_settable_per_thread() {
145        assert_eq!(default_stream().device, default_device());
146        let expected = if cfg!(all(target_os = "macos", not(miri))) && crate::gpu::available() {
147            Device::gpu(0)
148        } else {
149            Device::cpu()
150        };
151        assert_eq!(default_device(), expected);
152        set_default_stream(Stream::default_for(Device::cpu()));
153        assert_eq!(default_stream().device, Device::cpu());
154        std::thread::spawn(|| assert_eq!(default_stream().device, default_device()))
155            .join()
156            .unwrap();
157    }
158
159    #[test]
160    fn display() {
161        assert_eq!(Device::gpu(1).to_string(), "gpu:1");
162    }
163}
nuvec_core/device.rs

nuvec_core/
device.rs
