1use bytemuck::{Pod, Zeroable};
2use pfx_gpu::Gpu;
3use pfx_gpu::pace::{Pacer, Stats, Turned};
4use wgpu::util::DeviceExt;
5
6use crate::gpu::Trace;
7
8#[derive(Clone, Copy, Debug, PartialEq)]
9pub struct Adaptive {
10 pub threshold: f32,
11 pub min_samples: u32,
12 pub max_samples: u32,
13 pub growth: f32,
14}
15
16impl Default for Adaptive {
17 fn default() -> Self {
18 Self {
19 threshold: 0.005,
20 min_samples: 32,
21 max_samples: 1024,
22 growth: 1.5,
23 }
24 }
25}
26
27impl Adaptive {
28 pub fn valid(&self) -> bool {
29 self.threshold.is_finite()
30 && self.threshold > 0.0
31 && self.min_samples >= 2
32 && self.max_samples >= self.min_samples
33 && self.growth.is_finite()
34 && self.growth > 1.0
35 }
36
37 pub fn margin(&self) -> u32 {
38 self.rounds().len() as u32
39 }
40
41 pub fn rounds(&self) -> Vec<u32> {
42 let mut rounds = vec![(self.min_samples / 2).max(1)];
43 loop {
44 let last = *rounds.last().unwrap();
45 if last >= self.max_samples {
46 return rounds;
47 }
48 let next = ((last as f32 * self.growth).ceil() as u32)
49 .max(last + 1)
50 .max(if rounds.len() == 1 {
51 self.min_samples
52 } else {
53 0
54 })
55 .min(self.max_samples);
56 rounds.push(next);
57 }
58 }
59}
60
61#[derive(Clone, Debug, Default)]
62pub struct Converged {
63 pub rounds: Vec<Round>,
64 pub pixels: u64,
65 pub stats: Stats,
66}
67
68#[derive(Clone, Copy, Debug, PartialEq)]
69pub struct Round {
70 pub samples: u32,
71 pub active: u64,
72}
73
74impl Converged {
75 pub fn max_samples(&self) -> u32 {
76 self.rounds.last().map_or(0, |round| round.samples)
77 }
78
79 pub fn mean_samples(&self) -> f64 {
80 if self.pixels == 0 {
81 return 0.0;
82 }
83 let mut total = 0.0;
84 let mut before = 0;
85 let mut running = self.pixels;
86 for round in &self.rounds {
87 total += f64::from(round.samples - before) * running as f64;
88 before = round.samples;
89 running = round.active;
90 }
91 total / self.pixels as f64
92 }
93
94 pub fn converged(&self) -> bool {
95 self.rounds.last().is_some_and(|round| round.active == 0)
96 }
97}
98
99#[repr(C)]
100#[derive(Clone, Copy, Pod, Zeroable)]
101struct Check {
102 size: [u32; 4],
103 limits: [f32; 4],
104}
105
106impl Trace {
107 pub fn sample_adaptive<T: Turned>(
108 &mut self,
109 gpu: &Gpu,
110 adaptive: Adaptive,
111 seed: u32,
112 pacer: &mut Pacer,
113 between: impl FnMut(f64) -> T,
114 ) -> Result<Converged, String> {
115 self.sample_adaptive_with(gpu, adaptive, seed, pacer, between, |_| {})
116 }
117
118 pub fn sample_adaptive_with<T: Turned>(
119 &mut self,
120 gpu: &Gpu,
121 adaptive: Adaptive,
122 seed: u32,
123 pacer: &mut Pacer,
124 mut between: impl FnMut(f64) -> T,
125 mut reached: impl FnMut(Round),
126 ) -> Result<Converged, String> {
127 if !adaptive.valid() {
128 return Err("adaptive sampling needs a positive threshold, a growth over 1 and at least 2 samples".into());
129 }
130 self.samples = 0;
131 let pixels = u64::from(self.width) * u64::from(self.height);
132 let check = Check {
133 size: [self.width, self.height, adaptive.min_samples, 0],
134 limits: [adaptive.threshold, 0.0, 0.0, 0.0],
135 };
136 let module = gpu
137 .device
138 .create_shader_module(wgpu::ShaderModuleDescriptor {
139 label: Some("trace convergence"),
140 source: wgpu::ShaderSource::Wgsl(include_str!("converge.wgsl").into()),
141 });
142 let pipeline = |entry: &'static str| {
143 gpu.device
144 .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
145 label: Some(entry),
146 layout: None,
147 module: &module,
148 entry_point: Some(entry),
149 compilation_options: Default::default(),
150 cache: None,
151 })
152 };
153 let measure = pipeline("measure");
154 let mark = pipeline("mark");
155 let uniform = gpu
156 .device
157 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
158 label: Some("trace convergence check"),
159 contents: bytemuck::bytes_of(&check),
160 usage: wgpu::BufferUsages::UNIFORM,
161 });
162 let snaps = gpu.device.create_buffer(&wgpu::BufferDescriptor {
163 label: Some("trace convergence snapshots"),
164 size: pixels * 32,
165 usage: wgpu::BufferUsages::STORAGE,
166 mapped_at_creation: false,
167 });
168 let active = gpu.device.create_buffer(&wgpu::BufferDescriptor {
169 label: Some("trace convergence active"),
170 size: 16,
171 usage: wgpu::BufferUsages::STORAGE
172 | wgpu::BufferUsages::COPY_SRC
173 | wgpu::BufferUsages::COPY_DST,
174 mapped_at_creation: false,
175 });
176 let read = gpu.device.create_buffer(&wgpu::BufferDescriptor {
177 label: Some("trace convergence readback"),
178 size: 16,
179 usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
180 mapped_at_creation: false,
181 });
182 let bind = |pipeline: &wgpu::ComputePipeline, with_active: bool| {
183 let mut entries = vec![
184 wgpu::BindGroupEntry {
185 binding: 0,
186 resource: uniform.as_entire_binding(),
187 },
188 wgpu::BindGroupEntry {
189 binding: 1,
190 resource: self.accum.as_entire_binding(),
191 },
192 wgpu::BindGroupEntry {
193 binding: 2,
194 resource: snaps.as_entire_binding(),
195 },
196 ];
197 if with_active {
198 entries.push(wgpu::BindGroupEntry {
199 binding: 3,
200 resource: active.as_entire_binding(),
201 });
202 }
203 gpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
204 label: Some("trace convergence bindings"),
205 layout: &pipeline.get_bind_group_layout(0),
206 entries: &entries,
207 })
208 };
209 let measure_binding = bind(&measure, false);
210 let mark_binding = bind(&mark, true);
211 let groups = (self.width.div_ceil(8), self.height.div_ceil(8));
212 let mut converged = Converged {
213 pixels,
214 ..Converged::default()
215 };
216 for (step, target) in adaptive.rounds().into_iter().enumerate() {
217 let count = target - self.samples;
218 let stats = self.sample_paced(gpu, count, seed, pacer, &mut between)?;
219 converged.stats.milliseconds.extend(stats.milliseconds);
220 converged.stats.timings.extend(stats.timings);
221 converged.stats.wall_ms += stats.wall_ms;
222 let mut encoder = gpu
223 .device
224 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
225 label: Some("trace convergence"),
226 });
227 encoder.clear_buffer(&active, 0, None);
228 {
229 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
230 label: Some("trace convergence measure"),
231 timestamp_writes: None,
232 });
233 pass.set_pipeline(&measure);
234 pass.set_bind_group(0, &measure_binding, &[]);
235 pass.dispatch_workgroups(groups.0, groups.1, 1);
236 }
237 if step > 0 {
238 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
239 label: Some("trace convergence mark"),
240 timestamp_writes: None,
241 });
242 pass.set_pipeline(&mark);
243 pass.set_bind_group(0, &mark_binding, &[]);
244 pass.dispatch_workgroups(groups.0, groups.1, 1);
245 }
246 encoder.copy_buffer_to_buffer(&active, 0, &read, 0, 16);
247 gpu.queue.submit(Some(encoder.finish()));
248 let running = if step > 0 {
249 u64::from(read_count(gpu, &read)?)
250 } else {
251 pixels
252 };
253 let round = Round {
254 samples: target,
255 active: running,
256 };
257 converged.rounds.push(round);
258 reached(round);
259 if running == 0 {
260 break;
261 }
262 }
263 Ok(converged)
264 }
265}
266
267impl Trace {
268 pub fn sample_counts(&self, gpu: &Gpu) -> Result<Vec<u32>, String> {
269 let size = self.accum.size();
270 let read = gpu.device.create_buffer(&wgpu::BufferDescriptor {
271 label: Some("trace sample counts"),
272 size,
273 usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
274 mapped_at_creation: false,
275 });
276 let mut encoder = gpu
277 .device
278 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
279 label: Some("trace sample counts"),
280 });
281 encoder.copy_buffer_to_buffer(&self.accum, 0, &read, 0, size);
282 gpu.queue.submit(Some(encoder.finish()));
283 let bytes = mapped(gpu, &read)?;
284 Ok(bytes
285 .chunks_exact(48)
286 .map(|accum| f32::from_le_bytes(accum[12..16].try_into().unwrap()).abs() as u32)
287 .collect())
288 }
289}
290
291fn mapped(gpu: &Gpu, buffer: &wgpu::Buffer) -> Result<Vec<u8>, String> {
292 let (send, receive) = std::sync::mpsc::channel();
293 let slice = buffer.slice(..);
294 slice.map_async(wgpu::MapMode::Read, move |result| {
295 let _ = send.send(result);
296 });
297 gpu.device
298 .poll(wgpu::PollType::wait_indefinitely())
299 .map_err(|error| format!("{error:?}"))?;
300 receive
301 .recv()
302 .map_err(|error| error.to_string())?
303 .map_err(|error| error.to_string())?;
304 let bytes = slice.get_mapped_range().expect("a mapped range").to_vec();
305 buffer.unmap();
306 Ok(bytes)
307}
308
309fn read_count(gpu: &Gpu, buffer: &wgpu::Buffer) -> Result<u32, String> {
310 Ok(u32::from_le_bytes(
311 mapped(gpu, buffer)?[..4].try_into().unwrap(),
312 ))
313}
314
315pub fn error(prior: [f32; 4], prior_count: f32, now: [f32; 4], count: f32) -> f32 {
316 if prior_count <= 0.0 || count <= prior_count {
317 return 1e9;
318 }
319 let scale = (1.0 / count).sqrt() / (1.0 / prior_count + 1.0 / (count - prior_count)).sqrt();
320 let gap: [f32; 4] = std::array::from_fn(|k| {
321 let rest = (now[k] * count - prior[k] * prior_count) / (count - prior_count);
322 (prior[k] - rest).abs() * scale
323 });
324 (gap[0] + gap[1] + gap[2]) / (1e-4 + (now[0] + now[1] + now[2]).max(0.0).sqrt()) + gap[3]
325}
326
327#[cfg(test)]
328mod tests {
329 use super::*;
330
331 #[test]
332 fn rounds_grow_to_the_cap() {
333 let adaptive = Adaptive {
334 threshold: 0.01,
335 min_samples: 16,
336 max_samples: 100,
337 growth: 1.5,
338 };
339 assert_eq!(adaptive.rounds(), vec![8, 16, 24, 36, 54, 81, 100]);
340 assert!(adaptive.valid());
341 assert!(
342 !Adaptive {
343 growth: 1.0,
344 ..adaptive
345 }
346 .valid()
347 );
348 assert!(
349 !Adaptive {
350 max_samples: 8,
351 ..adaptive
352 }
353 .valid()
354 );
355 }
356
357 #[test]
358 fn mean_samples_count_the_pixels_still_running() {
359 let converged = Converged {
360 rounds: vec![
361 Round {
362 samples: 8,
363 active: 100,
364 },
365 Round {
366 samples: 16,
367 active: 40,
368 },
369 Round {
370 samples: 32,
371 active: 0,
372 },
373 ],
374 pixels: 100,
375 stats: Stats::default(),
376 };
377 assert!((converged.mean_samples() - (16.0 + 0.4 * 16.0)).abs() < 1e-9);
378 assert!(converged.converged());
379 assert_eq!(converged.max_samples(), 32);
380 }
381
382 #[test]
383 fn the_error_is_cycles_half_buffer_at_a_half_split() {
384 let even = [0.5, 0.4, 0.3, 1.0];
385 let odd = [0.52, 0.38, 0.33, 1.0];
386 let all: [f32; 4] = std::array::from_fn(|k| (even[k] + odd[k]) * 0.5);
387 let got = error(even, 64.0, all, 128.0);
388 let difference = (0..3).map(|k| (even[k] - odd[k]).abs()).sum::<f32>() * 0.5;
389 let want = difference / (1e-4 + (all[0] + all[1] + all[2]).sqrt());
390 assert!((got - want).abs() < 1e-6, "{got} against {want}");
391 assert_eq!(error(even, 0.0, all, 8.0), 1e9);
392 assert_eq!(error([0.2; 4], 10.0, [0.2; 4], 30.0), 0.0);
393 }
394
395 #[test]
396 fn convergence_shader_validates() {
397 let module = naga::front::wgsl::parse_str(include_str!("converge.wgsl")).unwrap();
398 naga::valid::Validator::new(
399 naga::valid::ValidationFlags::all(),
400 naga::valid::Capabilities::all(),
401 )
402 .validate(&module)
403 .unwrap();
404 }
405}