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
112#[derive(Clone, Copy, Debug, PartialEq)]
115pub struct ScenePassOps {
116 pub clear: Option<[f32; 4]>,
119 pub load_depth: bool,
121 pub store_depth: bool,
125}
126
127struct BlitShared {
128 shader: wgpu::ShaderModule,
129 layout: wgpu::BindGroupLayout,
130 sampler: wgpu::Sampler,
131 pipeline_layout: wgpu::PipelineLayout,
132}
133
134struct Target {
135 key: (wgpu::TextureFormat, u32, u32, u32),
136 #[allow(dead_code)]
137 scene: wgpu::Texture,
138 scene_view: wgpu::TextureView,
139 #[allow(dead_code)]
140 depth: wgpu::Texture,
141 depth_view: wgpu::TextureView,
142 blit_bind: wgpu::BindGroup,
143 last_used_tick: u64,
144 last_used_frame: u64,
145 bytes: u64,
146}
147
148impl DepthComposite {
149 pub fn get(resources: &mut CallbackResources) -> &mut Self {
152 resources.get_or_insert_with::<Self>()
153 }
154
155 pub fn begin_frame(&mut self) {
156 self.frame_index = self.frame_index.wrapping_add(1);
157 }
158
159 pub fn end_frame(&mut self) {
160 let frame_index = self.frame_index;
161 let mut removed = 0u64;
162 self.targets.retain(|_, target| {
163 if frame_index.saturating_sub(target.last_used_frame) <= TARGET_IDLE_FRAMES {
164 true
165 } else {
166 removed = removed.saturating_add(target.bytes);
167 false
168 }
169 });
170 self.bytes_total = self.bytes_total.saturating_sub(removed);
171 }
172
173 fn shared(&mut self, device: &wgpu::Device) -> &BlitShared {
174 if self.shared.is_none() {
175 let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
176 label: Some("depth_composite_blit"),
177 source: wgpu::ShaderSource::Wgsl(BLIT_WGSL.into()),
178 });
179 let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
180 label: Some("depth_composite_blit_bgl"),
181 entries: &[
182 wgpu::BindGroupLayoutEntry {
183 binding: 0,
184 visibility: wgpu::ShaderStages::FRAGMENT,
185 ty: wgpu::BindingType::Texture {
186 sample_type: wgpu::TextureSampleType::Float { filterable: true },
187 view_dimension: wgpu::TextureViewDimension::D2,
188 multisampled: false,
189 },
190 count: None,
191 },
192 wgpu::BindGroupLayoutEntry {
193 binding: 1,
194 visibility: wgpu::ShaderStages::FRAGMENT,
195 ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
196 count: None,
197 },
198 ],
199 });
200 let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
201 label: Some("depth_composite_blit_sampler"),
202 address_mode_u: wgpu::AddressMode::ClampToEdge,
203 address_mode_v: wgpu::AddressMode::ClampToEdge,
204 address_mode_w: wgpu::AddressMode::ClampToEdge,
205 mag_filter: wgpu::FilterMode::Linear,
206 min_filter: wgpu::FilterMode::Linear,
207 mipmap_filter: wgpu::MipmapFilterMode::Nearest,
208 lod_min_clamp: 0.0,
209 lod_max_clamp: 0.0,
210 compare: None,
211 anisotropy_clamp: 1,
212 border_color: None,
213 });
214 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
215 label: Some("depth_composite_blit_pl"),
216 bind_group_layouts: &[Some(&layout)],
217 immediate_size: 0,
218 });
219 self.shared = Some(BlitShared {
220 shader,
221 layout,
222 sampler,
223 pipeline_layout,
224 });
225 }
226 self.shared.as_ref().expect("shared blit resources")
227 }
228
229 fn pipeline(
230 &mut self,
231 device: &wgpu::Device,
232 target_format: wgpu::TextureFormat,
233 sample_count: u32,
234 ) -> &wgpu::RenderPipeline {
235 let pipeline_key = (target_format, sample_count);
236 if !self.pipelines.contains_key(&pipeline_key) {
237 let shared = self.shared(device);
238 let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
239 label: Some("depth_composite_blit"),
240 layout: Some(&shared.pipeline_layout),
241 vertex: wgpu::VertexState {
242 module: &shared.shader,
243 entry_point: Some("vs_main"),
244 buffers: &[],
245 compilation_options: Default::default(),
246 },
247 fragment: Some(wgpu::FragmentState {
248 module: &shared.shader,
249 entry_point: Some("fs_main"),
250 targets: &[Some(wgpu::ColorTargetState {
251 format: target_format,
252 blend: Some(wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING),
256 write_mask: wgpu::ColorWrites::ALL,
257 })],
258 compilation_options: Default::default(),
259 }),
260 primitive: wgpu::PrimitiveState {
261 topology: wgpu::PrimitiveTopology::TriangleList,
262 ..Default::default()
263 },
264 depth_stencil: Some(wgpu::DepthStencilState {
266 format: DEPTH_FORMAT,
267 depth_write_enabled: Some(false),
268 depth_compare: Some(wgpu::CompareFunction::Always),
269 stencil: wgpu::StencilState {
270 front: wgpu::StencilFaceState {
271 compare: wgpu::CompareFunction::LessEqual,
272 ..Default::default()
273 },
274 back: wgpu::StencilFaceState {
275 compare: wgpu::CompareFunction::LessEqual,
276 ..Default::default()
277 },
278 read_mask: 0xFF,
279 write_mask: 0,
280 },
281 bias: wgpu::DepthBiasState::default(),
282 }),
283 multisample: wgpu::MultisampleState {
284 count: sample_count,
285 mask: !0,
286 alpha_to_coverage_enabled: false,
287 },
288 multiview_mask: None,
289 cache: None,
290 });
291 self.pipelines.insert(pipeline_key, pipeline);
292 }
293 &self.pipelines[&pipeline_key]
294 }
295
296 #[allow(clippy::too_many_arguments)] pub fn ensure(
302 &mut self,
303 device: &wgpu::Device,
304 screen: &ScreenDescriptor,
305 id: &str,
306 w: u32,
307 h: u32,
308 ) -> anyhow::Result<()> {
309 let w = w.max(1);
310 let h = h.max(1);
311 let key = (screen.target_format, screen.sample_count, w, h);
312 self.next_tick = self.next_tick.wrapping_add(1);
313 let tick = self.next_tick;
314 if let Some(target) = self.targets.get_mut(id)
315 && target.key == key
316 {
317 target.last_used_tick = tick;
318 target.last_used_frame = self.frame_index;
319 return Ok(());
320 }
321 let scene_bytes = screen
322 .target_format
323 .theoretical_memory_footprint(wgpu::Extent3d {
324 width: w,
325 height: h,
326 depth_or_array_layers: 1,
327 });
328 let depth_bytes = DEPTH_FORMAT.theoretical_memory_footprint(wgpu::Extent3d {
329 width: w,
330 height: h,
331 depth_or_array_layers: 1,
332 });
333 let bytes = scene_bytes.saturating_add(depth_bytes);
334 let max_dimension = device.limits().max_texture_dimension_2d;
335 if w > max_dimension || h > max_dimension {
336 anyhow::bail!(
337 "depth composite {id}: {w}x{h} exceeds the device texture limit {max_dimension}"
338 );
339 }
340 if bytes > MAX_DEPTH_RESOURCE_BYTES {
341 anyhow::bail!(
342 "depth composite {id}: {w}x{h} needs {bytes} bytes, over the {MAX_DEPTH_RESOURCE_BYTES} byte limit"
343 );
344 }
345 let replaced = self.targets.get(id).map_or(0, |old| old.bytes);
348 while self
349 .bytes_total
350 .saturating_sub(replaced)
351 .saturating_add(bytes)
352 > MAX_DEPTH_RESOURCE_BYTES
353 {
354 let Some(victim) = self
355 .targets
356 .iter()
357 .filter(|(key, _)| key.as_str() != id)
358 .min_by_key(|(_, target)| target.last_used_tick)
359 .map(|(key, _)| key.clone())
360 else {
361 anyhow::bail!(
362 "depth composite {id}: no room in the {MAX_DEPTH_RESOURCE_BYTES} byte budget"
363 );
364 };
365 if let Some(old) = self.targets.remove(&victim) {
366 self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
367 }
368 }
369 let shared = self.shared(device);
370 let scene = device.create_texture(&wgpu::TextureDescriptor {
371 label: Some("depth_composite_scene"),
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: screen.target_format,
381 usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
382 view_formats: &[],
383 });
384 let scene_view = scene.create_view(&wgpu::TextureViewDescriptor::default());
385 let depth = device.create_texture(&wgpu::TextureDescriptor {
386 label: Some("depth_composite_depth"),
387 size: wgpu::Extent3d {
388 width: w,
389 height: h,
390 depth_or_array_layers: 1,
391 },
392 mip_level_count: 1,
393 sample_count: 1,
394 dimension: wgpu::TextureDimension::D2,
395 format: DEPTH_FORMAT,
396 usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
397 view_formats: &[],
398 });
399 let depth_view = depth.create_view(&wgpu::TextureViewDescriptor::default());
400 let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
401 label: Some("depth_composite_blit_bg"),
402 layout: &shared.layout,
403 entries: &[
404 wgpu::BindGroupEntry {
405 binding: 0,
406 resource: wgpu::BindingResource::TextureView(&scene_view),
407 },
408 wgpu::BindGroupEntry {
409 binding: 1,
410 resource: wgpu::BindingResource::Sampler(&shared.sampler),
411 },
412 ],
413 });
414 self.pipeline(device, screen.target_format, screen.sample_count);
415 let target = Target {
416 key,
417 scene,
418 scene_view,
419 depth,
420 depth_view,
421 blit_bind: bind,
422 last_used_tick: tick,
423 last_used_frame: self.frame_index,
424 bytes,
425 };
426 if let Some(old) = self.targets.insert(id.to_string(), target) {
427 self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
428 }
429 self.bytes_total = self.bytes_total.saturating_add(bytes);
430 while self.targets.len() > MAX_DEPTH_TARGETS {
431 let Some(victim) = self
432 .targets
433 .iter()
434 .filter(|(key, _)| key.as_str() != id)
435 .min_by_key(|(_, target)| target.last_used_tick)
436 .map(|(key, _)| key.clone())
437 else {
438 break;
439 };
440 if let Some(old) = self.targets.remove(&victim) {
441 self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
442 }
443 }
444 Ok(())
445 }
446
447 pub fn begin_scene<'a>(
461 &'a self,
462 id: &str,
463 encoder: &'a mut wgpu::CommandEncoder,
464 clear: [f32; 4],
465 ) -> Option<wgpu::RenderPass<'a>> {
466 self.begin_scene_ops(
467 id,
468 encoder,
469 ScenePassOps {
470 clear: Some(clear),
471 load_depth: false,
472 store_depth: false,
473 },
474 )
475 }
476
477 pub fn begin_scene_ops<'a>(
482 &'a self,
483 id: &str,
484 encoder: &'a mut wgpu::CommandEncoder,
485 ops: ScenePassOps,
486 ) -> Option<wgpu::RenderPass<'a>> {
487 let t = self.targets.get(id)?;
488 let (w, h) = (t.key.2 as f32, t.key.3 as f32);
489 let color_load = match ops.clear {
490 Some(clear) => wgpu::LoadOp::Clear(wgpu::Color {
491 r: clear[0] as f64,
492 g: clear[1] as f64,
493 b: clear[2] as f64,
494 a: clear[3] as f64,
495 }),
496 None => wgpu::LoadOp::Load,
497 };
498 let depth_ops = wgpu::Operations {
499 load: if ops.load_depth {
500 wgpu::LoadOp::Load
501 } else {
502 wgpu::LoadOp::Clear(1.0)
503 },
504 store: if ops.store_depth {
507 wgpu::StoreOp::Store
508 } else {
509 wgpu::StoreOp::Discard
510 },
511 };
512 let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
513 label: Some("depth_composite_scene"),
514 color_attachments: &[Some(wgpu::RenderPassColorAttachment {
515 view: &t.scene_view,
516 depth_slice: None,
517 resolve_target: None,
518 ops: wgpu::Operations {
519 load: color_load,
520 store: wgpu::StoreOp::Store,
521 },
522 })],
523 depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
524 view: &t.depth_view,
525 depth_ops: Some(depth_ops),
526 stencil_ops: Some(wgpu::Operations {
527 load: wgpu::LoadOp::Clear(0),
528 store: wgpu::StoreOp::Discard,
529 }),
530 }),
531 timestamp_writes: None,
532 occlusion_query_set: None,
533 multiview_mask: None,
534 });
535 pass.set_viewport(0.0, 0.0, w, h, 0.0, 1.0);
536 Some(pass)
537 }
538
539 pub fn depth_format(&self, id: &str) -> Option<wgpu::TextureFormat> {
541 self.targets.get(id).map(|_| DEPTH_FORMAT)
542 }
543
544 pub fn blit<P: BlitRenderPass>(&self, id: &str, rpass: &mut P) {
550 let Some(t) = self.targets.get(id) else {
551 return;
552 };
553 let Some(pipeline) = self.pipelines.get(&(t.key.0, t.key.1)) else {
554 log::warn!("depth composite {id}: blit pipeline missing");
555 return;
556 };
557 rpass.set_blit_pipeline(pipeline);
558 rpass.set_blit_bind_group(0, &t.blit_bind, &[]);
559 rpass.draw_blit_triangle();
560 }
561
562 pub fn blit_callback(&self, id: &str, rpass: &mut CallbackRenderPass<'_, '_>) {
563 self.blit(id, rpass);
564 }
565}