Skip to main content

baracuda_runtime/
device.rs

1//! Device enumeration + queries via the Runtime API.
2//!
3//! Unlike [`baracuda_driver::Device`], a [`Device`] in the Runtime API is
4//! just an ordinal — there's no separate `CUdevice` handle. Contexts are
5//! implicit (the "primary context" per device) and are shared with the
6//! Driver API on the same device.
7
8use baracuda_cuda_sys::runtime::{runtime, types::cudaDeviceAttr as Attr};
9
10use crate::error::{check, Result};
11
12/// A CUDA device (Runtime API view — a bare ordinal).
13#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
14pub struct Device(pub(crate) i32);
15
16impl Device {
17    /// Number of CUDA devices visible to the process.
18    pub fn count() -> Result<u32> {
19        let r = runtime()?;
20        let cu = r.cuda_get_device_count()?;
21        let mut n: core::ffi::c_int = 0;
22        check(unsafe { cu(&mut n) })?;
23        Ok(n as u32)
24    }
25
26    /// Construct a `Device` for the given ordinal. Does not validate — use
27    /// [`Device::all`] if you want a checked enumeration.
28    #[inline]
29    pub const fn from_ordinal(ordinal: u32) -> Self {
30        Self(ordinal as i32)
31    }
32
33    /// All visible devices, in ordinal order.
34    pub fn all() -> Result<Vec<Self>> {
35        let count = Self::count()?;
36        Ok((0..count).map(Self::from_ordinal).collect())
37    }
38
39    /// Set this device as current on the calling thread. Subsequent Runtime
40    /// API calls (allocations, launches, ...) operate on this device.
41    pub fn set_current(&self) -> Result<()> {
42        let r = runtime()?;
43        let cu = r.cuda_set_device()?;
44        check(unsafe { cu(self.0) })
45    }
46
47    /// Retrieve the device currently selected on the calling thread.
48    pub fn current() -> Result<Self> {
49        let r = runtime()?;
50        let cu = r.cuda_get_device()?;
51        let mut dev: core::ffi::c_int = 0;
52        check(unsafe { cu(&mut dev) })?;
53        Ok(Self(dev))
54    }
55
56    /// Ordinal of this device (`0`, `1`, ...).
57    #[inline]
58    pub fn ordinal(&self) -> i32 {
59        self.0
60    }
61
62    /// Compute capability as `(major, minor)`.
63    pub fn compute_capability(&self) -> Result<(u32, u32)> {
64        Ok((
65            self.attribute(Attr::COMPUTE_CAPABILITY_MAJOR)? as u32,
66            self.attribute(Attr::COMPUTE_CAPABILITY_MINOR)? as u32,
67        ))
68    }
69
70    /// Multiprocessor count.
71    pub fn multiprocessor_count(&self) -> Result<u32> {
72        Ok(self.attribute(Attr::MULTIPROCESSOR_COUNT)? as u32)
73    }
74
75    /// Warp size in threads.
76    pub fn warp_size(&self) -> Result<u32> {
77        Ok(self.attribute(Attr::WARP_SIZE)? as u32)
78    }
79
80    /// Raw device-attribute query. See [`baracuda_cuda_sys::runtime::types::cudaDeviceAttr`].
81    pub fn attribute(&self, attr: i32) -> Result<i32> {
82        let r = runtime()?;
83        let cu = r.cuda_device_get_attribute()?;
84        let mut val: core::ffi::c_int = 0;
85        check(unsafe { cu(&mut val, attr, self.0) })?;
86        Ok(val)
87    }
88
89    /// `true` if this device can peer-access `peer`'s allocations (P2P).
90    /// Call [`Device::enable_peer_access`] before actually using peer
91    /// pointers in kernels.
92    pub fn can_access_peer(&self, peer: &Device) -> Result<bool> {
93        let r = runtime()?;
94        let cu = r.cuda_device_can_access_peer()?;
95        let mut v: core::ffi::c_int = 0;
96        check(unsafe { cu(&mut v, self.0, peer.0) })?;
97        Ok(v != 0)
98    }
99
100    /// Enable peer access from the *current* device to `peer`'s
101    /// allocations. Call `Device::set_current()` on the accessing device
102    /// first.
103    pub fn enable_peer_access(peer: &Device) -> Result<()> {
104        let r = runtime()?;
105        let cu = r.cuda_device_enable_peer_access()?;
106        check(unsafe { cu(peer.0, 0) })
107    }
108
109    /// Disable peer access previously enabled via
110    /// [`Device::enable_peer_access`].
111    pub fn disable_peer_access(peer: &Device) -> Result<()> {
112        let r = runtime()?;
113        let cu = r.cuda_device_disable_peer_access()?;
114        check(unsafe { cu(peer.0) })
115    }
116}