Skip to main content

tessera_ui/graphics/
frame.rs

1use std::num::NonZeroU64;
2
3use rustc_hash::FxHashMap;
4
5#[derive(Debug)]
6struct Page {
7    buffer: wgpu::Buffer,
8    usage: wgpu::BufferUsages,
9    capacity: u64,
10    bytes: Vec<u8>,
11    last_used: u64,
12}
13
14#[derive(Clone, Debug, Hash, Eq, PartialEq)]
15enum Binding {
16    Buffer(wgpu::Buffer, u64, Option<NonZeroU64>),
17    Texture(wgpu::TextureView),
18    Sampler(wgpu::Sampler),
19}
20
21#[derive(Debug, Hash, Eq, PartialEq)]
22struct GroupKey {
23    layout: wgpu::BindGroupLayout,
24    entries: Vec<(u32, Binding)>,
25}
26
27#[derive(Debug, Default)]
28pub(crate) struct FrameResources {
29    frame: u64,
30    pages: Vec<Page>,
31    groups: FxHashMap<GroupKey, (wgpu::BindGroup, u64)>,
32}
33
34impl FrameResources {
35    pub(crate) fn begin_frame(&mut self) {
36        self.frame = self.frame.wrapping_add(1);
37        self.pages
38            .retain(|page| self.frame.wrapping_sub(page.last_used) <= 2);
39        for page in &mut self.pages {
40            page.bytes.clear();
41        }
42        self.groups
43            .retain(|_, (_, frame)| self.frame.wrapping_sub(*frame) <= 2);
44    }
45}
46
47/// A distinct buffer range for commands recorded in the current frame.
48/// Its contents must not be used by a later frame.
49pub struct FrameBuffer {
50    buffer: wgpu::Buffer,
51    offset: u64,
52    size: NonZeroU64,
53}
54
55impl FrameBuffer {
56    /// Binds this range without exposing other commands' data.
57    pub fn binding(&self) -> wgpu::BindingResource<'_> {
58        wgpu::BindingResource::Buffer(wgpu::BufferBinding {
59            buffer: &self.buffer,
60            offset: self.offset,
61            size: Some(self.size),
62        })
63    }
64}
65
66impl super::GraphicsDevice {
67    /// Stages uniform or storage bytes in a reusable frame buffer. Each call
68    /// owns a separate aligned range; pages are uploaded once before submission.
69    pub fn frame_buffer(&self, usage: wgpu::BufferUsages, bytes: &[u8]) -> FrameBuffer {
70        assert!(
71            usage == wgpu::BufferUsages::UNIFORM || usage == wgpu::BufferUsages::STORAGE,
72            "frame buffers must be uniform or storage buffers"
73        );
74        let limits = self.device.limits();
75        let (alignment, maximum) = if usage == wgpu::BufferUsages::UNIFORM {
76            (
77                limits.min_uniform_buffer_offset_alignment,
78                limits.max_uniform_buffer_binding_size,
79            )
80        } else {
81            (
82                limits.min_storage_buffer_offset_alignment,
83                limits.max_storage_buffer_binding_size,
84            )
85        };
86        let size = bytes.len() as u64;
87        assert!(
88            size > 0 && size <= u64::from(maximum),
89            "invalid frame buffer binding size"
90        );
91        let alignment = u64::from(alignment).max(4);
92        let mut resources = self.frame_resources.lock();
93        let index = resources.pages.iter().position(|page| {
94            page.usage == usage
95                && (page.bytes.len() as u64).next_multiple_of(alignment) + size <= page.capacity
96        });
97        let index = index.unwrap_or_else(|| {
98            let capacity = size
99                .max(64 << 10)
100                .next_power_of_two()
101                .min(limits.max_buffer_size);
102            let buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
103                label: Some("Frame Buffer Page"),
104                size: capacity,
105                usage: usage | wgpu::BufferUsages::COPY_DST,
106                mapped_at_creation: false,
107            });
108            self.work
109                .lock()
110                .groups
111                .entry("frame_buffers")
112                .or_default()
113                .buffer_allocations += 1;
114            let frame = resources.frame;
115            resources.pages.push(Page {
116                buffer,
117                usage,
118                capacity,
119                bytes: Vec::new(),
120                last_used: frame,
121            });
122            resources.pages.len() - 1
123        });
124        let frame = resources.frame;
125        let page = &mut resources.pages[index];
126        let offset = (page.bytes.len() as u64).next_multiple_of(alignment);
127        page.bytes.resize(offset as usize, 0);
128        page.bytes.extend_from_slice(bytes);
129        page.last_used = frame;
130        FrameBuffer {
131            buffer: page.buffer.clone(),
132            offset,
133            size: NonZeroU64::new(size).expect("nonempty frame buffer"),
134        }
135    }
136
137    pub(crate) fn upload_frame_buffers(&self) {
138        let mut resources = self.frame_resources.lock();
139        for page in resources
140            .pages
141            .iter_mut()
142            .filter(|page| !page.bytes.is_empty())
143        {
144            // Binding sizes need not be aligned, but queue writes do.
145            let padded = page.bytes.len().next_multiple_of(4);
146            page.bytes.resize(padded, 0);
147            self.queue.write_buffer(&page.buffer, 0, &page.bytes);
148            let mut work = self.work.lock();
149            let work = work.groups.entry("frame_buffers").or_default();
150            work.buffer_uploads += 1;
151            work.buffer_uploaded_bytes += padded as u64;
152        }
153    }
154
155    /// Reuses bindings for buffer ranges, textures and samplers across frames.
156    /// Array bindings are created directly rather than cached.
157    pub fn cached_bind_group(&self, descriptor: &wgpu::BindGroupDescriptor<'_>) -> wgpu::BindGroup {
158        let entries: Option<Vec<_>> = descriptor
159            .entries
160            .iter()
161            .map(|entry| {
162                let binding = match &entry.resource {
163                    wgpu::BindingResource::Buffer(buffer) => {
164                        Binding::Buffer(buffer.buffer.clone(), buffer.offset, buffer.size)
165                    }
166                    wgpu::BindingResource::TextureView(view) => Binding::Texture((*view).clone()),
167                    wgpu::BindingResource::Sampler(sampler) => Binding::Sampler((*sampler).clone()),
168                    _ => return None,
169                };
170                Some((entry.binding, binding))
171            })
172            .collect();
173        let Some(mut entries) = entries else {
174            return self.device.create_bind_group(descriptor);
175        };
176        entries.sort_unstable_by_key(|(binding, _)| *binding);
177        let key = GroupKey {
178            layout: descriptor.layout.clone(),
179            entries,
180        };
181        let mut resources = self.frame_resources.lock();
182        let frame = resources.frame;
183        if let Some((group, used)) = resources.groups.get_mut(&key) {
184            *used = frame;
185            return group.clone();
186        }
187        const MAX_GROUPS: usize = 512;
188        if resources.groups.len() >= MAX_GROUPS {
189            let oldest = resources
190                .groups
191                .iter()
192                .max_by_key(|(_, (_, used))| frame.wrapping_sub(*used))
193                .map(|(key, _)| GroupKey {
194                    layout: key.layout.clone(),
195                    entries: key.entries.clone(),
196                })
197                .expect("nonempty binding cache");
198            resources.groups.remove(&oldest);
199        }
200        let group = self.device.create_bind_group(descriptor);
201        resources.groups.insert(key, (group.clone(), frame));
202        group
203    }
204}
205
206#[cfg(all(test, feature = "headless", not(target_family = "wasm")))]
207mod tests {
208    #[test]
209    #[ignore = "requires a native GPU adapter"]
210    fn queued_frame_ranges_preserve_distinct_values_and_reuse_bindings_next_frame() {
211        let session = crate::RuntimeSession::new();
212        let _scope = session.enter();
213        let core = pollster::block_on(crate::renderer::core::RenderCore::new_offscreen(
214            winit::dpi::PhysicalSize::new(1, 1),
215            wgpu::TextureFormat::Rgba8Unorm,
216            1,
217            crate::renderer::DEFAULT_MAX_NON_SAMPLER_BINDINGS,
218        ));
219        let graphics = core.resources().graphics;
220        let device = graphics.device();
221        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
222            label: None,
223            source: wgpu::ShaderSource::Wgsl(
224                "\
225                @group(0) @binding(0) var<uniform> a: u32;
226                @group(0) @binding(1) var<uniform> b: u32;
227                @group(0) @binding(2) var<uniform> c: u32;
228                @group(0) @binding(3) var<storage, read_write> result: array<u32>;
229                @compute @workgroup_size(1) fn main() { result[0] = a; result[1] = b; result[2] = c; }
230            "
231                .into(),
232            ),
233        });
234        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
235            label: None,
236            layout: None,
237            module: &shader,
238            entry_point: Some("main"),
239            compilation_options: Default::default(),
240            cache: None,
241        });
242        let layout = pipeline.get_bind_group_layout(0);
243        let output = device.create_buffer(&wgpu::BufferDescriptor {
244            label: None,
245            size: 12,
246            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
247            mapped_at_creation: false,
248        });
249        let readback = device.create_buffer(&wgpu::BufferDescriptor {
250            label: None,
251            size: 12,
252            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
253            mapped_at_creation: false,
254        });
255        let mut previous = None;
256        for values in [[11_u32, 37, 91], [81, 5, 7]] {
257            graphics.begin_frame();
258            let a =
259                graphics.frame_buffer(wgpu::BufferUsages::UNIFORM, bytemuck::bytes_of(&values[0]));
260            let b =
261                graphics.frame_buffer(wgpu::BufferUsages::UNIFORM, bytemuck::bytes_of(&values[1]));
262            let mut padded = vec![0; 64 << 10];
263            padded[..4].copy_from_slice(bytemuck::bytes_of(&values[2]));
264            let c = graphics.frame_buffer(wgpu::BufferUsages::UNIFORM, &padded);
265            let group = graphics.cached_bind_group(&wgpu::BindGroupDescriptor {
266                label: None,
267                layout: &layout,
268                entries: &[
269                    wgpu::BindGroupEntry {
270                        binding: 0,
271                        resource: a.binding(),
272                    },
273                    wgpu::BindGroupEntry {
274                        binding: 1,
275                        resource: b.binding(),
276                    },
277                    wgpu::BindGroupEntry {
278                        binding: 2,
279                        resource: c.binding(),
280                    },
281                    wgpu::BindGroupEntry {
282                        binding: 3,
283                        resource: output.as_entire_binding(),
284                    },
285                ],
286            });
287            if let Some(previous) = previous {
288                assert_eq!(group, previous, "stable ranges reuse their bindings");
289            }
290            previous = Some(group.clone());
291            let mut encoder = device.create_command_encoder(&Default::default());
292            {
293                let mut pass = encoder.begin_compute_pass(&Default::default());
294                pass.set_pipeline(&pipeline);
295                pass.set_bind_group(0, &group, &[]);
296                pass.dispatch_workgroups(1, 1, 1);
297            }
298            encoder.copy_buffer_to_buffer(&output, 0, &readback, 0, 12);
299            graphics.upload_frame_buffers();
300            graphics.queue().submit([encoder.finish()]);
301            let (sender, receiver) = std::sync::mpsc::channel();
302            readback
303                .slice(..)
304                .map_async(wgpu::MapMode::Read, move |result| {
305                    sender.send(result).unwrap();
306                });
307            device.poll(wgpu::PollType::wait_indefinitely()).unwrap();
308            receiver.recv().unwrap().unwrap();
309            assert_eq!(
310                bytemuck::cast_slice::<u8, u32>(&readback.slice(..).get_mapped_range().unwrap()),
311                values
312            );
313            readback.unmap();
314        }
315        let work = graphics.diagnostics().group("frame_buffers");
316        assert_eq!(
317            work.buffer_allocations, 2,
318            "both frames reuse the pages on either side of a page boundary"
319        );
320        assert_eq!(work.buffer_uploads, 4, "each frame uploads each page once");
321    }
322}