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}
103
104impl Workload {
105    /// Bytes of the result.
106    pub fn output_bytes(&self) -> u64 {
107        match self.output {
108            WorkOutput::RawU16 => self.samples * 2,
109            WorkOutput::Normalized => self.samples * 4,
110            WorkOutput::Rgb { bytes_per_px, half_size } => {
111                (if half_size { self.samples / 4 } else { self.samples }) * bytes_per_px as u64
112            }
113        }
114    }
115}
116
117/// Outcome of [`decide`].
118#[derive(Debug, Clone, PartialEq)]
119pub struct Decision {
120    pub use_gpu: bool,
121    pub reason: String,
122    pub est_cpu_ms: f64,
123    pub est_gpu_ms: f64,
124}
125
126/// Frames below this size stay on the CPU (fixed GPU overheads dominate).
127pub const MIN_GPU_SAMPLES: u64 = 2_000_000;
128/// The GPU must win by this factor (estimates are noisy; CPU is the safer default).
129pub const GPU_MARGIN: f64 = 0.85;
130
131/// Estimate CPU and GPU times and choose.
132pub fn decide(class: AdapterClass, gpu: &GpuCalibration, cpu: &CpuCosts, w: &Workload) -> Decision {
133    let out_px = match w.output {
134        WorkOutput::Rgb { half_size: true, .. } => w.samples / 4,
135        _ => w.samples,
136    };
137    let (cpu_ms, gpu_kernel_ms, out_bytes) = match w.output {
138        WorkOutput::RawU16 => {
139            let c = if w.packed { w.samples as f64 / cpu.unpack_px_per_ms } else { 0.0 };
140            (c, w.samples as f64 / gpu.unpack_px_per_ms, w.samples * 2)
141        }
142        WorkOutput::Normalized => (
143            w.samples as f64 / cpu.normalize_px_per_ms
144                + if w.packed { w.samples as f64 / cpu.unpack_px_per_ms } else { 0.0 },
145            w.samples as f64 / gpu.normalize_px_per_ms,
146            w.samples * 4,
147        ),
148        WorkOutput::Rgb { bytes_per_px, half_size } => {
149            let unpack = if w.packed { w.samples as f64 / cpu.unpack_px_per_ms } else { 0.0 };
150            let (c, g) = if half_size {
151                (w.samples as f64 / cpu.develop_half_px_per_ms, w.samples as f64 / gpu.develop_half_px_per_ms)
152            } else {
153                (out_px as f64 / cpu.develop_px_per_ms, w.samples as f64 / gpu.develop_px_per_ms)
154            };
155            (c + unpack, g, out_px * bytes_per_px as u64)
156        }
157    };
158    let download_ms = if w.resident { 0.0 } else { out_bytes as f64 / gpu.download_bytes_per_ms };
159    // Tiled frames pay a per-band cost; halo rows add ~2 rows of upload per band.
160    let extra_bands = w.bands.max(1) as f64 - 1.0;
161    let gpu_ms = gpu.overhead_ms
162        + extra_bands * gpu.band_overhead_ms
163        + w.upload_bytes as f64 / gpu.upload_bytes_per_ms
164        + gpu_kernel_ms
165        + download_ms;
166    let mk = |use_gpu: bool, reason: String| Decision { use_gpu, reason, est_cpu_ms: cpu_ms, est_gpu_ms: gpu_ms };
167    if class == AdapterClass::Software {
168        return mk(false, "software adapter".into());
169    }
170    if w.resident {
171        return mk(true, "GPU-resident output requested".into());
172    }
173    if w.output == WorkOutput::RawU16 && !w.packed {
174        return mk(false, "raw output of an entropy-coded payload needs no GPU work".into());
175    }
176    if w.samples < MIN_GPU_SAMPLES {
177        return mk(false, format!("small frame ({:.1} MP)", w.samples as f64 / 1e6));
178    }
179    if gpu_ms < cpu_ms * GPU_MARGIN {
180        mk(true, format!("estimated GPU {gpu_ms:.1} ms < CPU {cpu_ms:.1} ms"))
181    } else {
182        mk(false, format!("estimated CPU {cpu_ms:.1} ms <= GPU {gpu_ms:.1} ms"))
183    }
184}
185
186fn calibration_disabled() -> bool {
187    #[cfg(not(target_arch = "wasm32"))]
188    {
189        std::env::var("FLASHRAW_GPU_CALIBRATION").map(|v| v.eq_ignore_ascii_case("static")).unwrap_or(false)
190    }
191    #[cfg(target_arch = "wasm32")]
192    {
193        false
194    }
195}
196
197impl GpuContext {
198    /// Measured characteristics of this device (one-time micro-benchmark, ~2048x2048
199    /// synthetic frame; cached). Returns static defaults if the benchmark fails or
200    /// `FLASHRAW_GPU_CALIBRATION=static`.
201    pub async fn calibration(self: &Arc<Self>) -> GpuCalibration {
202        if let Some(c) = self.calibration.get() {
203            return *c;
204        }
205        let c = if calibration_disabled() {
206            GpuCalibration::defaults(self.info.class)
207        } else {
208            self.measure().await.unwrap_or_else(|_| GpuCalibration::defaults(self.info.class))
209        };
210        *self.calibration.get_or_init(|| c)
211    }
212
213    /// The calibration if it has already been measured (never runs the benchmark).
214    pub fn cached_calibration(&self) -> Option<GpuCalibration> {
215        self.calibration.get().copied()
216    }
217
218    /// Blocking [`GpuContext::calibration`] (native only).
219    #[cfg(not(target_arch = "wasm32"))]
220    pub fn calibration_blocking(self: &Arc<Self>) -> GpuCalibration {
221        pollster::block_on(self.calibration())
222    }
223
224    async fn measure(self: &Arc<Self>) -> Result<GpuCalibration> {
225        const W: u32 = 2048;
226        const H: u32 = 2048;
227        let bits = 12u32;
228        let stride = (W * bits / 8) as usize;
229        let data: Vec<u8> = (0..stride * H as usize).map(|i| (i.wrapping_mul(2_654_435_761) >> 13) as u8).collect();
230        let payload = flashraw_core::PackedPayload {
231            data: &data,
232            width: W,
233            height: H,
234            bits,
235            order: PackOrder::MsbBytes,
236            row_stride: stride,
237        };
238        let frame = GpuFrame::from_packed(&payload);
239        let mut info = FrameInfo::new(W, H, bits as u8);
240        info.cfa = Some(CfaPattern::RGGB);
241        let dev = GpuOp::Develop(DevelopSpec {
242            crop: Rect::full(W, H),
243            half_size: false,
244            wb: [2.0, 1.0, 1.5, 1.0],
245            cam_to_rgb: [[1.6, -0.4, -0.2], [-0.2, 1.4, -0.2], [0.0, -0.5, 1.5]],
246            curve: None,
247            orientation: Orientation::Normal,
248            output: DevelopOutput::Rgba8,
249            levels: None,
250        });
251        // Warm-up compiles the pipelines and faults in allocations.
252        let warm = self.run(&frame, &info, &dev).await?;
253        drop(warm);
254        // Transfer = host copy + everything in the compute wall that is not kernel time
255        // (device copies on discrete GPUs, synchronisation).
256        let transfer =
257            |t: &crate::run::GpuTimings| t.upload_ms + (t.compute_ms - t.kernel_ms.unwrap_or(t.compute_ms)).max(0.0);
258        let mut best = (f64::MAX, f64::MAX, f64::MAX);
259        let mut wall_1 = f64::MAX;
260        for _ in 0..2 {
261            let r = self.run(&frame, &info, &dev).await?;
262            let (_, dl) = r.image.download_async().await?;
263            let k = r.timings.kernel_ms.unwrap_or(r.timings.compute_ms);
264            best.0 = best.0.min(transfer(&r.timings));
265            best.1 = best.1.min(k);
266            best.2 = best.2.min(dl);
267            wall_1 = wall_1.min(r.timings.upload_ms + r.timings.compute_ms);
268        }
269        // Per-band cost: the same frame in 8 bands (256 developed rows each).
270        let band_budget = Some(W as u64 * 4 * 256 + 4096);
271        let _ = self.run_banded(&frame, &info, &dev, band_budget).await?;
272        let mut wall_8 = f64::MAX;
273        let mut bands = 1;
274        for _ in 0..2 {
275            let r = self.run_banded(&frame, &info, &dev, band_budget).await?;
276            wall_8 = wall_8.min(r.timings.upload_ms + r.timings.compute_ms);
277            bands = r.timings.bands;
278        }
279        let band_overhead = if bands > 1 { ((wall_8 - wall_1) / (bands - 1) as f64).max(0.0) } else { 0.0 };
280        let (frame, info) = (&frame, &info);
281        let kernel_time = |op: GpuOp| async move {
282            let _ = self.run(frame, info, &op).await?; // warm-up
283            let mut t = f64::MAX;
284            for _ in 0..2 {
285                let r = self.run(frame, info, &op).await?;
286                t = t.min(r.timings.kernel_ms.unwrap_or(r.timings.compute_ms));
287            }
288            Ok::<f64, flashraw_core::Error>(t)
289        };
290        let unpack = kernel_time(GpuOp::Unpack).await?;
291        let normalize = kernel_time(GpuOp::Normalize).await?;
292        let GpuOp::Develop(spec) = dev else { unreachable!() };
293        let half = kernel_time(GpuOp::Develop(DevelopSpec { half_size: true, ..spec })).await?;
294        // Upload bandwidth from two sizes (separates the fixed per-upload cost).
295        let big: Vec<u8> = data.iter().cycle().take(data.len() * 4).copied().collect();
296        let big_payload = flashraw_core::PackedPayload { data: &big, height: H * 4, ..payload };
297        let big_frame = GpuFrame::from_packed(&big_payload);
298        let big_info = FrameInfo::new(W, H * 4, bits as u8);
299        let mut up_big = f64::MAX;
300        for _ in 0..3 {
301            let r = self.run(&big_frame, &big_info, &GpuOp::Unpack).await?;
302            up_big = up_big.min(transfer(&r.timings));
303        }
304        let upload_bw = (big.len() - data.len()) as f64 / (up_big - best.0).max(1e-2);
305        let upload_bw = upload_bw.min(big.len() as f64 / up_big.max(1e-3) * 4.0);
306        let upload_fixed = (best.0 - data.len() as f64 / upload_bw).max(0.0);
307        let px = (W * H) as f64;
308        Ok(GpuCalibration {
309            upload_bytes_per_ms: upload_bw,
310            download_bytes_per_ms: px * 4.0 / best.2.max(1e-3),
311            develop_px_per_ms: px / best.1.max(1e-3),
312            develop_half_px_per_ms: px / half.max(1e-3),
313            unpack_px_per_ms: px / unpack.max(1e-3),
314            normalize_px_per_ms: px / normalize.max(1e-3),
315            overhead_ms: upload_fixed + 0.2,
316            band_overhead_ms: band_overhead,
317            measured: true,
318        })
319    }
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325
326    fn cpu() -> CpuCosts {
327        CpuCosts {
328            unpack_px_per_ms: 2e6,
329            develop_px_per_ms: 0.3e6,
330            develop_half_px_per_ms: 0.5e6,
331            normalize_px_per_ms: 1e6,
332        }
333    }
334
335    fn work(output: WorkOutput, packed: bool, mp: f64) -> Workload {
336        let samples = (mp * 1e6) as u64;
337        Workload { samples, upload_bytes: samples * 14 / 8, packed, output, resident: false, bands: 1 }
338    }
339
340    #[test]
341    fn software_never_auto() {
342        let g = GpuCalibration::defaults(AdapterClass::Discrete);
343        let d = decide(
344            AdapterClass::Software,
345            &g,
346            &cpu(),
347            &work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 45.0),
348        );
349        assert!(!d.use_gpu);
350    }
351
352    #[test]
353    fn small_frames_stay_on_cpu() {
354        let g = GpuCalibration::defaults(AdapterClass::Discrete);
355        let d = decide(
356            AdapterClass::Discrete,
357            &g,
358            &cpu(),
359            &work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 1.0),
360        );
361        assert!(!d.use_gpu);
362    }
363
364    #[test]
365    fn raw_entropy_coded_stays_on_cpu() {
366        let g = GpuCalibration::defaults(AdapterClass::Discrete);
367        let d = decide(AdapterClass::Discrete, &g, &cpu(), &work(WorkOutput::RawU16, false, 45.0));
368        assert!(!d.use_gpu);
369    }
370
371    #[test]
372    fn large_rgb_on_fast_gpu() {
373        let g = GpuCalibration::defaults(AdapterClass::Discrete);
374        let d = decide(
375            AdapterClass::Discrete,
376            &g,
377            &cpu(),
378            &work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 45.0),
379        );
380        assert!(d.use_gpu, "{d:?}");
381        // A much faster CPU flips the decision.
382        let fast = CpuCosts {
383            unpack_px_per_ms: 20e6,
384            develop_px_per_ms: 10e6,
385            develop_half_px_per_ms: 20e6,
386            normalize_px_per_ms: 20e6,
387        };
388        let d = decide(
389            AdapterClass::Discrete,
390            &g,
391            &fast,
392            &work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 45.0),
393        );
394        assert!(!d.use_gpu, "{d:?}");
395    }
396
397    #[test]
398    fn tiled_frames_cost_more_but_still_use_the_gpu() {
399        let g = GpuCalibration::defaults(AdapterClass::Discrete);
400        let w = work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 150.0);
401        let one = decide(AdapterClass::Discrete, &g, &cpu(), &w);
402        let tiled = decide(AdapterClass::Discrete, &g, &cpu(), &Workload { bands: 5, ..w });
403        assert!(tiled.use_gpu, "{tiled:?}");
404        assert!((tiled.est_gpu_ms - one.est_gpu_ms - 4.0 * g.band_overhead_ms).abs() < 1e-9);
405    }
406
407    #[test]
408    fn resident_output_forces_gpu() {
409        let g = GpuCalibration::defaults(AdapterClass::Integrated);
410        let mut w = work(WorkOutput::Rgb { bytes_per_px: 4, half_size: false }, true, 0.5);
411        w.resident = true;
412        assert!(decide(AdapterClass::Integrated, &g, &cpu(), &w).use_gpu);
413    }
414}
flashraw_gpu/policy.rs

flashraw_gpu/
policy.rs
