henad_compute/gpu/primitives/
readback.rs

1//! Async readback of some `u32` counters reduced on the GPU.
2//!
3//! A GPU model answers `SimState::stats()` from these counters without copying its state back to the CPU. The map
4//! starts right after submission and completes on a later loop iteration, and [`CounterReadback::poll`] never blocks.
5//! A reported value is therefore a few milliseconds stale, as the display texture already is.
6
7use std::mem::size_of;
8
9/// State of a stats readback after a poll.
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum StatsPoll {
12    /// A readback has begun and not finished.
13    Pending,
14    /// No readback is in flight, and either no readback has begun or the latest one to finish succeeded.
15    Landed,
16    /// No readback is in flight, and the latest one to finish failed and left the earlier values in place.
17    Failed,
18}
19
20/// A GPU-side `u32` accumulator of `count` counters, plus the staging buffer used to read it back
21/// without blocking.
22#[derive(Debug)]
23pub struct CounterReadback {
24    /// The reduce shader's output. Cleared to 0 each time, accumulated into, then copied out.
25    storage: wgpu::Buffer,
26    staging: wgpu::Buffer,
27    /// `Some` while a `map_async` is in flight.
28    pending: Option<flume::Receiver<Result<(), wgpu::BufferAsyncError>>>,
29    /// Set once a fresh value is in `staging` and waiting to be mapped.
30    copied: bool,
31    /// Whether the latest map to finish failed.
32    failed: bool,
33    values: Vec<u32>,
34}
35
36impl CounterReadback {
37    /// Creates the accumulator and its staging buffer for `count` counters.
38    ///
39    /// `count` must match the length of the `atomic<u32>` array the reduce shader declares.
40    pub fn new(device: &wgpu::Device, label: &str, count: usize) -> Self {
41        let size = (count * size_of::<u32>()) as u64;
42        let storage = device.create_buffer(&wgpu::BufferDescriptor {
43            label: Some(&format!("{label}_storage")),
44            size,
45            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
46            mapped_at_creation: false,
47        });
48        let staging = device.create_buffer(&wgpu::BufferDescriptor {
49            label: Some(&format!("{label}_staging")),
50            size,
51            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
52            mapped_at_creation: false,
53        });
54        Self {
55            storage,
56            staging,
57            pending: None,
58            copied: false,
59            failed: false,
60            values: vec![0; count],
61        }
62    }
63
64    fn size(&self) -> u64 {
65        (self.values.len() * size_of::<u32>()) as u64
66    }
67
68    /// Returns the accumulator, to bind as the reduce shader's `read_write` storage target.
69    pub fn binding(&self) -> wgpu::BindingResource<'_> {
70        self.storage.as_entire_binding()
71    }
72
73    /// Clears the accumulator to zero. Record it before the model's reduce pass.
74    pub fn encode_clear(&self, encoder: &mut wgpu::CommandEncoder) {
75        encoder.clear_buffer(&self.storage, 0, None);
76    }
77
78    /// Copies the accumulated totals into the staging buffer. Record it after the model's
79    /// reduce pass, in the same encoder, where wgpu inserts the barrier between the pass and the copy.
80    ///
81    /// Note that the copy is skipped while a previous map is still in flight, since writing into a buffer that is
82    /// mapped or pending a map is invalid. The sample is then dropped, and the stats repeat the older values with no
83    /// error.
84    pub fn encode_copy(&mut self, encoder: &mut wgpu::CommandEncoder) {
85        if self.pending.is_some() {
86            return;
87        }
88        encoder.copy_buffer_to_buffer(&self.storage, 0, &self.staging, 0, self.size());
89        self.copied = true;
90    }
91
92    /// Starts the async map. Call it once, immediately after submitting the encoder that
93    /// [`Self::encode_copy`] was recorded into. Mapping any earlier races the copy.
94    pub fn begin_map(&mut self) {
95        if self.pending.is_some() || !self.copied {
96            return;
97        }
98        self.copied = false;
99        let (tx, rx) = flume::bounded(1);
100        self.staging.slice(..).map_async(wgpu::MapMode::Read, move |result| {
101            drop(tx.send(result));
102        });
103        self.pending = Some(rx);
104    }
105
106    /// Returns whether a map started by [`Self::begin_map`] has not been consumed yet.
107    pub fn is_pending(&self) -> bool {
108        self.pending.is_some()
109    }
110
111    /// Consumes the in-flight map if it has completed, updating [`Self::values`], and returns the state of the
112    /// readback after the poll. It never blocks.
113    ///
114    /// On native, `device.poll` runs wgpu's map callbacks. Call this on every loop iteration, whether or not a value
115    /// is expected.
116    pub fn poll(&mut self, device: &wgpu::Device) -> StatsPoll {
117        let Some(rx) = self.pending.as_ref() else {
118            return self.status();
119        };
120
121        drop(device.poll(wgpu::PollType::Poll));
122
123        match rx.try_recv() {
124            Ok(result) => {
125                self.pending = None;
126                self.failed = !self.finish_map(result);
127            }
128            // The callback was dropped without running, so no value can arrive.
129            Err(flume::TryRecvError::Disconnected) => self.abandon_map(),
130            Err(flume::TryRecvError::Empty) => {}
131        }
132        self.status()
133    }
134
135    /// Waits for the GPU to drain, then consumes the map, and returns the state of the readback after the wait. It is
136    /// the blocking counterpart of [`Self::poll`].
137    ///
138    /// It is meant for one-shot snapshots and samples. Never call it from the hot batching loop.
139    #[cfg(not(target_arch = "wasm32"))]
140    pub fn poll_blocking(&mut self, device: &wgpu::Device) -> StatsPoll {
141        if let Some(rx) = self.pending.as_ref() {
142            let received = match device.poll(wgpu::PollType::wait_indefinitely()) {
143                Ok(_) => rx.recv().ok(),
144                // The map can have finished before the wait failed.
145                Err(_) => rx.try_recv().ok(),
146            };
147            match received {
148                Some(result) => {
149                    self.pending = None;
150                    self.failed = !self.finish_map(result);
151                }
152                None => self.abandon_map(),
153            }
154        }
155        self.status()
156    }
157
158    /// Polls once without waiting. WebGPU's `poll` is a no-op and the map resolves on the JS microtask
159    /// queue, so a wait here hangs the tab. The value arrives on a later [`Self::poll`] instead.
160    #[cfg(target_arch = "wasm32")]
161    pub fn poll_blocking(&mut self, device: &wgpu::Device) -> StatsPoll {
162        self.poll(device)
163    }
164
165    fn status(&self) -> StatsPoll {
166        if self.pending.is_some() {
167            StatsPoll::Pending
168        } else if self.failed {
169            StatsPoll::Failed
170        } else {
171            StatsPoll::Landed
172        }
173    }
174
175    /// Gives up the map in flight as failed, and unmaps the staging buffer.
176    ///
177    /// Otherwise the buffer stays mapped or pending, and the next [`Self::encode_copy`] records a copy into it.
178    fn abandon_map(&mut self) {
179        self.pending = None;
180        self.failed = true;
181        self.staging.unmap();
182    }
183
184    /// Reads and unmaps the staging buffer after a completed `map_async`.
185    fn finish_map(&mut self, result: Result<(), wgpu::BufferAsyncError>) -> bool {
186        if let Err(err) = result {
187            log::warn!("GPU stat readback failed: {err}");
188            return false;
189        }
190
191        let slice = self.staging.slice(..);
192        // Unmap on the error path too, or the next `map_async` finds the buffer still mapped.
193        let data = match slice.get_mapped_range() {
194            Ok(data) => data,
195            Err(err) => {
196                log::warn!("GPU stat readback failed: {err}");
197                self.staging.unmap();
198                return false;
199            }
200        };
201        let words: &[u32] = bytemuck::cast_slice(&data);
202        let n = self.values.len();
203        self.values.copy_from_slice(&words[..n]);
204        drop(data);
205        self.staging.unmap();
206
207        true
208    }
209
210    /// Most recently read-back values, all zero until the first readback completes.
211    pub fn values(&self) -> &[u32] {
212        &self.values
213    }
214
215    /// Most recently read-back values, each word's bits reinterpreted as an `f32`, for
216    /// [`crate::gpu::primitives::reduce`].
217    pub fn values_f32(&self) -> impl Iterator<Item = f32> + '_ {
218        self.values.iter().copied().map(f32::from_bits)
219    }
220}
221
222#[cfg(all(test, not(target_arch = "wasm32")))]
223mod tests {
224    use super::{CounterReadback, StatsPoll};
225    use crate::gpu::{GpuContext, headless_context};
226
227    /// Writes `values` into the accumulator, submits their copy to the staging buffer and begins the map.
228    fn begin_read_back(ctx: &GpuContext, readback: &mut CounterReadback, values: &[u32]) {
229        ctx.queue
230            .write_buffer(&readback.storage, 0, bytemuck::cast_slice(values));
231        let mut encoder = ctx
232            .device
233            .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
234        readback.encode_copy(&mut encoder);
235        ctx.queue.submit(Some(encoder.finish()));
236        readback.begin_map();
237    }
238
239    #[test]
240    fn a_completed_map_polls_as_landed() {
241        let Some(ctx) = headless_context("henad_readback_landed_test", wgpu::Features::empty()) else {
242            log::warn!("skipping a_completed_map_polls_as_landed: no adapter");
243            return;
244        };
245        let mut readback = CounterReadback::new(&ctx.device, "henad_readback_landed_test", 3);
246
247        begin_read_back(&ctx, &mut readback, &[1, 2, 3]);
248        assert!(readback.is_pending(), "a begun map is pending");
249        assert_eq!(readback.poll_blocking(&ctx.device), StatsPoll::Landed);
250        assert_eq!(readback.values(), [1, 2, 3]);
251
252        begin_read_back(&ctx, &mut readback, &[4, 5, 6]);
253        ctx.device
254            .poll(wgpu::PollType::wait_indefinitely())
255            .expect("the copy runs");
256        assert_eq!(
257            readback.poll(&ctx.device),
258            StatsPoll::Landed,
259            "a map the device completed lands on the next poll"
260        );
261        assert_eq!(readback.values(), [4, 5, 6]);
262        assert_eq!(
263            readback.poll(&ctx.device),
264            StatsPoll::Landed,
265            "a poll with nothing in flight reports the latest map"
266        );
267    }
268
269    #[test]
270    fn a_failed_map_polls_as_failed_and_keeps_the_values() {
271        let Some(ctx) = headless_context("henad_readback_failed_test", wgpu::Features::empty()) else {
272            log::warn!("skipping a_failed_map_polls_as_failed_and_keeps_the_values: no adapter");
273            return;
274        };
275        let mut readback = CounterReadback::new(&ctx.device, "henad_readback_failed_test", 3);
276        begin_read_back(&ctx, &mut readback, &[1, 2, 3]);
277        assert_eq!(readback.poll_blocking(&ctx.device), StatsPoll::Landed);
278
279        // A destroyed staging buffer rejects the map.
280        let mut encoder = ctx
281            .device
282            .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
283        readback.encode_copy(&mut encoder);
284        ctx.queue.submit(Some(encoder.finish()));
285        readback.staging.destroy();
286        readback.begin_map();
287
288        assert_eq!(readback.poll_blocking(&ctx.device), StatsPoll::Failed);
289        assert_eq!(readback.values(), [1, 2, 3], "a failed map leaves the values it found");
290        assert_eq!(readback.poll(&ctx.device), StatsPoll::Failed);
291        assert!(ctx.faults.take().is_some(), "the device reports the refused map");
292    }
293
294    /// A map given up without being read must not leave the staging buffer mapped or pending. The next copy into it
295    /// would be a validation error.
296    #[test]
297    fn an_abandoned_map_leaves_the_next_readback_working() {
298        let Some(ctx) = headless_context("henad_readback_abandoned_test", wgpu::Features::empty()) else {
299            log::warn!("skipping an_abandoned_map_leaves_the_next_readback_working: no adapter");
300            return;
301        };
302        let mut readback = CounterReadback::new(&ctx.device, "henad_readback_abandoned_test", 3);
303
304        begin_read_back(&ctx, &mut readback, &[1, 2, 3]);
305        readback.abandon_map();
306        assert_eq!(
307            readback.poll(&ctx.device),
308            StatsPoll::Failed,
309            "an abandoned map polls as failed"
310        );
311
312        // This time the map completes before it is given up.
313        begin_read_back(&ctx, &mut readback, &[4, 5, 6]);
314        ctx.device
315            .poll(wgpu::PollType::wait_indefinitely())
316            .expect("the map runs");
317        readback.abandon_map();
318
319        begin_read_back(&ctx, &mut readback, &[7, 8, 9]);
320        assert_eq!(readback.poll_blocking(&ctx.device), StatsPoll::Landed);
321        assert_eq!(readback.values(), [7, 8, 9]);
322        assert!(ctx.faults.take().is_none(), "a copy went into a mapped staging buffer");
323    }
324}
henad_compute/gpu/primitives/readback.rs

henad_compute/gpu/primitives/
readback.rs
