unlit_wgpu/readback.rs
1//! Copying GPU resources back to the host.
2//!
3//! A readback is the one GPU operation that produces something on the CPU:
4//! render a frame into a texture, or fill a buffer, then copy it into a
5//! `MAP_READ` buffer and map it. The two functions here do that in one
6//! blocking call each, so a caller that wants to inspect a frame — a test, a
7//! screenshot, a debug tool — does not have to write the staging dance.
8//!
9//! Both are deliberately low-level: they take the wgpu handles and return raw
10//! bytes, with no opinion about what the bytes mean. Decoding them into an
11//! image, a vertex stream or anything else is the caller's.
12//!
13//! # Blocking
14//!
15//! Both calls submit the copy, poll the device to completion and map the
16//! staging buffer, so they block until the GPU has finished. That is what a
17//! readback is for: the bytes cannot exist before then. A frame loop should
18//! only call one where it actually needs the pixels.
19
20use ringbuf::StaticRb;
21use ringbuf::traits::{Consumer, Producer, Split};
22
23/// Copy `size` bytes of `source`, starting at `offset`, back to the host.
24///
25/// The bytes are copied into a staging buffer and mapped, so the call blocks
26/// until the copy has run. `offset` and `size` are the same pair
27/// [`wgpu::CommandEncoder::copy_buffer_to_buffer`] takes, and must satisfy the
28/// same alignment: a multiple of [`wgpu::COPY_BUFFER_ALIGNMENT`] each, and
29/// within `source`.
30///
31/// # Panics
32///
33/// If the copy is misaligned or out of range, if the device poll fails, or if
34/// the buffer cannot be mapped.
35pub fn readback_buffer(
36 device: &wgpu::Device,
37 queue: &wgpu::Queue,
38 source: &wgpu::Buffer,
39 offset: u64,
40 size: u64,
41) -> Vec<u8> {
42 let staging = device.create_buffer(&wgpu::BufferDescriptor {
43 label: Some("unlit_wgpu::readback::buffer"),
44 size,
45 usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
46 mapped_at_creation: false,
47 });
48 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
49 label: Some("unlit_wgpu::readback::buffer"),
50 });
51 encoder.copy_buffer_to_buffer(source, offset, &staging, 0, size);
52 queue.submit([encoder.finish()]);
53
54 let slice = staging.slice(..);
55 // The callback hands the result over once, so a one-slot ring is all the
56 // handoff needs; the poll below is what makes the callback run. A
57 // `StaticRb` keeps the slot inline instead of allocating it on the heap.
58 let (mut sender, mut receiver) =
59 StaticRb::<Result<(), wgpu::BufferAsyncError>, 1>::default().split();
60 slice.map_async(wgpu::MapMode::Read, move |result| {
61 let _ = sender.try_push(result);
62 });
63 device
64 .poll(wgpu::PollType::wait_indefinitely())
65 .expect("the device polls to completion during a readback");
66 receiver
67 .try_pop()
68 .expect("the map callback fires once the device has polled")
69 .expect("the staging buffer maps for reading");
70 let bytes = {
71 let view = slice.get_mapped_range().expect("the mapping is held");
72 view.to_vec()
73 };
74 staging.unmap();
75 bytes
76}
77
78/// Copy `texture` back to the host as tight rows.
79///
80/// The texture's own dimensions and format decide the copy, so the result is
81/// its whole mip level 0 with no row padding: the buffer rows wgpu requires to
82/// be aligned to [`wgpu::COPY_BYTES_PER_ROW_ALIGNMENT`] are stripped back off.
83/// A row is the texture's width rounded up to whole format blocks, so a
84/// block-compressed texture reads back block-aligned exactly as wgpu lays it
85/// out.
86///
87/// The bytes are in the texture's format — an sRGB texture reads back sRGB
88/// values, a `Rgba8Unorm` one linear ones — and the call blocks until the copy
89/// has run.
90///
91/// # Panics
92///
93/// If the texture's format has no block copy size (a depth-stencil format read
94/// without choosing an aspect, say), or if the copy cannot run.
95pub fn readback_texture(
96 device: &wgpu::Device,
97 queue: &wgpu::Queue,
98 texture: &wgpu::Texture,
99) -> Vec<u8> {
100 let format = texture.format();
101 let (block_width, block_height) = format.block_dimensions();
102 let block_bytes = format
103 .block_copy_size(None)
104 .expect("the texture's format has a block copy size to read it back in");
105 // wgpu copies whole blocks: a row is the texture's width rounded up to a
106 // block, and the copy runs for as many whole rows as the height needs.
107 let rows = texture.height().div_ceil(block_height);
108 let tight_bytes_per_row =
109 u64::from(texture.width().div_ceil(block_width)) * u64::from(block_bytes);
110 let padded_bytes_per_row =
111 tight_bytes_per_row.next_multiple_of(u64::from(wgpu::COPY_BYTES_PER_ROW_ALIGNMENT));
112 let size = padded_bytes_per_row * u64::from(rows);
113 let staging = device.create_buffer(&wgpu::BufferDescriptor {
114 label: Some("unlit_wgpu::readback::texture"),
115 size,
116 usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC,
117 mapped_at_creation: false,
118 });
119 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
120 label: Some("unlit_wgpu::readback::texture"),
121 });
122 encoder.copy_texture_to_buffer(
123 texture.as_image_copy(),
124 wgpu::TexelCopyBufferInfo {
125 buffer: &staging,
126 layout: wgpu::TexelCopyBufferLayout {
127 offset: 0,
128 bytes_per_row: Some(padded_bytes_per_row as u32),
129 rows_per_image: Some(rows),
130 },
131 },
132 // Layer 0 only: the buffer is sized for one image, so a 3D or array
133 // texture reads back its first layer.
134 wgpu::Extent3d {
135 width: texture.width(),
136 height: texture.height(),
137 depth_or_array_layers: 1,
138 },
139 );
140 queue.submit([encoder.finish()]);
141 let padded = readback_buffer(device, queue, &staging, 0, size);
142 if padded_bytes_per_row == tight_bytes_per_row {
143 return padded;
144 }
145 let mut tight = Vec::with_capacity((tight_bytes_per_row * u64::from(rows)) as usize);
146 for row in padded.chunks_exact(padded_bytes_per_row as usize) {
147 tight.extend_from_slice(&row[..tight_bytes_per_row as usize]);
148 }
149 tight
150}