1use std::collections::HashMap;
19
20use super::{CallbackRenderPass, CallbackResources, ScreenDescriptor};
21
22const MAX_DEPTH_RESOURCE_BYTES: u64 = 256 * 1024 * 1024;
23const MAX_DEPTH_TARGETS: usize = 64;
24const TARGET_IDLE_FRAMES: u64 = 3;
29const DEPTH_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Depth24PlusStencil8;
30
31pub trait BlitRenderPass {
32 fn set_blit_pipeline(&mut self, pipeline: &wgpu::RenderPipeline);
33 fn set_blit_bind_group(&mut self, index: u32, bind_group: &wgpu::BindGroup, offsets: &[u32]);
34 fn draw_blit_triangle(&mut self);
35}
36
37impl BlitRenderPass for wgpu::RenderPass<'_> {
38 fn set_blit_pipeline(&mut self, pipeline: &wgpu::RenderPipeline) {
39 self.set_pipeline(pipeline);
40 }
41
42 fn set_blit_bind_group(&mut self, index: u32, bind_group: &wgpu::BindGroup, offsets: &[u32]) {
43 self.set_bind_group(index, bind_group, offsets);
44 }
45
46 fn draw_blit_triangle(&mut self) {
47 self.draw(0..3, 0..1);
48 }
49}
50
51impl BlitRenderPass for CallbackRenderPass<'_, '_> {
52 fn set_blit_pipeline(&mut self, pipeline: &wgpu::RenderPipeline) {
53 self.set_pipeline(pipeline);
54 }
55
56 fn set_blit_bind_group(&mut self, index: u32, bind_group: &wgpu::BindGroup, offsets: &[u32]) {
57 self.set_bind_group(index, bind_group, offsets);
58 }
59
60 fn draw_blit_triangle(&mut self) {
61 self.draw(0..3, 0..1);
62 }
63}
64
65const BLIT_WGSL: &str = r#"
72@group(0) @binding(0) var scene_tex: texture_2d<f32>;
73@group(0) @binding(1) var scene_smp: sampler;
74struct VsOut {
75 @builtin(position) pos: vec4<f32>,
76 @location(0) uv: vec2<f32>,
77};
78@vertex
79fn vs_main(@builtin(vertex_index) i: u32) -> VsOut {
80 let x = f32(i / 2u) * 4.0 - 1.0;
81 let y = f32(i % 2u) * 4.0 - 1.0;
82 var out: VsOut;
83 out.pos = vec4<f32>(x, y, 0.0, 1.0);
84 out.uv = vec2<f32>((x + 1.0) * 0.5, (1.0 - y) * 0.5);
85 return out;
86}
87@fragment
88fn fs_main(in: VsOut) -> @location(0) vec4<f32> {
89 let c = textureSample(scene_tex, scene_smp, in.uv);
90 return vec4<f32>(c.rgb * c.a, c.a);
91}
92"#;
93
94#[derive(Default)]
101pub struct DepthComposite {
102 targets: HashMap<String, Target>,
103 shared: Option<BlitShared>,
104 pipelines: HashMap<(wgpu::TextureFormat, u32), wgpu::RenderPipeline>,
107 next_tick: u64,
108 frame_index: u64,
109 bytes_total: u64,
110}
111
112struct BlitShared {
113 shader: wgpu::ShaderModule,
114 layout: wgpu::BindGroupLayout,
115 sampler: wgpu::Sampler,
116 pipeline_layout: wgpu::PipelineLayout,
117}
118
119struct Target {
120 key: (wgpu::TextureFormat, u32, u32, u32),
121 #[allow(dead_code)]
122 scene: wgpu::Texture,
123 scene_view: wgpu::TextureView,
124 #[allow(dead_code)]
125 depth: wgpu::Texture,
126 depth_view: wgpu::TextureView,
127 blit_bind: wgpu::BindGroup,
128 last_used_tick: u64,
129 last_used_frame: u64,
130 bytes: u64,
131}
132
133impl DepthComposite {
134 pub fn get(resources: &mut CallbackResources) -> &mut Self {
137 resources.get_or_insert_with::<Self>()
138 }
139
140 pub fn begin_frame(&mut self) {
141 self.frame_index = self.frame_index.wrapping_add(1);
142 }
143
144 pub fn end_frame(&mut self) {
145 let frame_index = self.frame_index;
146 let mut removed = 0u64;
147 self.targets.retain(|_, target| {
148 if frame_index.saturating_sub(target.last_used_frame) <= TARGET_IDLE_FRAMES {
149 true
150 } else {
151 removed = removed.saturating_add(target.bytes);
152 false
153 }
154 });
155 self.bytes_total = self.bytes_total.saturating_sub(removed);
156 }
157
158 fn shared(&mut self, device: &wgpu::Device) -> &BlitShared {
159 if self.shared.is_none() {
160 let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
161 label: Some("depth_composite_blit"),
162 source: wgpu::ShaderSource::Wgsl(BLIT_WGSL.into()),
163 });
164 let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
165 label: Some("depth_composite_blit_bgl"),
166 entries: &[
167 wgpu::BindGroupLayoutEntry {
168 binding: 0,
169 visibility: wgpu::ShaderStages::FRAGMENT,
170 ty: wgpu::BindingType::Texture {
171 sample_type: wgpu::TextureSampleType::Float { filterable: true },
172 view_dimension: wgpu::TextureViewDimension::D2,
173 multisampled: false,
174 },
175 count: None,
176 },
177 wgpu::BindGroupLayoutEntry {
178 binding: 1,
179 visibility: wgpu::ShaderStages::FRAGMENT,
180 ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
181 count: None,
182 },
183 ],
184 });
185 let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
186 label: Some("depth_composite_blit_sampler"),
187 address_mode_u: wgpu::AddressMode::ClampToEdge,
188 address_mode_v: wgpu::AddressMode::ClampToEdge,
189 address_mode_w: wgpu::AddressMode::ClampToEdge,
190 mag_filter: wgpu::FilterMode::Linear,
191 min_filter: wgpu::FilterMode::Linear,
192 mipmap_filter: wgpu::MipmapFilterMode::Nearest,
193 lod_min_clamp: 0.0,
194 lod_max_clamp: 0.0,
195 compare: None,
196 anisotropy_clamp: 1,
197 border_color: None,
198 });
199 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
200 label: Some("depth_composite_blit_pl"),
201 bind_group_layouts: &[Some(&layout)],
202 immediate_size: 0,
203 });
204 self.shared = Some(BlitShared {
205 shader,
206 layout,
207 sampler,
208 pipeline_layout,
209 });
210 }
211 self.shared.as_ref().expect("shared blit resources")
212 }
213
214 fn pipeline(
215 &mut self,
216 device: &wgpu::Device,
217 target_format: wgpu::TextureFormat,
218 sample_count: u32,
219 ) -> &wgpu::RenderPipeline {
220 let pipeline_key = (target_format, sample_count);
221 if !self.pipelines.contains_key(&pipeline_key) {
222 let shared = self.shared(device);
223 let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
224 label: Some("depth_composite_blit"),
225 layout: Some(&shared.pipeline_layout),
226 vertex: wgpu::VertexState {
227 module: &shared.shader,
228 entry_point: Some("vs_main"),
229 buffers: &[],
230 compilation_options: Default::default(),
231 },
232 fragment: Some(wgpu::FragmentState {
233 module: &shared.shader,
234 entry_point: Some("fs_main"),
235 targets: &[Some(wgpu::ColorTargetState {
236 format: target_format,
237 blend: Some(wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING),
241 write_mask: wgpu::ColorWrites::ALL,
242 })],
243 compilation_options: Default::default(),
244 }),
245 primitive: wgpu::PrimitiveState {
246 topology: wgpu::PrimitiveTopology::TriangleList,
247 ..Default::default()
248 },
249 depth_stencil: Some(wgpu::DepthStencilState {
251 format: DEPTH_FORMAT,
252 depth_write_enabled: Some(false),
253 depth_compare: Some(wgpu::CompareFunction::Always),
254 stencil: wgpu::StencilState {
255 front: wgpu::StencilFaceState {
256 compare: wgpu::CompareFunction::LessEqual,
257 ..Default::default()
258 },
259 back: wgpu::StencilFaceState {
260 compare: wgpu::CompareFunction::LessEqual,
261 ..Default::default()
262 },
263 read_mask: 0xFF,
264 write_mask: 0,
265 },
266 bias: wgpu::DepthBiasState::default(),
267 }),
268 multisample: wgpu::MultisampleState {
269 count: sample_count,
270 mask: !0,
271 alpha_to_coverage_enabled: false,
272 },
273 multiview_mask: None,
274 cache: None,
275 });
276 self.pipelines.insert(pipeline_key, pipeline);
277 }
278 &self.pipelines[&pipeline_key]
279 }
280
281 #[allow(clippy::too_many_arguments)] pub fn ensure(
287 &mut self,
288 device: &wgpu::Device,
289 screen: &ScreenDescriptor,
290 id: &str,
291 w: u32,
292 h: u32,
293 ) -> anyhow::Result<()> {
294 let w = w.max(1);
295 let h = h.max(1);
296 let key = (screen.target_format, screen.sample_count, w, h);
297 self.next_tick = self.next_tick.wrapping_add(1);
298 let tick = self.next_tick;
299 if let Some(target) = self.targets.get_mut(id)
300 && target.key == key
301 {
302 target.last_used_tick = tick;
303 target.last_used_frame = self.frame_index;
304 return Ok(());
305 }
306 let scene_bytes = screen
307 .target_format
308 .theoretical_memory_footprint(wgpu::Extent3d {
309 width: w,
310 height: h,
311 depth_or_array_layers: 1,
312 });
313 let depth_bytes = DEPTH_FORMAT.theoretical_memory_footprint(wgpu::Extent3d {
314 width: w,
315 height: h,
316 depth_or_array_layers: 1,
317 });
318 let bytes = scene_bytes.saturating_add(depth_bytes);
319 let max_dimension = device.limits().max_texture_dimension_2d;
320 if w > max_dimension || h > max_dimension {
321 anyhow::bail!(
322 "depth composite {id}: {w}x{h} exceeds the device texture limit {max_dimension}"
323 );
324 }
325 if bytes > MAX_DEPTH_RESOURCE_BYTES {
326 anyhow::bail!(
327 "depth composite {id}: {w}x{h} needs {bytes} bytes, over the {MAX_DEPTH_RESOURCE_BYTES} byte limit"
328 );
329 }
330 let replaced = self.targets.get(id).map_or(0, |old| old.bytes);
333 while self
334 .bytes_total
335 .saturating_sub(replaced)
336 .saturating_add(bytes)
337 > MAX_DEPTH_RESOURCE_BYTES
338 {
339 let Some(victim) = self
340 .targets
341 .iter()
342 .filter(|(key, _)| key.as_str() != id)
343 .min_by_key(|(_, target)| target.last_used_tick)
344 .map(|(key, _)| key.clone())
345 else {
346 anyhow::bail!(
347 "depth composite {id}: no room in the {MAX_DEPTH_RESOURCE_BYTES} byte budget"
348 );
349 };
350 if let Some(old) = self.targets.remove(&victim) {
351 self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
352 }
353 }
354 let shared = self.shared(device);
355 let scene = device.create_texture(&wgpu::TextureDescriptor {
356 label: Some("depth_composite_scene"),
357 size: wgpu::Extent3d {
358 width: w,
359 height: h,
360 depth_or_array_layers: 1,
361 },
362 mip_level_count: 1,
363 sample_count: 1,
364 dimension: wgpu::TextureDimension::D2,
365 format: screen.target_format,
366 usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
367 view_formats: &[],
368 });
369 let scene_view = scene.create_view(&wgpu::TextureViewDescriptor::default());
370 let depth = device.create_texture(&wgpu::TextureDescriptor {
371 label: Some("depth_composite_depth"),
372 size: wgpu::Extent3d {
373 width: w,
374 height: h,
375 depth_or_array_layers: 1,
376 },
377 mip_level_count: 1,
378 sample_count: 1,
379 dimension: wgpu::TextureDimension::D2,
380 format: DEPTH_FORMAT,
381 usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
382 view_formats: &[],
383 });
384 let depth_view = depth.create_view(&wgpu::TextureViewDescriptor::default());
385 let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
386 label: Some("depth_composite_blit_bg"),
387 layout: &shared.layout,
388 entries: &[
389 wgpu::BindGroupEntry {
390 binding: 0,
391 resource: wgpu::BindingResource::TextureView(&scene_view),
392 },
393 wgpu::BindGroupEntry {
394 binding: 1,
395 resource: wgpu::BindingResource::Sampler(&shared.sampler),
396 },
397 ],
398 });
399 self.pipeline(device, screen.target_format, screen.sample_count);
400 let target = Target {
401 key,
402 scene,
403 scene_view,
404 depth,
405 depth_view,
406 blit_bind: bind,
407 last_used_tick: tick,
408 last_used_frame: self.frame_index,
409 bytes,
410 };
411 if let Some(old) = self.targets.insert(id.to_string(), target) {
412 self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
413 }
414 self.bytes_total = self.bytes_total.saturating_add(bytes);
415 while self.targets.len() > MAX_DEPTH_TARGETS {
416 let Some(victim) = self
417 .targets
418 .iter()
419 .filter(|(key, _)| key.as_str() != id)
420 .min_by_key(|(_, target)| target.last_used_tick)
421 .map(|(key, _)| key.clone())
422 else {
423 break;
424 };
425 if let Some(old) = self.targets.remove(&victim) {
426 self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
427 }
428 }
429 Ok(())
430 }
431
432 pub fn begin_scene<'a>(
445 &'a self,
446 id: &str,
447 encoder: &'a mut wgpu::CommandEncoder,
448 clear: [f32; 4],
449 ) -> Option<wgpu::RenderPass<'a>> {
450 let t = self.targets.get(id)?;
451 let (w, h) = (t.key.2 as f32, t.key.3 as f32);
452 let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
453 label: Some("depth_composite_scene"),
454 color_attachments: &[Some(wgpu::RenderPassColorAttachment {
455 view: &t.scene_view,
456 depth_slice: None,
457 resolve_target: None,
458 ops: wgpu::Operations {
459 load: wgpu::LoadOp::Clear(wgpu::Color {
460 r: clear[0] as f64,
461 g: clear[1] as f64,
462 b: clear[2] as f64,
463 a: clear[3] as f64,
464 }),
465 store: wgpu::StoreOp::Store,
466 },
467 })],
468 depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
469 view: &t.depth_view,
470 depth_ops: Some(wgpu::Operations {
474 load: wgpu::LoadOp::Clear(1.0),
475 store: wgpu::StoreOp::Discard,
476 }),
477 stencil_ops: Some(wgpu::Operations {
478 load: wgpu::LoadOp::Clear(0),
479 store: wgpu::StoreOp::Discard,
480 }),
481 }),
482 timestamp_writes: None,
483 occlusion_query_set: None,
484 multiview_mask: None,
485 });
486 pass.set_viewport(0.0, 0.0, w, h, 0.0, 1.0);
487 Some(pass)
488 }
489
490 pub fn depth_format(&self, id: &str) -> Option<wgpu::TextureFormat> {
492 self.targets.get(id).map(|_| DEPTH_FORMAT)
493 }
494
495 pub fn blit<P: BlitRenderPass>(&self, id: &str, rpass: &mut P) {
501 let Some(t) = self.targets.get(id) else {
502 return;
503 };
504 let Some(pipeline) = self.pipelines.get(&(t.key.0, t.key.1)) else {
505 log::warn!("depth composite {id}: blit pipeline missing");
506 return;
507 };
508 rpass.set_blit_pipeline(pipeline);
509 rpass.set_blit_bind_group(0, &t.blit_bind, &[]);
510 rpass.draw_blit_triangle();
511 }
512
513 pub fn blit_callback(&self, id: &str, rpass: &mut CallbackRenderPass<'_, '_>) {
514 self.blit(id, rpass);
515 }
516}