flashraw_gpu/
policy.rs

1//! CPU-vs-GPU backend selection.
2//!
3//! The decision compares *estimated* end-to-end times built from measured
4//! throughputs: a one-time GPU micro-benchmark ([`GpuCalibration`], cached per
5//! context) and CPU throughputs supplied by the caller ([`CpuCosts`], measured by
6//! the `flashraw` facade the same way). Hard rules come first: software adapters are
7//! never auto-selected, tiny frames stay on the CPU, and outputs that need no
8//! processing never go through the GPU.
9
10use crate::context::{AdapterClass, GpuContext};
11use crate::run::{DevelopOutput, DevelopSpec, GpuFrame, GpuOp};
12use flashraw_core::{CfaPattern, FrameInfo, Orientation, PackOrder, Rect, Result};
13use std::sync::Arc;
14
15/// Measured GPU characteristics.
16#[derive(Debug, Clone, Copy, PartialEq)]
17pub struct GpuCalibration {
18    /// Host → device bandwidth, bytes per ms.
19    pub upload_bytes_per_ms: f64,
20    /// Device → host bandwidth, bytes per ms.
21    pub download_bytes_per_ms: f64,
22    /// Fused develop (packed input → RGBA8), pixels per ms.
23    pub develop_px_per_ms: f64,
24    /// Half-size (2x2 superpixel) develop, *input* samples per ms.
25    pub develop_half_px_per_ms: f64,
26    /// Unpack kernel, samples per ms.
27    pub unpack_px_per_ms: f64,
28    /// Normalise kernel, samples per ms.
29    pub normalize_px_per_ms: f64,
30    /// Fixed cost per run (submissions, synchronisation, allocations), ms.
31    pub overhead_ms: f64,
32    /// Extra cost of each additional band when a frame is tiled (submission, bind
33    /// group, halo rows, lost overlap), ms.
34    pub band_overhead_ms: f64,
35    /// True when measured; false for the static defaults.
36    pub measured: bool,
37}
38
39impl GpuCalibration {
40    /// Conservative static defaults per adapter class (used when calibration is
41    /// disabled with `FLASHRAW_GPU_CALIBRATION=static`).
42    pub fn defaults(class: AdapterClass) -> Self {
43        let (up, down, dev, unp, ovh) = match class {
44            AdapterClass::Discrete => (10e6, 8e6, 4e6, 10e6, 1.0),
45            AdapterClass::Integrated => (8e6, 8e6, 1e6, 3e6, 1.0),
46            _ => (3e6, 3e6, 0.2e6, 0.5e6, 2.0),
47        };
48        GpuCalibration {
49            upload_bytes_per_ms: up,
50            download_bytes_per_ms: down,
51            develop_px_per_ms: dev,
52            develop_half_px_per_ms: dev * 2.0,
53            unpack_px_per_ms: unp,
54            normalize_px_per_ms: unp,
55            overhead_ms: ovh,
56            band_overhead_ms: ovh * 0.25,
57            measured: false,
58        }
59    }
60}
61
62/// CPU throughputs (pixels per ms with all threads) for the same operations.
63#[derive(Debug, Clone, Copy, PartialEq)]
64pub struct CpuCosts {
65    /// Packed unpack (SIMD), samples per ms.
66    pub unpack_px_per_ms: f64,
67    /// Fused develop + RGBA8 encode (bilinear), output pixels per ms.
68    pub develop_px_per_ms: f64,
69    /// Half-size develop + RGBA8 encode, *input* samples per ms.
70    pub develop_half_px_per_ms: f64,
71    /// Normalisation, samples per ms.
72    pub normalize_px_per_ms: f64,
73}
74
75/// What the caller wants produced.
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub enum WorkOutput {
78    RawU16,
79    Normalized,
80    /// Developed RGB with this many output bytes per pixel (4 for RGBA8, 12 for f32 RGB).
81    Rgb {
82        bytes_per_px: u32,
83        half_size: bool,
84    },
85}
86
87/// Description of one decode for the cost model.
88#[derive(Debug, Clone, Copy, PartialEq)]
89pub struct Workload {
90    /// Samples in the frame.
91    pub samples: u64,
92    /// Bytes uploaded on the GPU path (packed payload or u16 samples).
93    pub upload_bytes: u64,
94    /// The payload is bit-packed (GPU can unpack it; otherwise the CPU decodes first).
95    pub packed: bool,
96    pub output: WorkOutput,
97    /// The result stays on the GPU (no download).
98    pub resident: bool,
99    /// Bands the frame is processed in on this device (1 = fits one storage binding;
100    /// see [`GpuContext::estimate_bands`]).
101    pub bands: u32,
102    /// The edge-aware demosaic instead of the bilinear one (RGB outputs).
103    pub edge_aware: bool,
104    /// The AHD demosaic instead of the bilinear one (RGB outputs).
105    pub ahd: bool,
106}
107
108impl Workload {
109    /// Bytes of the result.
110    pub fn output_bytes(&self) -> u64 {
111        match self.output {
112            WorkOutput::RawU16 => self.samples * 2,
113            WorkOutput::Normalized => self.samples * 4,
114            WorkOutput::Rgb { bytes_per_px, half_size } => {
115                (if half_size { self.samples / 4 } else { self.samples }) * bytes_per_px as u64
116            }
117        }
118    }
119}
120
121/// Outcome of [`decide`].
122#[derive(Debug, Clone, PartialEq)]
123pub struct Decision {
124    pub use_gpu: bool,
125    pub reason: String,
126    pub est_cpu_ms: f64,
127    pub est_gpu_ms: f64,
128}
129
130/// Cost of the edge-aware demosaic relative to the bilinear one, for the develop step alone: on
131/// the CPU (measured 4.4x to 9.2x: Apple M1 Ultra, X-Trans and Bayer; 3.7x to 5.5x on x86-64) and
132/// on the GPU (measured 6.2x to 6.3x on an Apple M1 Ultra, kernel time). Rounded; only the
133/// CPU/GPU ratio matters.
134pub const EDGE_CPU_FACTOR: f64 = 5.0;
135pub const EDGE_GPU_FACTOR: f64 = 6.0;
136/// The same for the AHD demosaic: on the CPU 2.3x (x86-64, 45 MP Bayer, develop to RGBA8 on one
137/// thread), on the GPU 5.2x (Apple M1 Ultra, 24 MP Bayer, kernel time 26 ms against 5 ms).
138pub const AHD_CPU_FACTOR: f64 = 2.3;
139pub const AHD_GPU_FACTOR: f64 = 5.0;
140
141/// Frames below this size stay on the CPU (fixed GPU overheads dominate).
142pub const MIN_GPU_SAMPLES: u64 = 2_000_000;
143/// The GPU must win by this factor (estimates are noisy; CPU is the safer default).
144pub const GPU_MARGIN: f64 = 0.85;
145
146/// Estimate CPU and GPU times and choose.
147pub fn decide(class: AdapterClass, gpu: &GpuCalibration, cpu: &CpuCosts, w: &Workload) -> Decision {
148    let out_px = match w.output {
149        WorkOutput::Rgb { half_size: true, .. } => w.samples / 4,
150        _ => w.samples,
151    };
152    let (cpu_ms, gpu_kernel_ms, out_bytes) = match w.output {
153        WorkOutput::RawU16 => {
154            let c = if w.packed { w.samples as f64 / cpu.unpack_px_per_ms } else { 0.0 };
155            (c, w.samples as f64 / gpu.unpack_px_per_ms, w.samples * 2)
156        }
157        WorkOutput::Normalized => (
158            w.samples as f64 / cpu.normalize_px_per_ms
159                + if w.packed { w.samples as f64 / cpu.unpack_px_per_ms } else { 0.0 },
160            w.samples as f64 / gpu.normalize_px_per_ms,
161            w.samples * 4,
162        ),
163        WorkOutput::Rgb { bytes_per_px, half_size } => {
164            let unpack = if w.packed { w.samples as f64 / cpu.unpack_px_per_ms } else { 0.0 };
165            let (c, g) = if half_size {
166                (w.samples as f64 / cpu.develop_half_px_per_ms, w.samples as f64 / gpu.develop_half_px_per_ms)
167            } else if w.edge_aware || w.ahd {
168                let (fc, fg) =
169                    if w.ahd { (AHD_CPU_FACTOR, AHD_GPU_FACTOR) } else { (EDGE_CPU_FACTOR, EDGE_GPU_FACTOR) };
170                (out_px as f64 / cpu.develop_px_per_ms * fc, w.samples as f64 / gpu.develop_px_per_ms * fg)
171            } else {
172                (out_px as f64 / cpu.develop_px_per_ms, w.samples as f64 / gpu.develop_px_per_ms)
173            };
174            (c + unpack, g, out_px * bytes_per_px as u64)
175        }
176    };
177    let download_ms = if w.resident { 0.0 } else { out_bytes as f64 / gpu.download_bytes_per_ms };
178    // Tiled frames pay a per-band cost; halo rows add ~2 rows of upload per band.
179    let extra_bands = w.bands.max(1) as f64 - 1.0;
180    let gpu_ms = gpu.overhead_ms
181        + extra_bands * gpu.band_overhead_ms
182        + w.upload_bytes as f64 / gpu.upload_bytes_per_ms
183        + gpu_kernel_ms
184        + download_ms;
185    let mk = |use_gpu: bool, reason: String| Decision { use_gpu, reason, est_cpu_ms: cpu_ms, est_gpu_ms: gpu_ms };
186    if class == AdapterClass::Software {
187        return mk(false, "software adapter".into());
188    }
189    if w.resident {
190        return mk(true, "GPU-resident output requested".into());
191    }
192    if w.output == WorkOutput::RawU16 && !w.packed {
193        return mk(false, "raw output of an entropy-coded payload needs no GPU work".into());
194    }
195    if w.samples < MIN_GPU_SAMPLES {
196        return mk(false, format!("small frame ({:.1} MP)", w.samples as f64 / 1e6));
197    }
198    if gpu_ms < cpu_ms * GPU_MARGIN {
199        mk(true, format!("estimated GPU {gpu_ms:.1} ms < CPU {cpu_ms:.1} ms"))
200    } else {
201        mk(false, format!("estimated CPU {cpu_ms:.1} ms <= GPU {gpu_ms:.1} ms"))
202    }
203}
204
205fn calibration_disabled() -> bool {
206    #[cfg(not(target_arch = "wasm32"))]
207    {
208        std::env::var("FLASHRAW_GPU_CALIBRATION").map(|v| v.eq_ignore_ascii_case("static")).unwrap_or(false)
209    }
210    #[cfg(target_arch = "wasm32")]
211    {
212        false
213    }
214}
215
216impl GpuContext {
217    /// Measured characteristics of this device (one-time micro-benchmark, ~2048x2048
218    /// synthetic frame; cached). Returns static defaults if the benchmark fails or
219    /// `FLASHRAW_GPU_CALIBRATION=static`.
220    pub async fn calibration(self: &Arc<Self>) -> GpuCalibration {
221        if let Some(c) = self.calibration.get() {
222            return *c;
223        }
224        let c = if calibration_disabled() {
225            GpuCalibration::defaults(self.info.class)
226        } else {
227            self.measure().await.unwrap_or_else(|_| GpuCalibration::defaults(self.info.class))
228        };
229        *self.calibration.get_or_init(|| c)
230    }
231
232    /// The calibration if it has already been measured (never runs the benchmark).
233    pub fn cached_calibration(&self) -> Option<GpuCalibration> {
234        self.calibration.get().copied()
235    }
236
237    /// Blocking [`GpuContext::calibration`] (native only).
238    #[cfg(not(target_arch = "wasm32"))]
239    pub fn calibration_blocking(self: &Arc<Self>) -> GpuCalibration {
240        pollster::block_on(self.calibration())
241    }
242
243    async fn measure(self: &Arc<Self>) -> Result<GpuCalibration> {
244        const W: u32 = 2048;
245        const H: u32 = 2048;
246        let bits = 12u32;
247        let stride = (W * bits / 8) as usize;
248        let data: Vec<u8> = (0..stride * H as usize).map(|i| (i.wrapping_mul(2_654_435_761) >> 13) as u8).collect();
249        let payload = flashraw_core::PackedPayload {
250            data: &data,
251            width: W,
252            height: H,
253            bits,
254            order: PackOrder::MsbBytes,
255            row_stride: stride,
256        };
257        let frame = GpuFrame::from_packed(&payload);
258        let mut info = FrameInfo::new(W, H, bits as u8);
259        info.cfa = Some(CfaPattern::RGGB);
260        let dev = GpuOp::Develop(DevelopSpec {
261            crop: Rect::full(W, H),
262            half_size: false,
263            wb: [2.0, 1.0, 1.5, 1.0],
264            cam_to_rgb: [[1.6, -0.4, -0.2], [-0.2, 1.4, -0.2], [0.0, -0.5, 1.5]],
265            curve: None,
266            orientation: Orientation::Normal,
267            output: DevelopOutput::Rgba8,
268            levels: None,
269            edge: None,
270        });
271        // Warm-up compiles the pipelines and faults in allocations.
272        let warm = self.run(&frame, &info, &dev).await?;
273        drop(warm);
274        // Transfer = host copy + everything in the compute wall that is not kernel time
275        // (device copies on discrete GPUs, synchronisation).
276        let transfer =
277            |t: &crate::run::GpuTimings| t.upload_ms + (t.compute_ms - t.kernel_ms.unwrap_or(t.compute_ms)).max(0.0);
278        let mut best = (f64::MAX, f64::MAX, f64::MAX);
279        let mut wall_1 = f64::MAX;
280        for _ in 0..2 {
281            let r = self.run(&frame, &info, &dev).await?;
282            let (_, dl) = r.image.download_async().await?;
283            let k = r.timings.kernel_ms.unwrap_or(r.timings.compute_ms);
284            best.0 = best.0.min(transfer(&r.timings));
285            best.1 = best.1.min(k);
286            best.2 = best.2.min(dl);
287            wall_1 = wall_1.min(r.timings.upload_ms + r.timings.compute_ms);
288        }
289        // Per-band cost: the same frame in 8 bands (256 developed rows each).
290        let band_budget = Some(W as u64 * 4 * 256 + 4096);
291        let _ = self.run_banded(&frame, &info, &dev, band_budget).await?;
292        let mut wall_8 = f64::MAX;
293        let mut bands = 1;
294        for _ in 0..2 {
295            let r = self.run_banded(&frame, &info, &dev, band_budget).await?;
296            wall_8 = wall_8.min(r.timings.upload_ms + r.timings.compute_ms);
297            bands = r.timings.bands;
298        }
299        let band_overhead = if bands > 1 { ((wall_8 - wall_1) / (bands - 1) as f64).max(0.0) } else { 0.0 };
300        let (frame, info) = (&frame, &info);
301        let kernel_time = |op: GpuOp| async move {
302            let _ = self.run(frame, info, &op).await?; // warm-up
303            let mut t = f64::MAX;
304            for _ in 0..2 {
305                let r = self.run(frame, info, &op).await?;
306                t = t.min(r.timings.kernel_ms.unwrap_or(r.timings.compute_ms));
307            }
308            Ok::<f64, flashraw_core::Error>(t)
309        };
310        let unpack = kernel_time(GpuOp::Unpack).await?;
311        let normalize = kernel_time(GpuOp::Normalize).await?;
312        let GpuOp::Develop(spec) = dev else { unreachable!() };
313        let half = kernel_time(GpuOp::Develop(DevelopSpec { half_size: true, ..spec })).await?;
314        // Upload bandwidth from two sizes (separates the fixed per-upload cost).
315        let big: Vec<u8> = data.iter().cycle().take(data.len() * 4).copied().collect();
316        let big_payload = flashraw_core::PackedPayload { data: &big, height: H * 4, ..payload };
317        let big_frame = GpuFrame::from_packed(&big_payload);
318        let big_info = FrameInfo::new(W, H * 4, bits as u8);
319        let mut up_big = f64::MAX;
320        for _ in 0..3 {
321            let r = self.run(&big_frame, &big_info, &GpuOp::Unpack).await?;
322            up_big = up_big.min(transfer(&r.timings));
323        }
324        let upload_bw = (big.len() - data.len()) as f64 / (up_big - best.0).max(1e-2);
325        let upload_bw = upload_bw.min(big.len() as f64 / up_big.max(1e-3) * 4.0);
326        let upload_fixed = (best.0 - data.len() as f64 / upload_bw).max(0.0);
327        let px = (W * H) as f64;
328        Ok(GpuCalibration {
329            upload_bytes_per_ms: upload_bw,
330            download_bytes_per_ms: px * 4.0 / best.2.max(1e-3),
331            develop_px_per_ms: px / best.1.max(1e-3),
332            develop_half_px_per_ms: px / half.max(1e-3),
333            unpack_px_per_ms: px / unpack.max(1e-3),
334            normalize_px_per_ms: px / normalize.max(1e-3),
335            overhead_ms: upload_fixed + 0.2,
336            band_overhead_ms: band_overhead,
337            measured: true,
338        })
339    }
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345
346    fn cpu() -> CpuCosts {
347        CpuCosts {
348            unpack_px_per_ms: 2e6,
349            develop_px_per_ms: 0.3e6,
350            develop_half_px_per_ms: 0.5e6,
351            normalize_px_per_ms: 1e6,
352        }
353    }
354
355    fn work(output: WorkOutput, packed: bool, mp: f64) -> Workload {
356        let samples = (mp * 1e6) as u64;
357        Workload {
358            samples,
359            upload_bytes: samples * 14 / 8,
360            packed,
361            output,
362            resident: false,
363            bands: 1,
364            edge_aware: false,
365            ahd: false,
366        }
367    }
368
369    #[test]
370    fn software_never_auto() {
371        let g = GpuCalibration::defaults(AdapterClass::Discrete);
372        let d = decide(
373            AdapterClass::Software,
374            &g,
375            &cpu(),
376            &work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 45.0),
377        );
378        assert!(!d.use_gpu);
379    }
380
381    #[test]
382    fn small_frames_stay_on_cpu() {
383        let g = GpuCalibration::defaults(AdapterClass::Discrete);
384        let d = decide(
385            AdapterClass::Discrete,
386            &g,
387            &cpu(),
388            &work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 1.0),
389        );
390        assert!(!d.use_gpu);
391    }
392
393    #[test]
394    fn raw_entropy_coded_stays_on_cpu() {
395        let g = GpuCalibration::defaults(AdapterClass::Discrete);
396        let d = decide(AdapterClass::Discrete, &g, &cpu(), &work(WorkOutput::RawU16, false, 45.0));
397        assert!(!d.use_gpu);
398    }
399
400    #[test]
401    fn large_rgb_on_fast_gpu() {
402        let g = GpuCalibration::defaults(AdapterClass::Discrete);
403        let d = decide(
404            AdapterClass::Discrete,
405            &g,
406            &cpu(),
407            &work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 45.0),
408        );
409        assert!(d.use_gpu, "{d:?}");
410        // A much faster CPU flips the decision.
411        let fast = CpuCosts {
412            unpack_px_per_ms: 20e6,
413            develop_px_per_ms: 10e6,
414            develop_half_px_per_ms: 20e6,
415            normalize_px_per_ms: 20e6,
416        };
417        let d = decide(
418            AdapterClass::Discrete,
419            &g,
420            &fast,
421            &work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 45.0),
422        );
423        assert!(!d.use_gpu, "{d:?}");
424    }
425
426    #[test]
427    fn edge_aware_costs_more_on_both_sides() {
428        let g = GpuCalibration::defaults(AdapterClass::Discrete);
429        let w = work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 45.0);
430        let plain = decide(AdapterClass::Discrete, &g, &cpu(), &w);
431        let edge = decide(AdapterClass::Discrete, &g, &cpu(), &Workload { edge_aware: true, ..w });
432        assert!(edge.est_cpu_ms > plain.est_cpu_ms * 1.5, "{edge:?} {plain:?}");
433        assert!(edge.est_gpu_ms > plain.est_gpu_ms, "{edge:?} {plain:?}");
434    }
435
436    #[test]
437    fn tiled_frames_cost_more_but_still_use_the_gpu() {
438        let g = GpuCalibration::defaults(AdapterClass::Discrete);
439        let w = work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 150.0);
440        let one = decide(AdapterClass::Discrete, &g, &cpu(), &w);
441        let tiled = decide(AdapterClass::Discrete, &g, &cpu(), &Workload { bands: 5, ..w });
442        assert!(tiled.use_gpu, "{tiled:?}");
443        assert!((tiled.est_gpu_ms - one.est_gpu_ms - 4.0 * g.band_overhead_ms).abs() < 1e-9);
444    }
445
446    #[test]
447    fn resident_output_forces_gpu() {
448        let g = GpuCalibration::defaults(AdapterClass::Integrated);
449        let mut w = work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 0.5);
450        w.resident = true;
451        assert!(decide(AdapterClass::Integrated, &g, &cpu(), &w).use_gpu);
452    }
453}
flashraw_gpu/policy.rs

flashraw_gpu/
policy.rs
