gpu_async_reactor/backend/mapped.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! The parent of the CUDA and HIP lines: one page-locked, device-mapped host page holding one
4//! [`CompletionLine`], which a stream stores sequence numbers into behind its work. Each vendor
5//! supplies only its driver table ([`LineDriver`]); the ownership rules live here once:
6//! - the line holds its own reference to its domain (the CUDA primary context, or the HIP
7//! device's primary context) from before the page is allocated until after it is freed, so no
8//! other owner letting the domain go can free the page under the line (CWE-416; 0.1.0 freed it
9//! with a context the line did not keep alive);
10//! - a store is accepted only from a stream of that domain, so `Drop` can wait for it;
11//! - `Drop` frees the page only after that domain settled, then releases the domain; if it cannot
12//! settle, the page and the reference are leaked, never freed under a pending store;
13//! - the host reads the line only through atomics; the device writes only the words named.
14
15#![allow(unsafe_code)]
16
17use core::cell::Cell;
18use core::ffi::c_void;
19use core::ptr::NonNull;
20use core::sync::atomic::Ordering;
21
22use crate::core::line::{CompletionLine, Slot};
23
24/// The page the line lives in: a whole page, so no other allocation shares the line with a
25/// device writer.
26pub(crate) const PAGE: usize = 4096;
27
28/// Why a line could not be made or a store not queued.
29#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30pub enum LineError {
31 /// The driver returned this nonzero status.
32 Status(i32),
33 /// The driver library or one of its entry points is missing.
34 Unavailable,
35 /// No current context/device on the calling thread.
36 NoDomain,
37 /// The stream belongs to another context/device than the line's.
38 ForeignStream,
39}
40
41impl LineError {
42 /// One integer per error, as the line's first users logged them: the driver's status, `-1`
43 /// unavailable, `-3` no domain, `-4` foreign stream.
44 #[must_use]
45 pub fn status(self) -> i32 {
46 match self {
47 Self::Status(s) => s,
48 Self::Unavailable => -1,
49 Self::NoDomain => -3,
50 Self::ForeignStream => -4,
51 }
52 }
53}
54
55impl core::fmt::Display for LineError {
56 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
57 match self {
58 Self::Status(s) => write!(f, "driver status {s}"),
59 Self::Unavailable => f.write_str("the driver or one of its entry points is missing"),
60 Self::NoDomain => f.write_str("no current context on this thread"),
61 Self::ForeignStream => f.write_str("the stream belongs to another context"),
62 }
63 }
64}
65
66impl core::error::Error for LineError {}
67
68impl From<LineError> for crate::Error {
69 fn from(e: LineError) -> Self {
70 crate::error!(Backend, "completion line", e.status())
71 }
72}
73
74/// What a line needs of a vendor driver. Crate-private: the line is its only caller and passes
75/// only values the driver handed out.
76pub(crate) trait LineDriver {
77 /// The vendor's stream handle.
78 type Stream: Copy;
79 /// A domain the line holds a reference to (the vendor handles needed to use and release it).
80 type Domain: Copy;
81 /// Take a reference to the domain of the calling thread's current context or device, which
82 /// keeps it alive until [`Self::release`]; [`LineError::NoDomain`] when none is current.
83 fn retain_current(&self) -> Result<Self::Domain, LineError>;
84 /// Drop one reference from [`Self::retain_current`], once.
85 fn release(&self, domain: Self::Domain);
86 /// A page-locked, device-mapped host page of `bytes`, owned by `domain`.
87 fn alloc_host(&self, domain: Self::Domain, bytes: usize) -> Result<NonNull<c_void>, LineError>;
88 /// The device address of `host`, a page from [`Self::alloc_host`] not yet freed.
89 fn device_address(&self, domain: Self::Domain, host: NonNull<c_void>) -> Result<u64, LineError>;
90 /// Free a page from [`Self::alloc_host`], once, while `domain` is still retained.
91 fn free_host(&self, domain: Self::Domain, host: NonNull<c_void>);
92 /// Whether `stream` belongs to `domain`.
93 fn stream_in(&self, stream: Self::Stream, domain: Self::Domain) -> Result<bool, LineError>;
94 /// Queue a store of `value` at device address `addr` on `stream`.
95 fn write_value(&self, stream: Self::Stream, addr: u64, value: u64) -> Result<(), LineError>;
96 /// Wait for every queued work of `domain`; `false` if the driver could not.
97 fn settle(&self, domain: Self::Domain) -> bool;
98}
99
100/// A device-mapped host line over driver `D`.
101pub(crate) struct Line<D: LineDriver> {
102 host: NonNull<CompletionLine>,
103 device: u64,
104 driver: D,
105 domain: D::Domain,
106 queued: Cell<bool>,
107}
108
109// WHY: a line owns a raw mapped page, which withholds the auto `Send`; the executor moves the owner between threads.
110// SAFETY: the page and the domain reference are owned by this value alone and the host touches
111// the page only through atomics (which are `Sync`); the driver table is `Send` by the bound, and a
112// domain reference is a plain handle every driver call takes explicitly (no thread-current
113// state). `queued` is a `Cell`, so the value is `Send` but not `Sync`, which is all a line needs.
114// NO SAFE ALTERNATIVE: a raw-pointer owner is never `Send` by derivation, and wrapping the page would add an indirection on every host read.
115unsafe impl<D: LineDriver + Send> Send for Line<D> {}
116
117impl<D: LineDriver> Line<D> {
118 /// A fresh line reading all 0, mapped for (and holding a reference to) the calling thread's
119 /// current domain.
120 pub(crate) fn with_driver(driver: D) -> Result<Self, LineError> {
121 let domain = driver.retain_current()?;
122 let page = match driver.alloc_host(domain, PAGE) {
123 Ok(page) => page,
124 Err(e) => {
125 driver.release(domain);
126 return Err(e);
127 }
128 };
129 // Built before the last check, so its error frees the page and releases through `Drop`.
130 let mut line = Self {
131 host: page.cast::<CompletionLine>(),
132 device: 0,
133 driver,
134 domain,
135 queued: Cell::new(false),
136 };
137 line.device = line.driver.device_address(domain, page)?;
138 for w in &line.line().words {
139 w.store(0, Ordering::Release);
140 }
141 Ok(line)
142 }
143
144 /// The driver table.
145 #[cfg(feature = "cuda")]
146 pub(crate) fn driver(&self) -> &D {
147 &self.driver
148 }
149
150 /// The device address of word 0 (for a stream-side wait on the line).
151 pub(crate) fn device_address(&self) -> u64 {
152 self.device
153 }
154
155 /// The line, as the host reads it.
156 pub(crate) fn line(&self) -> &CompletionLine {
157 // WHY: reads the owned mapped page as a `CompletionLine` without copying it.
158 // SAFETY: `host` is the live, page-aligned (so 64-aligned, as `CompletionLine` needs) page
159 // this value owns until `Drop`, in a domain this value keeps alive until then; the device
160 // writes it only with aligned 8-byte stores to its words and every host access goes
161 // through the atomics, so nothing aliases it non-atomically.
162 // NO SAFE ALTERNATIVE: a `NonNull<T>` deref is unsafe; the page must stay in place because the device writes it.
163 unsafe { self.host.as_ref() }
164 }
165
166 /// Queue a store of `value` into word `slot` on `stream`, behind the work queued there.
167 pub(crate) fn enqueue_store(
168 &self,
169 stream: D::Stream,
170 slot: Slot,
171 value: u64,
172 ) -> Result<(), LineError> {
173 let offset = 8 * u64::try_from(slot.index()).map_err(|_| LineError::Unavailable)?;
174 if !self.driver.stream_in(stream, self.domain)? {
175 return Err(LineError::ForeignStream);
176 }
177 // Set before the call: a store the driver queued despite a failing status still makes
178 // `Drop` settle first.
179 self.queued.set(true);
180 self.driver.write_value(stream, self.device + offset, value)
181 }
182
183 /// Every store queued into the line has landed (`true` at once when none was queued).
184 fn settled(&self) -> bool {
185 !self.queued.get() || self.driver.settle(self.domain)
186 }
187}
188
189impl<D: LineDriver> Drop for Line<D> {
190 fn drop(&mut self) {
191 if !self.settled() {
192 // A store may still be pending into the page: leak 4 KiB (and the domain reference
193 // that keeps the mapping valid) rather than free memory the device may yet write.
194 return;
195 }
196 self.driver.free_host(self.domain, self.host.cast::<c_void>());
197 self.driver.release(self.domain);
198 }
199}
200
201#[cfg(test)]
202#[path = "mapped_tests.rs"]
203mod tests;