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<Stream> = const { Cell::new(Stream::default_for(Device::cpu())) };
69}
70
71/// The stream new operations on this thread are recorded on.
72pub fn default_stream() -> Stream {
73    DEFAULT_STREAM.get()
74}
75
76/// Sets the stream new operations on this thread are recorded on.
77pub fn set_default_stream(stream: Stream) {
78    DEFAULT_STREAM.set(stream);
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    #[test]
86    fn default_is_cpu_and_settable_per_thread() {
87        assert_eq!(default_stream().device, Device::cpu());
88        set_default_stream(Stream::default_for(Device::gpu(0)));
89        assert_eq!(default_stream().device, Device::gpu(0));
90        std::thread::spawn(|| assert_eq!(default_stream().device, Device::cpu()))
91            .join()
92            .unwrap();
93        set_default_stream(Stream::default_for(Device::cpu()));
94    }
95
96    #[test]
97    fn display() {
98        assert_eq!(Device::gpu(1).to_string(), "gpu:1");
99    }
100}
nuvec_core/device.rs

nuvec_core/
device.rs
