pfx_trace/
probe.rs

1use pfx_gpu::Gpu;
2
3use crate::gpu::{Trace, texture};
4
5pub(crate) struct Parts {
6    pub buffers: [(u32, wgpu::Buffer); 9],
7    pub lit: bool,
8}
9
10#[derive(Clone, Debug, PartialEq)]
11pub struct Probe {
12    pub width: u32,
13    pub height: u32,
14    pub depth: Vec<f32>,
15    pub distance: Vec<f32>,
16    pub id: Vec<u32>,
17    pub normal: Vec<[f32; 3]>,
18}
19
20impl Probe {
21    pub fn without(&self, ids: &[u32]) -> Vec<u32> {
22        self.id
23            .iter()
24            .map(|&id| if ids.contains(&id) { 0 } else { id })
25            .collect()
26    }
27
28    pub fn material_id(material: u32) -> u32 {
29        material + 1
30    }
31}
32
33pub fn shader(lit: bool) -> String {
34    format!(
35        "{}\n{}",
36        crate::gpu::shader(lit),
37        include_str!("probe.wgsl")
38    )
39}
40
41impl Trace {
42    pub fn probe(&self, gpu: &Gpu) -> Result<Probe, String> {
43        let module = gpu
44            .device
45            .create_shader_module(wgpu::ShaderModuleDescriptor {
46                label: Some("trace probe shader"),
47                source: wgpu::ShaderSource::Wgsl(shader(self.probe_parts.lit).into()),
48            });
49        let pipeline = gpu
50            .device
51            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
52                label: Some("trace probe pipeline"),
53                layout: None,
54                module: &module,
55                entry_point: Some("probe"),
56                compilation_options: Default::default(),
57                cache: None,
58            });
59        let depth = texture(gpu, self.width, self.height, "trace probe depth");
60        let normal = texture(gpu, self.width, self.height, "trace probe normal");
61        let mut frame = self.base;
62        frame.size[2] = 0;
63        frame.size[3] = 1;
64        let uniform = gpu.device.create_buffer(&wgpu::BufferDescriptor {
65            label: Some("trace probe frame"),
66            size: std::mem::size_of_val(&frame) as u64,
67            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
68            mapped_at_creation: false,
69        });
70        gpu.queue
71            .write_buffer(&uniform, 0, bytemuck::bytes_of(&frame));
72        let mut entries = vec![
73            wgpu::BindGroupEntry {
74                binding: 0,
75                resource: uniform.as_entire_binding(),
76            },
77            wgpu::BindGroupEntry {
78                binding: 17,
79                resource: wgpu::BindingResource::TextureView(&depth.view),
80            },
81            wgpu::BindGroupEntry {
82                binding: 18,
83                resource: wgpu::BindingResource::TextureView(&normal.view),
84            },
85        ];
86        entries.extend(self.probe_parts.buffers.iter().map(|(binding, buffer)| {
87            wgpu::BindGroupEntry {
88                binding: *binding,
89                resource: buffer.as_entire_binding(),
90            }
91        }));
92        let binding = gpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
93            label: Some("trace probe bindings"),
94            layout: &pipeline.get_bind_group_layout(0),
95            entries: &entries,
96        });
97        let mut encoder = gpu
98            .device
99            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
100                label: Some("trace probe"),
101            });
102        {
103            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
104                label: Some("trace probe"),
105                timestamp_writes: None,
106            });
107            pass.set_pipeline(&pipeline);
108            pass.set_bind_group(0, &binding, &[]);
109            pass.dispatch_workgroups(self.width.div_ceil(8), self.height.div_ceil(8), 1);
110        }
111        gpu.queue.submit(Some(encoder.finish()));
112        let depth = texels(&gpu.readback_bytes(&depth)?);
113        let normal = texels(&gpu.readback_bytes(&normal)?);
114        Ok(Probe {
115            width: self.width,
116            height: self.height,
117            depth: depth.iter().map(|texel| texel[0]).collect(),
118            distance: depth.iter().map(|texel| texel[2]).collect(),
119            id: depth.iter().map(|texel| texel[1].round() as u32).collect(),
120            normal: normal
121                .iter()
122                .map(|texel| [texel[0], texel[1], texel[2]])
123                .collect(),
124        })
125    }
126}
127
128pub fn texels(bytes: &[u8]) -> Vec<[f32; 4]> {
129    bytes
130        .chunks_exact(16)
131        .map(|chunk| {
132            std::array::from_fn(|k| f32::from_le_bytes(chunk[k * 4..k * 4 + 4].try_into().unwrap()))
133        })
134        .collect()
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    #[test]
142    fn probe_shader_validates() {
143        for lit in [false, true] {
144            let source = shader(lit);
145            let module = naga::front::wgsl::parse_str(&source).unwrap();
146            naga::valid::Validator::new(
147                naga::valid::ValidationFlags::all(),
148                naga::valid::Capabilities::all(),
149            )
150            .validate(&module)
151            .unwrap();
152        }
153    }
154
155    #[test]
156    fn ids_drop_what_is_not_outlined() {
157        let probe = Probe {
158            width: 3,
159            height: 1,
160            depth: vec![1.0; 3],
161            distance: vec![1.0; 3],
162            id: vec![0, 1, 3],
163            normal: vec![[0.0, 0.0, 1.0]; 3],
164        };
165        assert_eq!(probe.without(&[Probe::material_id(2)]), vec![0, 1, 0]);
166        assert_eq!(texels(&[0u8; 32]).len(), 2);
167    }
168}
pfx_trace/probe.rs

pfx_trace/
probe.rs
