Skip to main content

av_denoise_core/
denoiser.rs

1use std::collections::VecDeque;
2
3use cubecl::Runtime;
4use cubecl::prelude::ComputeClient;
5
6use crate::accelerate::Accelerator;
7use crate::device::Device;
8use crate::nl4d::{Nl4dDenoiser, Nl4dParams};
9#[cfg(test)]
10use crate::nlmeans::MotionEstimation;
11use crate::nlmeans::{
12    ChannelMode,
13    Depth,
14    HqParams,
15    MotionCompensationMode,
16    MotionSearch,
17    NlmDenoiser,
18    NlmParams,
19    Pending,
20    PrefilterMode,
21    TryWait,
22    hq_default_strength,
23    validate_dimensions,
24};
25use crate::sniff::sniff_best_accelerator;
26
27/// How a [`Denoiser`] should be set up.
28///
29/// Build one with `DenoiserOptions::builder()`. Every field has a
30/// default, so only the parts you care about need naming.
31///
32/// Only the settings every algorithm reads live here. Everything else
33/// belongs to whichever [`Algorithm`] variant actually uses it.
34#[derive(Debug, Clone, bon::Builder)]
35pub struct DenoiserOptions {
36    /// Which channels of the frame to denoise.
37    #[builder(default = ChannelMode::Yuv)]
38    pub channel_mode: ChannelMode,
39    /// Whether to clean each frame on its own or across a temporal
40    /// window.
41    #[builder(default = DenoisingMode::Spacial)]
42    pub mode: DenoisingMode,
43    /// Which algorithm to run, along with the settings only that
44    /// algorithm reads.
45    #[builder(default)]
46    pub algorithm: Algorithm,
47    /// What format denoised frames come back in.
48    #[builder(default = OutputFormat::F32)]
49    pub output_format: OutputFormat,
50}
51
52/// What a denoiser hands back from [`Denoiser::recv_frame`],
53/// [`Denoiser::try_recv_frame`] and [`Denoiser::flush`].
54///
55/// The GPU only ever holds normalised `f32`. [`Depth`] is a wire concept,
56/// so a denoiser that returns wire bytes is told its depth when it is
57/// built rather than at each call.
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum OutputFormat {
60    /// Normalised `f32`, one value per channel per pixel.
61    F32,
62    /// Wire bytes at this depth, quantised on the GPU.
63    Wire { depth: Depth },
64}
65
66/// One denoised frame, in whichever format its denoiser was built for.
67#[derive(Debug, Clone, PartialEq)]
68pub enum FrameOutput {
69    /// Normalised `f32`, `width * height * channels` values.
70    F32(Vec<f32>),
71    /// Wire bytes, interleaved except for a chroma pair, which is laid
72    /// out as U's whole region followed by V's.
73    Wire(Vec<u8>),
74}
75
76impl FrameOutput {
77    /// The `f32` frame, or `None` if this came from a wire-mode denoiser.
78    pub fn into_f32(self) -> Option<Vec<f32>> {
79        match self {
80            Self::F32(v) => Some(v),
81            Self::Wire(_) => None,
82        }
83    }
84
85    /// The wire bytes, or `None` if this came from an `f32`-mode denoiser.
86    pub fn into_wire(self) -> Option<Vec<u8>> {
87        match self {
88            Self::Wire(v) => Some(v),
89            Self::F32(_) => None,
90        }
91    }
92
93    /// Borrows the `f32` frame, or `None` if this came from a wire-mode
94    /// denoiser.
95    pub fn as_f32(&self) -> Option<&[f32]> {
96        match self {
97            Self::F32(v) => Some(v),
98            Self::Wire(_) => None,
99        }
100    }
101
102    /// Borrows the wire bytes, or `None` if this came from an `f32`-mode
103    /// denoiser.
104    pub fn as_wire(&self) -> Option<&[u8]> {
105        match self {
106            Self::Wire(v) => Some(v),
107            Self::F32(_) => None,
108        }
109    }
110}
111
112/// Which denoising algorithm to run.
113///
114/// Each variant carries its own settings, so a knob one algorithm has no
115/// use for cannot be set on it.
116#[derive(Debug, Copy, Clone, PartialEq)]
117pub enum Algorithm {
118    /// The fast NLMeans path, with fixed weighting and no noise
119    /// measurement.
120    Nlmeans(NlmeansOptions),
121    /// NLMeans with its weighting matched to the measured noise level.
122    ///
123    /// This also uses a different default `strength`, one that adapts to
124    /// the temporal radius and the plane being denoised. See
125    /// [`crate::nlmeans::hq_default_strength`].
126    NlmeansHq(NlmeansHqOptions),
127    /// Groups 8x8 patches across the motion-compensated temporal window
128    /// itself, rather than filtering with NLM first and grouping within
129    /// one frame afterward.
130    ///
131    /// No NLM weighting pass ever runs, so none of the NLM knobs appear
132    /// on [`Nl4dOptions`].
133    Nl4d(Nl4dOptions),
134}
135
136impl Default for Algorithm {
137    fn default() -> Self {
138        Self::Nlmeans(NlmeansOptions::default())
139    }
140}
141
142/// Settings for [`Algorithm::Nlmeans`].
143#[derive(Debug, Copy, Clone, Default, PartialEq)]
144pub struct NlmeansOptions {
145    /// Which reference image the NLM weights are computed against.
146    ///
147    /// `None`, the default, compares patches on the noisy input
148    /// directly. Every other mode costs one extra GPU pass per frame.
149    pub prefilter: PrefilterMode,
150    /// Whether temporal denoising follows motion between frames.
151    ///
152    /// `None`, the default, turns motion compensation off. `Mvtools`
153    /// warps temporal neighbours into line with the centre frame before
154    /// the NLM weighting runs.
155    ///
156    /// Only has an effect when [`DenoiserOptions::mode`] is
157    /// `Temporal { .. }`.
158    pub motion_compensation: MotionCompensationMode,
159    /// Overrides for the NLM search radius, patch radius, strength, and
160    /// self-weight.
161    pub tuning: NlmTuning,
162}
163
164/// Settings for [`Algorithm::NlmeansHq`].
165#[derive(Debug, Copy, Clone, Default, PartialEq)]
166pub struct NlmeansHqOptions {
167    /// Everything the fast path takes, which HQ takes too.
168    pub nlm: NlmeansOptions,
169    /// The noise measurement and confidence weighting HQ adds on top.
170    pub hq: HqParams,
171}
172
173/// Settings for [`Algorithm::Nl4d`].
174///
175/// nl4d runs the HQ front end only for its machinery, the frame ring,
176/// the motion field, and the noise estimate. Nothing weights or averages
177/// patches the NLM way, so the NLM knobs are absent here and the fields
178/// below are the whole surface.
179///
180/// The temporal radius comes from [`DenoiserOptions::mode`], which has to
181/// be `Temporal { .. }`. Motion tracking is always on, because the
182/// grouping kernel reads the motion field and confidence scores it
183/// produces.
184///
185/// `lambda_ht` has a per-plane default. `None` resolves through
186/// [`nl4d_default_lambda_ht`] once the plane being denoised is known.
187/// `lambda_ht_scale` then multiplies whichever value that resolves to.
188///
189/// Every other default comes from [`Nl4dParams::default`].
190#[derive(Debug, Copy, Clone, PartialEq)]
191pub struct Nl4dOptions {
192    /// How motion between frames is tracked.
193    pub motion: MotionSearch,
194    /// A fixed noise standard deviation in `[0, 1]` units, replacing the
195    /// automatic per-frame estimate.
196    ///
197    /// `None`, the default, measures the noise in each pushed frame and
198    /// smooths it over time.
199    pub sigma: Option<f32>,
200    /// A multiplier applied to the measured noise level before anything
201    /// reads it. Defaults to 1.0.
202    ///
203    /// This does nothing when `sigma` pins the noise level, because the
204    /// estimator never runs in that case.
205    pub sigma_scale: f32,
206    /// A multiplier on the per-block mismatch threshold, which sets how
207    /// much extra SAD a block tolerates before its confidence starts to
208    /// fall. Defaults to 1.0.
209    ///
210    /// Higher values tolerate larger mismatches.
211    pub thsad_scale: f32,
212    /// Half-width of the refine window searched around each neighbour
213    /// frame's motion-predicted position, in `1..=4`. Defaults to 2.
214    pub refine: u32,
215    /// Half-width of the spatial candidate window searched in the centre
216    /// frame, in `1..=16`. Defaults to 9.
217    pub spatial_radius: u32,
218    /// Hard-threshold multiplier on the propagated coefficient sigma.
219    /// Higher removes more noise and more fine detail.
220    ///
221    /// `None` resolves through [`nl4d_default_lambda_ht`], which returns
222    /// a different value for luma than for chroma.
223    pub lambda_ht: Option<f32>,
224    /// A multiplier applied to the resolved `lambda_ht`. Defaults to
225    /// 1.0.
226    ///
227    /// It scales an explicit `lambda_ht` and the calibrated per-plane
228    /// default alike, so one value moves both planes together. Has to
229    /// be finite and in `[0.1, 10.0]`.
230    pub lambda_ht_scale: f32,
231    /// The confidence floor below which a whole neighbour block is
232    /// skipped rather than scored, in `[0, 1)`. Defaults to 0.05. Only
233    /// affects how much compute a submit spends, never which candidates
234    /// are admitted once they are scored.
235    pub c_min: f32,
236    /// A multiplier on the mismatch variance a poorly matched temporal
237    /// member carries into the hard threshold. Defaults to 1.0.
238    ///
239    /// The variance grows with the square of this. The mechanism
240    /// saturates well before the top of its accepted range, see
241    /// [`crate::nl4d::Nl4dParams::mismatch_scale`].
242    pub mismatch_scale: f32,
243    /// Whether a temporal member's mismatch variance reaches the
244    /// hard-threshold shrinkage at all. Defaults to `true`. See
245    /// [`crate::nl4d::Nl4dParams::confidence_variance`].
246    pub confidence_variance: bool,
247    /// Estimates noise fresh from each frame's own window instead of
248    /// smoothing it across the whole stream's history. Defaults to
249    /// `false`, matching every calibrated preset.
250    ///
251    /// `av-denoise-vs` turns this on unconditionally, because a
252    /// VapourSynth filter has to return the same pixels for a frame no
253    /// matter what order frames were requested in, and history-dependent
254    /// estimation breaks that guarantee under random access. See
255    /// [`HqParams::windowed_noise_estimation`].
256    pub windowed_noise_estimation: bool,
257}
258
259impl Default for Nl4dOptions {
260    fn default() -> Self {
261        let defaults = Nl4dParams::default();
262        let hq = HqParams::default();
263        Self {
264            motion: MotionSearch::default(),
265            sigma: hq.sigma_override,
266            sigma_scale: hq.sigma_scale,
267            thsad_scale: hq.thsad_scale,
268            refine: defaults.refine,
269            spatial_radius: defaults.spatial_radius,
270            // Resolved per plane by `nl4d_default_lambda_ht` at
271            // construction time, once the plane being denoised is
272            // known.
273            lambda_ht: None,
274            lambda_ht_scale: 1.0,
275            c_min: defaults.c_min,
276            mismatch_scale: defaults.mismatch_scale,
277            confidence_variance: defaults.confidence_variance,
278            windowed_noise_estimation: false,
279        }
280    }
281}
282
283impl Nl4dOptions {
284    /// The front end's HQ parameters for this configuration.
285    ///
286    /// `temporal_confidence` is always on, because the grouping kernel
287    /// reads the confidence scores it produces. The two strength-related
288    /// switches keep their defaults, since nl4d never runs a weighting
289    /// pass for them to affect.
290    fn to_hq_params(self) -> HqParams {
291        HqParams {
292            sigma_override: self.sigma,
293            sigma_scale: self.sigma_scale,
294            thsad_scale: self.thsad_scale,
295            temporal_confidence: true,
296            windowed_noise_estimation: self.windowed_noise_estimation,
297            ..HqParams::default()
298        }
299    }
300}
301
302/// The default `lambda_ht` for nl4d's hard-threshold stage, per plane.
303///
304/// `lambda_ht` is how many standard deviations of estimated noise a
305/// transform coefficient has to clear to survive. Raising it removes more
306/// noise and more fine detail with it, so the value is a trade rather
307/// than an optimum.
308///
309/// Luma gets 5.3, picked by eye from rendered comparisons on real grain
310/// and deliberately biased toward keeping detail. Higher values remove
311/// visibly more noise, but not enough to be worth what they cost in
312/// texture.
313///
314/// `ChannelMode::Yuv` reads the luma value, on the same "a fused pass is
315/// dominated by luma" assumption [`hq_default_strength`]
316/// makes for its own Yuv case.
317///
318/// Chroma gets 4.2, picked the same way from the chroma residuals with
319/// luma pinned at 5.3.
320pub fn nl4d_default_lambda_ht(channels: ChannelMode) -> f32 {
321    match channels {
322        ChannelMode::Luma | ChannelMode::Yuv => 5.3,
323        ChannelMode::Chroma => 4.2,
324    }
325}
326
327/// Resolves `Nl4dOptions.lambda_ht` for one plane, falling back to
328/// [`nl4d_default_lambda_ht`] when the caller left it unset, then
329/// applies `lambda_ht_scale`.
330///
331/// The scale multiplies an explicit value and the calibrated default
332/// alike, so it moves both planes together whether or not one of them
333/// is pinned.
334///
335/// The range check lives here rather than in [`Nl4dParams`],
336/// which only ever sees the product. A scale of 0 would surface there as
337/// a complaint about `lambda_ht`, naming a knob the caller never set.
338fn resolve_lambda_ht(opts: &Nl4dOptions, channels: ChannelMode) -> Result<f32, String> {
339    if !(opts.lambda_ht_scale.is_finite() && (0.1..=10.0).contains(&opts.lambda_ht_scale)) {
340        return Err(format!(
341            "lambda_ht_scale must be finite and in [0.1, 10.0], got {}",
342            opts.lambda_ht_scale
343        ));
344    }
345
346    let lambda_ht = opts.lambda_ht.unwrap_or_else(|| nl4d_default_lambda_ht(channels));
347
348    Ok(lambda_ht * opts.lambda_ht_scale)
349}
350
351/// Speed vs quality dial.
352///
353/// Each denoising family reads the same dial and fills in its own knobs
354/// from it. For `nlmeans` that is [`nlmeans_variant_for`],
355/// [`nlmeans_temporal_radius_for`], and [`nlmeans_search_radius_for`].
356/// For `nl4d` it is [`nl4d_temporal_radius_for`] and
357/// [`nl4d_spatial_radius_for`].
358///
359/// Both front ends parse the same names from this one type, so a preset
360/// resolves to the same dials everywhere it is used.
361#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, strum_macros::EnumString)]
362#[strum(ascii_case_insensitive)]
363pub enum Preset {
364    /// Fastest and lowest quality.
365    Veryfast,
366    /// One step up from `veryfast`.
367    Fast,
368    /// The default, favouring quality over speed.
369    #[default]
370    Base,
371    /// One step down from `veryslow`.
372    Slow,
373    /// Slowest and highest quality.
374    Veryslow,
375}
376
377/// Which nlmeans implementation a preset, or an explicit choice, selects.
378#[derive(Debug, Copy, Clone, PartialEq, Eq, strum_macros::EnumString)]
379#[strum(ascii_case_insensitive)]
380pub enum NlmeansVariant {
381    /// The fast path. Fixed weighting, no noise measurement.
382    Fast,
383    /// Quality focused. Calibrates its weighting to the noise level,
384    /// measured automatically per frame.
385    Hq,
386}
387
388/// Which [`NlmeansVariant`] a preset runs.
389pub fn nlmeans_variant_for(preset: Preset) -> NlmeansVariant {
390    match preset {
391        Preset::Veryfast => NlmeansVariant::Fast,
392        Preset::Fast | Preset::Base | Preset::Slow | Preset::Veryslow => NlmeansVariant::Hq,
393    }
394}
395
396/// How many neighbouring frames on each side `nlmeans` looks at, at a
397/// preset.
398pub fn nlmeans_temporal_radius_for(preset: Preset) -> u32 {
399    match preset {
400        Preset::Veryfast => 0,
401        Preset::Fast => 1,
402        Preset::Base => 2,
403        Preset::Slow => 4,
404        Preset::Veryslow => 8,
405    }
406}
407
408/// How far `nlmeans` looks for similar patches inside a frame, at a
409/// preset.
410pub fn nlmeans_search_radius_for(preset: Preset) -> u32 {
411    match preset {
412        Preset::Veryfast | Preset::Fast | Preset::Base => 2,
413        Preset::Slow | Preset::Veryslow => 4,
414    }
415}
416
417/// How far the temporal window reaches at each preset, for `nl4d`.
418///
419/// Unlike `nlmeans`, `veryfast` keeps a 1-frame window rather than
420/// dropping to 0, because nl4d has nothing to do without neighbouring
421/// frames to group against.
422pub fn nl4d_temporal_radius_for(preset: Preset) -> u32 {
423    match preset {
424        Preset::Veryfast | Preset::Fast => 1,
425        Preset::Base => 2,
426        Preset::Slow => 4,
427        Preset::Veryslow => 8,
428    }
429}
430
431/// How wide the centre frame's candidate search is at each preset, for
432/// `nl4d`.
433///
434/// `veryfast` shares its temporal radius with `fast`, so this is what
435/// separates them. The window covers `(2 * radius + 1)^2` positions, so
436/// 6 searches a little over half the candidates 9 does.
437///
438/// Every preset from `fast` up uses the library default. Widening it
439/// further at the slow end costs quadratically and has not been measured
440/// to be worth it.
441pub fn nl4d_spatial_radius_for(preset: Preset) -> u32 {
442    match preset {
443        Preset::Veryfast => 6,
444        Preset::Fast | Preset::Base | Preset::Slow | Preset::Veryslow => {
445            Nl4dOptions::default().spatial_radius
446        },
447    }
448}
449
450/// Whether a frame is cleaned on its own or alongside its neighbours.
451#[derive(Debug, Copy, Clone, Eq, PartialEq)]
452pub enum DenoisingMode {
453    /// Cleans each frame using only its own pixels.
454    Spacial,
455    /// Cleans each frame using a window of `2 * radius + 1` frames.
456    Temporal { radius: u32 },
457}
458
459/// NLM tuning knobs.
460///
461/// Every field is optional. Whatever is left unset falls back to the
462/// library default.
463#[derive(Debug, Copy, Clone, Default, PartialEq)]
464pub struct NlmTuning {
465    pub search_radius: Option<u32>,
466    pub patch_radius: Option<u32>,
467    pub strength: Option<f32>,
468    pub self_weight: Option<f32>,
469}
470
471impl DenoiserOptions {
472    /// Turns this option set into the low-level [`NlmParams`] a backend
473    /// denoiser is built from.
474    ///
475    /// For nl4d this describes the front end only, since nl4d's own
476    /// grouping stage is configured from [`Nl4dOptions`] separately in
477    /// [`build_engine`].
478    ///
479    /// Whichever default `strength` applies is folded in here. For the
480    /// HQ algorithm that comes from
481    /// [`crate::nlmeans::hq_default_strength`].
482    ///
483    /// This is public so callers building per-plane options, and tests,
484    /// can read the resolved values without building a real `Denoiser`.
485    #[doc(hidden)]
486    pub fn to_nlm_params(&self) -> NlmParams {
487        let temporal_radius = match self.mode {
488            DenoisingMode::Spacial => 0,
489            DenoisingMode::Temporal { radius } => radius,
490        };
491
492        match self.algorithm {
493            Algorithm::Nlmeans(opts) => self.nlm_params_for(opts, None, temporal_radius),
494            Algorithm::NlmeansHq(opts) => self.nlm_params_for(opts.nlm, Some(opts.hq), temporal_radius),
495            // nl4d never runs a weighting pass, so `strength`,
496            // `search_radius`, `patch_radius`, and `self_weight` stay at
497            // their library defaults and no prefilter is built.
498            Algorithm::Nl4d(opts) => NlmParams {
499                channels: self.channel_mode,
500                motion_compensation: opts.motion.into(),
501                temporal_radius,
502                hq: Some(opts.to_hq_params()),
503                ..NlmParams::default()
504            },
505        }
506    }
507
508    /// [`Self::to_nlm_params`] for whichever of the two NLM algorithms
509    /// is running, with `hq` set only for the quality one.
510    fn nlm_params_for(&self, opts: NlmeansOptions, hq: Option<HqParams>, temporal_radius: u32) -> NlmParams {
511        // An explicit `strength` always wins, whether it came straight
512        // from `NlmTuning` or from a per-plane override the caller
513        // already folded in.
514        //
515        // Otherwise the default depends on `auto_strength`. With it on,
516        // HQ reads `strength` as a multiplier on the measured noise
517        // level, so it needs its own calibrated default rather than the
518        // fast path's absolute FFmpeg-style one. That calibrated default
519        // also varies with the temporal radius and with the plane
520        // `channel_mode` names, because each per-plane `Denoiser`
521        // carries its own channel mode.
522        //
523        // With auto-strength off, HQ reads `strength` as an absolute
524        // value just like the fast path, so it falls back to the same
525        // absolute default.
526        let strength = opts.tuning.strength.unwrap_or(match hq {
527            Some(hq) if hq.auto_strength => hq_default_strength(self.channel_mode, temporal_radius),
528            _ => NlmParams::default().strength,
529        });
530
531        let defaults = NlmParams::default();
532        NlmParams {
533            channels: self.channel_mode,
534            prefilter: opts.prefilter,
535            motion_compensation: opts.motion_compensation,
536            temporal_radius,
537            hq,
538            strength,
539            search_radius: opts.tuning.search_radius.unwrap_or(defaults.search_radius),
540            patch_radius: opts.tuning.patch_radius.unwrap_or(defaults.patch_radius),
541            self_weight: opts.tuning.self_weight.unwrap_or(defaults.self_weight),
542        }
543    }
544}
545
546/// Errors reported by the high-level [`Denoiser`].
547#[derive(Debug, thiserror::Error)]
548pub enum DenoiserError {
549    /// An earlier denoised frame has not been collected yet, so pushing
550    /// again would overwrite it in the double-buffered output slot.
551    ///
552    /// Call [`Denoiser::recv_frame`] or [`Denoiser::try_recv_frame`],
553    /// then retry the same `push_frame` call.
554    #[error("denoiser queue is full, collect the pending frame before pushing more")]
555    QueueFull,
556    /// An earlier call failed, so how many frames are in flight is no
557    /// longer known and later output would not line up with its input.
558    ///
559    /// Call [`Denoiser::reset_stream`] to start a fresh stream, or drop the denoiser.
560    #[error("denoiser failed earlier, reset the stream before using it again")]
561    Poisoned,
562    /// None of the accelerators in the priority list could be started.
563    #[error("no accelerator from the priority list is available")]
564    NoAcceleratorAvailable,
565    /// Anything else, wrapping the internal `anyhow` errors raised by
566    /// kernel dispatch and readback.
567    #[error(transparent)]
568    Other(#[from] anyhow::Error),
569}
570
571/// Either denoiser a `Backend` runtime arm can hold.
572///
573/// This keeps `Backend`'s own match arms at one line each. Without it,
574/// adding a second denoiser type would multiply the runtime arms instead
575/// of fanning out once here.
576enum Engine<R: Runtime> {
577    Nlm(Box<NlmDenoiser<R>>),
578    Nl4d(Box<Nl4dDenoiser<R>>),
579}
580
581impl<R: Runtime> Engine<R> {
582    fn is_nl4d(&self) -> bool {
583        matches!(self, Self::Nl4d(_))
584    }
585
586    fn push_frame(&mut self, frame: &[f32]) {
587        match self {
588            Self::Nlm(d) => d.push_frame(frame),
589            Self::Nl4d(d) => d.push_frame(frame),
590        }
591    }
592
593    fn push_frame_wire(&mut self, planes: &[&[u8]], depth: Depth) {
594        match self {
595            Self::Nlm(d) => d.push_frame_wire(planes, depth),
596            Self::Nl4d(d) => d.push_frame_wire(planes, depth),
597        }
598    }
599
600    fn denoise_submit(&mut self) -> Result<Option<Pending<R>>, anyhow::Error> {
601        match self {
602            Self::Nlm(d) => d.denoise_submit(),
603            // `Nl4dDenoiser::denoise_submit` already returns
604            // `DenoiserError` rather than `anyhow::Error`, so this leans
605            // on `DenoiserError`'s own `anyhow::Error` conversion instead
606            // of re-wrapping it.
607            Self::Nl4d(d) => d.denoise_submit().map_err(anyhow::Error::from),
608        }
609    }
610
611    #[cfg(test)]
612    fn wire_outputs(&self) -> Option<&[cubecl::server::Handle; 2]> {
613        match self {
614            Self::Nlm(d) => d.wire_outputs_for_test(),
615            Self::Nl4d(d) => d.wire_outputs_for_test(),
616        }
617    }
618
619    fn flush(&mut self, sink: impl FnMut(&FrameOutput)) -> Result<(), anyhow::Error> {
620        match self {
621            Self::Nlm(d) => d.flush(sink),
622            Self::Nl4d(d) => d.flush(sink).map_err(anyhow::Error::from),
623        }
624    }
625
626    fn reset_stream(&mut self) {
627        match self {
628            Self::Nlm(d) => d.reset_stream_state(),
629            Self::Nl4d(d) => d.reset_stream(),
630        }
631    }
632}
633
634/// Builds whichever [`Engine`] `algorithm` calls for.
635///
636/// `Algorithm::Nl4d` carries its own grouping tuning, which is not part
637/// of `NlmParams`, so it is read from `algorithm` directly rather than
638/// from `params`. This is also where an unset `lambda_ht` picks up its
639/// calibrated per-plane default (`resolve_lambda_ht`), the same way
640/// `to_nlm_params` resolves HQ's calibrated `strength`, since this is
641/// the first point construction has both `opts` and `params.channels`
642/// together.
643fn build_engine<R: Runtime>(
644    client: &ComputeClient<R>,
645    algorithm: &Algorithm,
646    params: NlmParams,
647    width: u32,
648    height: u32,
649    output_format: OutputFormat,
650) -> Result<Engine<R>, DenoiserError> {
651    match algorithm {
652        Algorithm::Nl4d(opts) => {
653            // nl4d groups patches across neighbouring frames, so there
654            // is nothing for it to do without a temporal window.
655            if params.temporal_radius == 0 {
656                return Err(DenoiserError::Other(anyhow::anyhow!(
657                    "nl4d needs a temporal window, set DenoiserOptions::mode to \
658                     DenoisingMode::Temporal"
659                )));
660            }
661
662            let lambda_ht = resolve_lambda_ht(opts, params.channels)
663                .map_err(|e| DenoiserError::Other(anyhow::anyhow!(e)))?;
664            let nl4d_params = Nl4dParams {
665                temporal_radius: params.temporal_radius,
666                nlm: params,
667                refine: opts.refine,
668                spatial_radius: opts.spatial_radius,
669                lambda_ht,
670                c_min: opts.c_min,
671                mismatch_scale: opts.mismatch_scale,
672                confidence_variance: opts.confidence_variance,
673            };
674            let denoiser =
675                Nl4dDenoiser::with_output_format(client, nl4d_params, width, height, output_format)
676                    .map_err(|e| DenoiserError::Other(anyhow::anyhow!(e)))?;
677            Ok(Engine::Nl4d(Box::new(denoiser)))
678        },
679        Algorithm::Nlmeans(_) | Algorithm::NlmeansHq(_) => Ok(Engine::Nlm(Box::new(
680            NlmDenoiser::with_output_format(client, params, width, height, output_format),
681        ))),
682    }
683}
684
685enum Backend {
686    #[cfg(feature = "cuda")]
687    Cuda(Engine<cubecl::cuda::CudaRuntime>),
688    #[cfg(feature = "rocm")]
689    Rocm(Engine<cubecl::hip::HipRuntime>),
690    #[cfg(any(feature = "vulkan", feature = "metal"))]
691    Wgpu(Engine<cubecl::wgpu::WgpuRuntime>),
692}
693
694impl Backend {
695    fn is_nl4d(&self) -> bool {
696        match self {
697            #[cfg(feature = "cuda")]
698            Self::Cuda(e) => e.is_nl4d(),
699            #[cfg(feature = "rocm")]
700            Self::Rocm(e) => e.is_nl4d(),
701            #[cfg(any(feature = "vulkan", feature = "metal"))]
702            Self::Wgpu(e) => e.is_nl4d(),
703        }
704    }
705
706    #[cfg(test)]
707    fn wire_outputs(&self) -> Option<&[cubecl::server::Handle; 2]> {
708        match self {
709            #[cfg(feature = "cuda")]
710            Self::Cuda(e) => e.wire_outputs(),
711            #[cfg(feature = "rocm")]
712            Self::Rocm(e) => e.wire_outputs(),
713            #[cfg(any(feature = "vulkan", feature = "metal"))]
714            Self::Wgpu(e) => e.wire_outputs(),
715        }
716    }
717}
718
719enum BackendPending {
720    #[cfg(feature = "cuda")]
721    Cuda(Pending<cubecl::cuda::CudaRuntime>),
722    #[cfg(feature = "rocm")]
723    Rocm(Pending<cubecl::hip::HipRuntime>),
724    #[cfg(any(feature = "vulkan", feature = "metal"))]
725    Wgpu(Pending<cubecl::wgpu::WgpuRuntime>),
726}
727
728impl BackendPending {
729    fn wait(self) -> Result<FrameOutput, anyhow::Error> {
730        match self {
731            #[cfg(feature = "cuda")]
732            Self::Cuda(p) => p.wait(),
733            #[cfg(feature = "rocm")]
734            Self::Rocm(p) => p.wait(),
735            #[cfg(any(feature = "vulkan", feature = "metal"))]
736            Self::Wgpu(p) => p.wait(),
737        }
738    }
739
740    /// Polls the readback once. `Ok(Ok(frame))` is a landed frame,
741    /// `Ok(Err(self))` is a readback still in flight.
742    fn try_wait(self) -> Result<Result<FrameOutput, Self>, anyhow::Error> {
743        match self {
744            #[cfg(feature = "cuda")]
745            Self::Cuda(p) => match p.try_wait()? {
746                TryWait::Ready(frame) => Ok(Ok(frame)),
747                TryWait::NotReady(p) => Ok(Err(Self::Cuda(p))),
748            },
749            #[cfg(feature = "rocm")]
750            Self::Rocm(p) => match p.try_wait()? {
751                TryWait::Ready(frame) => Ok(Ok(frame)),
752                TryWait::NotReady(p) => Ok(Err(Self::Rocm(p))),
753            },
754            #[cfg(any(feature = "vulkan", feature = "metal"))]
755            Self::Wgpu(p) => match p.try_wait()? {
756                TryWait::Ready(frame) => Ok(Ok(frame)),
757                TryWait::NotReady(p) => Ok(Err(Self::Wgpu(p))),
758            },
759        }
760    }
761}
762
763/// How many readbacks the high-level [`Denoiser`] keeps in flight at
764/// once.
765///
766/// This has to match the backend's output-handle count, which is two.
767/// Going past it would reuse the oldest pending frame's output handle
768/// and quietly corrupt the results.
769pub const MAX_PENDING: usize = 2;
770
771/// How many source frames a windowed operation needs behind and ahead
772/// of its target frame, target frame itself not counted in either
773/// number.
774///
775/// `reseed` needs exactly `behind + 1 + ahead` frames, oldest first,
776/// with the target frame sitting at index `behind`. This is what tells
777/// a caller like `reseed` how wide a window to build, and it varies by
778/// algorithm because nl4d's own cross-frame accumulator needs more
779/// forward context than the NLM algorithms do. See
780/// [`Denoiser::window_span`].
781#[derive(Debug, Clone, Copy, PartialEq, Eq)]
782pub struct WindowSpan {
783    /// How many frames older than the target the window must include.
784    pub behind: usize,
785    /// How many frames newer than the target the window must include.
786    pub ahead: usize,
787}
788
789impl WindowSpan {
790    /// The full window size this span describes, target frame
791    /// included: `behind + 1 + ahead`.
792    pub fn frame_count(&self) -> usize {
793        self.behind + 1 + self.ahead
794    }
795}
796
797/// A stateful denoiser that cleans a stream of frames.
798///
799/// Push frames in order with [`push_frame`](Self::push_frame) and
800/// collect the cleaned ones with [`recv_frame`](Self::recv_frame) or
801/// [`try_recv_frame`](Self::try_recv_frame).
802///
803/// At the end of the stream call [`flush`](Self::flush) to drain
804/// whatever temporal context is left.
805///
806/// Input frames are `f32` values in `[0, 1]`, laid out as
807/// `width * height * channels`. Output comes back as a [`FrameOutput`]
808/// in whichever [`OutputFormat`] the options named.
809///
810/// ```no_run
811/// use av_denoise_core::accelerate::Accelerator;
812/// use av_denoise_core::{ChannelMode, Denoiser, DenoiserOptions, DenoisingMode, Device};
813///
814/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
815/// let options = DenoiserOptions::builder()
816///     .channel_mode(ChannelMode::Luma)
817///     .mode(DenoisingMode::Temporal { radius: 2 })
818///     .build();
819///
820/// let mut denoiser = Denoiser::create(
821///     &[Accelerator::Vulkan],
822///     &Device::Default,
823///     1920,
824///     1080,
825///     options,
826/// )?;
827///
828/// let frames: Vec<Vec<f32>> = read_my_frames();
829/// let mut cleaned: Vec<Vec<f32>> = Vec::new();
830///
831/// for frame in &frames {
832///     denoiser.push_frame(frame)?;
833///
834///     // Temporal denoising runs a few frames behind the input, so
835///     // there is not always one ready to collect.
836///     if let Some(out) = denoiser.recv_frame()? {
837///         cleaned.push(out.into_f32().expect("built for f32 output"));
838///     }
839/// }
840///
841/// // Drain the frames still inside the temporal window.
842/// denoiser.flush(|out| cleaned.push(out.into_f32().expect("built for f32 output")))?;
843/// # Ok(())
844/// # }
845/// # fn read_my_frames() -> Vec<Vec<f32>> { Vec::new() }
846/// ```
847pub struct Denoiser {
848    backend: Backend,
849    pending: VecDeque<BackendPending>,
850    accelerator: Accelerator,
851    width: u32,
852    height: u32,
853    temporal_radius: u32,
854    output_format: OutputFormat,
855    frames_pushed: u32,
856    /// Set once any call other than a `QueueFull` push has failed, so how
857    /// many frames are in flight is no longer known.
858    ///
859    /// Every entry point refuses to run while this is set.
860    /// [`Self::reset_stream`] clears it.
861    poisoned: bool,
862}
863
864impl Denoiser {
865    /// Tries each accelerator in `accelerators` in order and builds a
866    /// denoiser on the first one that works.
867    ///
868    /// `device` picks a non-default device on the chosen runtime.
869    ///
870    /// # Thread stack size
871    ///
872    /// cubecl spawns its own per-device worker thread, named
873    /// `DS{U,D}-…`, and runs GPU kernel codegen on it. That thread gets
874    /// Rust's default stack, which is `RUST_MIN_STACK` or 2 MiB when
875    /// that is unset.
876    ///
877    /// The windowed NLM kernels unroll their body
878    /// `(2 * search_radius + 1)^2` times, so a `search_radius` of about
879    /// 5 or more can overflow the 2 MiB default and abort the process.
880    ///
881    /// Callers using a `search_radius` above 4 should call
882    /// [`crate::raise_codegen_stack_limit`] before any cubecl thread
883    /// spawns, usually right at the top of `main`.
884    pub fn create(
885        accelerators: &[Accelerator],
886        device: &Device,
887        width: u32,
888        height: u32,
889        options: DenoiserOptions,
890    ) -> Result<Self, DenoiserError> {
891        let accelerator =
892            sniff_best_accelerator(accelerators, device).ok_or(DenoiserError::NoAcceleratorAvailable)?;
893
894        let params = options.to_nlm_params();
895        params.validate()?;
896        validate_dimensions(width, height)?;
897
898        let temporal_radius = params.temporal_radius;
899        let backend = build_backend(
900            accelerator,
901            device,
902            &options.algorithm,
903            params,
904            width,
905            height,
906            options.output_format,
907        )?;
908
909        Ok(Self {
910            backend,
911            pending: VecDeque::with_capacity(MAX_PENDING),
912            accelerator,
913            width,
914            height,
915            temporal_radius,
916            output_format: options.output_format,
917            frames_pushed: 0,
918            poisoned: false,
919        })
920    }
921
922    /// The accelerator [`sniff_best_accelerator`] picked.
923    pub fn selected_accelerator(&self) -> Accelerator {
924        self.accelerator
925    }
926
927    /// The width passed at construction.
928    pub fn width(&self) -> u32 {
929        self.width
930    }
931
932    /// The height passed at construction.
933    pub fn height(&self) -> u32 {
934        self.height
935    }
936
937    /// The temporal radius the resolved parameters run at.
938    pub fn temporal_radius(&self) -> u32 {
939        self.temporal_radius
940    }
941
942    /// The format every collected frame comes back in.
943    pub fn output_format(&self) -> OutputFormat {
944        self.output_format
945    }
946
947    /// How many frames behind and ahead of a target frame this
948    /// denoiser needs pushed, in order, to produce that frame's
949    /// output through [`PlanarDenoiser::reseed`](crate::PlanarDenoiser::reseed).
950    ///
951    /// Both NLM algorithms only ever need their own `2 * radius + 1`
952    /// sliding window, symmetric around the target frame:
953    /// `WindowSpan { behind: radius, ahead: radius }`.
954    ///
955    /// nl4d's cross-frame accumulator scatters every pass's
956    /// contribution across the `2 * radius + 1` frames the pass
957    /// reaches, and a frame's own region only starts collecting once
958    /// the pass that first reaches it, the one centred `radius` frames
959    /// behind it, has run. That earliest pass is itself only real once
960    /// the front end's own window is full at that centre, which needs
961    /// `radius` more frames behind it again. So nl4d needs the target's
962    /// own `radius`-wide neighbourhood doubled on both sides:
963    /// `WindowSpan { behind: 2 * radius, ahead: 2 * radius }`.
964    pub fn window_span(&self) -> WindowSpan {
965        let radius = self.temporal_radius as usize;
966        let span = if self.backend.is_nl4d() {
967            2 * radius
968        } else {
969            radius
970        };
971        WindowSpan {
972            behind: span,
973            ahead: span,
974        }
975    }
976
977    /// Uploads one frame into the temporal window.
978    ///
979    /// `frame` holds `width * height * channels` `f32` values in
980    /// `[0, 1]`.
981    ///
982    /// Once the window is full and the pipeline has room, this also
983    /// starts the kernels for the next denoised frame.
984    ///
985    /// Up to `MAX_PENDING` outputs can be in flight at once, so the GPU
986    /// runs one frame's kernels while the previous frame's readback is
987    /// still travelling. At that ceiling this returns
988    /// [`DenoiserError::QueueFull`], and the caller has to drain a frame
989    /// with [`Self::recv_frame`] before pushing more.
990    ///
991    /// Any other failure poisons the denoiser, so every further call
992    /// returns [`DenoiserError::Poisoned`] until [`Self::reset_stream`]
993    /// clears it. `QueueFull` does not poison, since it is the documented
994    /// retry signal above.
995    pub fn push_frame(&mut self, frame: &[f32]) -> Result<(), DenoiserError> {
996        if self.poisoned {
997            return Err(DenoiserError::Poisoned);
998        }
999        self.push_frame_inner(frame).inspect_err(|err| {
1000            if !matches!(err, DenoiserError::QueueFull) {
1001                self.poisoned = true;
1002            }
1003        })
1004    }
1005
1006    fn push_frame_inner(&mut self, frame: &[f32]) -> Result<(), DenoiserError> {
1007        // After `temporal_radius` real pushes the leading-edge mirror
1008        // has primed the window, so the next push produces a pending
1009        // frame. From then on every push takes a pending slot.
1010        let window_full = self.frames_pushed > self.temporal_radius;
1011        if window_full && self.pending.len() >= MAX_PENDING {
1012            return Err(DenoiserError::QueueFull);
1013        }
1014
1015        match &mut self.backend {
1016            #[cfg(feature = "cuda")]
1017            Backend::Cuda(d) => {
1018                d.push_frame(frame);
1019                if let Some(p) = d.denoise_submit()? {
1020                    self.pending.push_back(BackendPending::Cuda(p));
1021                }
1022            },
1023            #[cfg(feature = "rocm")]
1024            Backend::Rocm(d) => {
1025                d.push_frame(frame);
1026                if let Some(p) = d.denoise_submit()? {
1027                    self.pending.push_back(BackendPending::Rocm(p));
1028                }
1029            },
1030            #[cfg(any(feature = "vulkan", feature = "metal"))]
1031            Backend::Wgpu(d) => {
1032                d.push_frame(frame);
1033                if let Some(p) = d.denoise_submit()? {
1034                    self.pending.push_back(BackendPending::Wgpu(p));
1035                }
1036            },
1037        }
1038
1039        self.frames_pushed = self.frames_pushed.saturating_add(1);
1040        Ok(())
1041    }
1042
1043    /// Uploads one frame held as wire bytes into the temporal window.
1044    ///
1045    /// `planes` holds one `width * height` plane per channel at `depth`,
1046    /// which the GPU normalises and interleaves. The planes run Y, U, V
1047    /// for a fused frame and U, V for a chroma pair.
1048    ///
1049    /// Queueing, poisoning, and the `QueueFull` retry signal work exactly
1050    /// as they do for [`Self::push_frame`].
1051    pub fn push_frame_wire(&mut self, planes: &[&[u8]], depth: Depth) -> Result<(), DenoiserError> {
1052        if self.poisoned {
1053            return Err(DenoiserError::Poisoned);
1054        }
1055        self.push_frame_wire_inner(planes, depth).inspect_err(|err| {
1056            if !matches!(err, DenoiserError::QueueFull) {
1057                self.poisoned = true;
1058            }
1059        })
1060    }
1061
1062    fn push_frame_wire_inner(&mut self, planes: &[&[u8]], depth: Depth) -> Result<(), DenoiserError> {
1063        // The same window accounting `push_frame_inner` does.
1064        let window_full = self.frames_pushed > self.temporal_radius;
1065        if window_full && self.pending.len() >= MAX_PENDING {
1066            return Err(DenoiserError::QueueFull);
1067        }
1068
1069        match &mut self.backend {
1070            #[cfg(feature = "cuda")]
1071            Backend::Cuda(d) => {
1072                d.push_frame_wire(planes, depth);
1073                if let Some(p) = d.denoise_submit()? {
1074                    self.pending.push_back(BackendPending::Cuda(p));
1075                }
1076            },
1077            #[cfg(feature = "rocm")]
1078            Backend::Rocm(d) => {
1079                d.push_frame_wire(planes, depth);
1080                if let Some(p) = d.denoise_submit()? {
1081                    self.pending.push_back(BackendPending::Rocm(p));
1082                }
1083            },
1084            #[cfg(any(feature = "vulkan", feature = "metal"))]
1085            Backend::Wgpu(d) => {
1086                d.push_frame_wire(planes, depth);
1087                if let Some(p) = d.denoise_submit()? {
1088                    self.pending.push_back(BackendPending::Wgpu(p));
1089                }
1090            },
1091        }
1092
1093        self.frames_pushed = self.frames_pushed.saturating_add(1);
1094        Ok(())
1095    }
1096
1097    /// Uploads one frame held as wire bytes into the temporal window
1098    /// without starting a denoise.
1099    ///
1100    /// The wire counterpart of [`Self::push_frame_priming`].
1101    pub fn push_frame_wire_priming(&mut self, planes: &[&[u8]], depth: Depth) -> Result<(), DenoiserError> {
1102        if self.poisoned {
1103            return Err(DenoiserError::Poisoned);
1104        }
1105        match &mut self.backend {
1106            #[cfg(feature = "cuda")]
1107            Backend::Cuda(d) => d.push_frame_wire(planes, depth),
1108            #[cfg(feature = "rocm")]
1109            Backend::Rocm(d) => d.push_frame_wire(planes, depth),
1110            #[cfg(any(feature = "vulkan", feature = "metal"))]
1111            Backend::Wgpu(d) => d.push_frame_wire(planes, depth),
1112        }
1113
1114        self.frames_pushed = self.frames_pushed.saturating_add(1);
1115        Ok(())
1116    }
1117
1118    /// Uploads one frame into the temporal window without starting a
1119    /// denoise.
1120    ///
1121    /// The ring advances exactly as it does for [`Self::push_frame`], so
1122    /// the window still fills, but no kernels are submitted and no
1123    /// output is queued. This is how a caller that can hand over a whole
1124    /// window at once, rather than a strictly ordered stream, fills the
1125    /// window in one go and lets only the last push in it submit.
1126    ///
1127    /// A failure elsewhere poisons the denoiser, so this refuses to run
1128    /// until [`Self::reset_stream`] clears it.
1129    pub fn push_frame_priming(&mut self, frame: &[f32]) -> Result<(), DenoiserError> {
1130        if self.poisoned {
1131            return Err(DenoiserError::Poisoned);
1132        }
1133        match &mut self.backend {
1134            #[cfg(feature = "cuda")]
1135            Backend::Cuda(d) => d.push_frame(frame),
1136            #[cfg(feature = "rocm")]
1137            Backend::Rocm(d) => d.push_frame(frame),
1138            #[cfg(any(feature = "vulkan", feature = "metal"))]
1139            Backend::Wgpu(d) => d.push_frame(frame),
1140        }
1141
1142        self.frames_pushed = self.frames_pushed.saturating_add(1);
1143        Ok(())
1144    }
1145
1146    /// Drops the current stream and returns to the state a fresh
1147    /// denoiser starts in, keeping every GPU allocation.
1148    ///
1149    /// Anything still in flight is discarded. This also clears the
1150    /// poison an earlier failure left, so it is the recovery path for
1151    /// [`DenoiserError::Poisoned`].
1152    pub fn reset_stream(&mut self) {
1153        self.pending.clear();
1154        self.frames_pushed = 0;
1155        self.poisoned = false;
1156
1157        match &mut self.backend {
1158            #[cfg(feature = "cuda")]
1159            Backend::Cuda(d) => d.reset_stream(),
1160            #[cfg(feature = "rocm")]
1161            Backend::Rocm(d) => d.reset_stream(),
1162            #[cfg(any(feature = "vulkan", feature = "metal"))]
1163            Backend::Wgpu(d) => d.reset_stream(),
1164        }
1165    }
1166
1167    /// Blocks until the in-flight denoise finishes and returns the
1168    /// cleaned frame.
1169    ///
1170    /// Returns `Ok(None)` when nothing is in flight, which happens while
1171    /// the temporal window is still filling up.
1172    ///
1173    /// A failure poisons the denoiser, so every further call returns
1174    /// [`DenoiserError::Poisoned`] until [`Self::reset_stream`] clears it.
1175    pub fn recv_frame(&mut self) -> Result<Option<FrameOutput>, DenoiserError> {
1176        if self.poisoned {
1177            return Err(DenoiserError::Poisoned);
1178        }
1179        self.recv_frame_inner().inspect_err(|_| self.poisoned = true)
1180    }
1181
1182    fn recv_frame_inner(&mut self) -> Result<Option<FrameOutput>, DenoiserError> {
1183        let Some(pending) = self.pending.pop_front() else {
1184            return Ok(None);
1185        };
1186        Ok(Some(pending.wait()?))
1187    }
1188
1189    /// Polls the in-flight denoise once.
1190    ///
1191    /// Returns `Ok(None)` both when nothing is in flight and when the in-flight readback
1192    /// has not landed yet, so `None` alone does not tell those two cases apart.
1193    ///
1194    /// A caller that needs the frame rather than just checking on it should
1195    /// use [`Self::recv_frame`] instead.
1196    ///
1197    /// This only avoids blocking on the wgpu backends, meaning Vulkan and Metal.
1198    /// On CUDA and ROCm the readback completes synchronously on its first poll,
1199    /// so this call blocks until the readback lands there, the same as `recv_frame`.
1200    ///
1201    /// A failure poisons the denoiser, so every further call returns
1202    /// [`DenoiserError::Poisoned`] until [`Self::reset_stream`] clears it.
1203    pub fn try_recv_frame(&mut self) -> Result<Option<FrameOutput>, DenoiserError> {
1204        if self.poisoned {
1205            return Err(DenoiserError::Poisoned);
1206        }
1207        self.try_recv_frame_inner().inspect_err(|_| self.poisoned = true)
1208    }
1209
1210    fn try_recv_frame_inner(&mut self) -> Result<Option<FrameOutput>, DenoiserError> {
1211        let Some(pending) = self.pending.pop_front() else {
1212            return Ok(None);
1213        };
1214
1215        match pending.try_wait()? {
1216            Ok(frame) => Ok(Some(frame)),
1217            Err(pending) => {
1218                self.pending.push_front(pending);
1219                Ok(None)
1220            },
1221        }
1222    }
1223
1224    /// Drains the in-flight frames and the trailing temporal tail,
1225    /// handing each frame it produces to `sink`.
1226    ///
1227    /// The tail is padded by repeating the last pushed frame.
1228    ///
1229    /// On success the denoiser is ready for a fresh, unrelated stream of
1230    /// the same size and parameters. Pushing again after a flush starts
1231    /// a new temporal window from scratch, and flushing more than once
1232    /// is fine.
1233    ///
1234    /// A failure poisons the denoiser, so every further call returns
1235    /// [`DenoiserError::Poisoned`] until [`Self::reset_stream`] clears it.
1236    pub fn flush(&mut self, sink: impl FnMut(FrameOutput)) -> Result<(), DenoiserError> {
1237        if self.poisoned {
1238            return Err(DenoiserError::Poisoned);
1239        }
1240        self.flush_inner(sink).inspect_err(|_| self.poisoned = true)
1241    }
1242
1243    fn flush_inner(&mut self, mut sink: impl FnMut(FrameOutput)) -> Result<(), DenoiserError> {
1244        // Drain the whole pending pipeline, up to MAX_PENDING frames,
1245        // before submitting the trailing-tail mirrors. This also leaves
1246        // every output slot free, so the tail's own readbacks cannot be
1247        // handed a slot a streaming readback is still reading.
1248        while let Some(frame) = self.recv_frame_inner()? {
1249            sink(frame);
1250        }
1251
1252        // The tail frames come back through each algorithm's own
1253        // blocking readback, in the same format as every streaming
1254        // frame, so they are quantised by the same pack kernel.
1255        match &mut self.backend {
1256            #[cfg(feature = "cuda")]
1257            Backend::Cuda(d) => d.flush(|frame| sink(frame.clone()))?,
1258            #[cfg(feature = "rocm")]
1259            Backend::Rocm(d) => d.flush(|frame| sink(frame.clone()))?,
1260            #[cfg(any(feature = "vulkan", feature = "metal"))]
1261            Backend::Wgpu(d) => d.flush(|frame| sink(frame.clone()))?,
1262        }
1263
1264        // The backend has already reset its own stream indices. Reset
1265        // the outer push counter too, so the next push re-arms the
1266        // window-priming check at the top of `push_frame`.
1267        self.frames_pushed = 0;
1268
1269        Ok(())
1270    }
1271
1272    /// Sets the poison flag directly, without going through a failing
1273    /// call, so a test can check what a caller sees once the flag is
1274    /// already set and what recovers it.
1275    #[cfg(test)]
1276    pub(crate) fn poison_for_test(&mut self) {
1277        self.poisoned = true;
1278    }
1279
1280    /// The backend's packed-word output buffers, which only exist in
1281    /// wire mode.
1282    #[cfg(test)]
1283    pub(crate) fn wire_outputs_for_test(&self) -> Option<&[cubecl::server::Handle; 2]> {
1284        self.backend.wire_outputs()
1285    }
1286}
1287
1288fn build_backend(
1289    accel: Accelerator,
1290    device: &Device,
1291    algorithm: &Algorithm,
1292    params: NlmParams,
1293    width: u32,
1294    height: u32,
1295    output_format: OutputFormat,
1296) -> Result<Backend, DenoiserError> {
1297    match accel {
1298        #[cfg(feature = "cuda")]
1299        Accelerator::Cuda => {
1300            let dev = device.to_cuda()?;
1301            let client = <cubecl::cuda::CudaRuntime as Runtime>::client(&dev);
1302            Ok(Backend::Cuda(build_engine(
1303                &client,
1304                algorithm,
1305                params,
1306                width,
1307                height,
1308                output_format,
1309            )?))
1310        },
1311        #[cfg(feature = "rocm")]
1312        Accelerator::Rocm => {
1313            let dev = device.to_amd()?;
1314            let client = <cubecl::hip::HipRuntime as Runtime>::client(&dev);
1315            Ok(Backend::Rocm(build_engine(
1316                &client,
1317                algorithm,
1318                params,
1319                width,
1320                height,
1321                output_format,
1322            )?))
1323        },
1324        #[cfg(feature = "vulkan")]
1325        Accelerator::Vulkan => {
1326            let dev = device.to_wgpu()?;
1327            let client = <cubecl::wgpu::WgpuRuntime as Runtime>::client(&dev);
1328            Ok(Backend::Wgpu(build_engine(
1329                &client,
1330                algorithm,
1331                params,
1332                width,
1333                height,
1334                output_format,
1335            )?))
1336        },
1337        #[cfg(feature = "metal")]
1338        Accelerator::Metal => {
1339            let dev = device.to_wgpu()?;
1340            let client = <cubecl::wgpu::WgpuRuntime as Runtime>::client(&dev);
1341            Ok(Backend::Wgpu(build_engine(
1342                &client,
1343                algorithm,
1344                params,
1345                width,
1346                height,
1347                output_format,
1348            )?))
1349        },
1350        // Keeps the match exhaustive on docs.rs, where `cfg(docsrs)`
1351        // widens the `Accelerator` enum to include variants whose
1352        // backend feature is not enabled. Never reached at runtime.
1353        #[cfg(docsrs)]
1354        #[expect(
1355            unreachable_patterns,
1356            reason = "the arm only keeps the match exhaustive on docs.rs"
1357        )]
1358        _ => unreachable!(),
1359    }
1360}
1361
1362#[cfg(test)]
1363mod options_tests {
1364    use super::*;
1365
1366    /// `Algorithm::NlmeansHq` with `hq` overridden and everything else
1367    /// left at its default.
1368    fn hq(hq: HqParams) -> Algorithm {
1369        Algorithm::NlmeansHq(NlmeansHqOptions {
1370            hq,
1371            ..NlmeansHqOptions::default()
1372        })
1373    }
1374
1375    /// `Algorithm::Nlmeans` with `tuning` overridden.
1376    fn fast_tuned(tuning: NlmTuning) -> Algorithm {
1377        Algorithm::Nlmeans(NlmeansOptions {
1378            tuning,
1379            ..NlmeansOptions::default()
1380        })
1381    }
1382
1383    #[test]
1384    fn nl4d_default_lambda_ht_differs_between_luma_and_chroma() {
1385        let luma = nl4d_default_lambda_ht(ChannelMode::Luma);
1386        let chroma = nl4d_default_lambda_ht(ChannelMode::Chroma);
1387
1388        assert!((luma - 5.3).abs() < f32::EPSILON);
1389        assert!((chroma - 4.2).abs() < f32::EPSILON);
1390        assert!(
1391            (chroma - luma).abs() > f32::EPSILON,
1392            "the two planes should not resolve to the same default"
1393        );
1394    }
1395
1396    #[test]
1397    fn nl4d_default_lambda_ht_yuv_reads_the_luma_value() {
1398        let yuv = nl4d_default_lambda_ht(ChannelMode::Yuv);
1399        let luma = nl4d_default_lambda_ht(ChannelMode::Luma);
1400
1401        assert!((yuv - luma).abs() < f32::EPSILON);
1402    }
1403
1404    #[test]
1405    fn resolve_lambda_ht_unset_uses_the_per_plane_default() {
1406        let opts = Nl4dOptions::default();
1407
1408        let luma = resolve_lambda_ht(&opts, ChannelMode::Luma).expect("the default scale is in range");
1409        let chroma = resolve_lambda_ht(&opts, ChannelMode::Chroma).expect("the default scale is in range");
1410
1411        assert!((luma - 5.3).abs() < f32::EPSILON, "got {luma}");
1412        assert!((chroma - 4.2).abs() < f32::EPSILON, "got {chroma}");
1413    }
1414
1415    #[test]
1416    fn resolve_lambda_ht_explicit_value_overrides_every_plane() {
1417        let opts = Nl4dOptions {
1418            lambda_ht: Some(4.4),
1419            ..Nl4dOptions::default()
1420        };
1421
1422        for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1423            let got = resolve_lambda_ht(&opts, channels).expect("the default scale is in range");
1424            assert!(
1425                (got - 4.4).abs() < f32::EPSILON,
1426                "channels {channels:?} got {got}"
1427            );
1428        }
1429    }
1430
1431    #[test]
1432    fn resolve_lambda_ht_default_scale_leaves_the_value_alone() {
1433        let opts = Nl4dOptions::default();
1434
1435        for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1436            let got = resolve_lambda_ht(&opts, channels).expect("the default scale is in range");
1437            let want = nl4d_default_lambda_ht(channels);
1438            assert!(
1439                (got - want).abs() < f32::EPSILON,
1440                "channels {channels:?} got {got}"
1441            );
1442        }
1443    }
1444
1445    #[test]
1446    fn resolve_lambda_ht_scale_multiplies_the_per_plane_default() {
1447        let opts = Nl4dOptions {
1448            lambda_ht_scale: 1.1,
1449            ..Nl4dOptions::default()
1450        };
1451
1452        for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1453            let got = resolve_lambda_ht(&opts, channels).expect("1.1 is in range");
1454            let want = nl4d_default_lambda_ht(channels) * 1.1;
1455            assert!(
1456                (got - want).abs() < 1e-5,
1457                "channels {channels:?} got {got}, want {want}"
1458            );
1459        }
1460    }
1461
1462    /// The scale is not limited to the defaults. Pinning one plane and
1463    /// scaling both is the combination this exists for.
1464    #[test]
1465    fn resolve_lambda_ht_scale_multiplies_an_explicit_value() {
1466        let opts = Nl4dOptions {
1467            lambda_ht: Some(5.0),
1468            lambda_ht_scale: 0.9,
1469            ..Nl4dOptions::default()
1470        };
1471
1472        let got = resolve_lambda_ht(&opts, ChannelMode::Luma).expect("0.9 is in range");
1473        assert!((got - 4.5).abs() < 1e-5, "got {got}");
1474    }
1475
1476    #[test]
1477    fn resolve_lambda_ht_rejects_an_out_of_range_scale() {
1478        for bad in [0.0, -1.0, 0.05, 10.5, f32::NAN, f32::INFINITY] {
1479            let opts = Nl4dOptions {
1480                lambda_ht_scale: bad,
1481                ..Nl4dOptions::default()
1482            };
1483            let err = resolve_lambda_ht(&opts, ChannelMode::Luma).unwrap_err();
1484            assert!(
1485                err.contains("lambda_ht_scale"),
1486                "lambda_ht_scale={bad} should be rejected, got {err}"
1487            );
1488        }
1489    }
1490
1491    #[test]
1492    fn the_default_algorithm_is_the_fast_nlmeans_path() {
1493        let opts = DenoiserOptions::builder().build();
1494        assert_eq!(opts.algorithm, Algorithm::Nlmeans(NlmeansOptions::default()));
1495    }
1496
1497    #[test]
1498    fn spatial_mode_maps_to_zero_temporal_radius() {
1499        let opts = DenoiserOptions::builder()
1500            .channel_mode(ChannelMode::Yuv)
1501            .mode(DenoisingMode::Spacial)
1502            .build();
1503        let params = opts.to_nlm_params();
1504
1505        assert_eq!(params.temporal_radius, 0);
1506        assert_eq!(params.channels, ChannelMode::Yuv);
1507    }
1508
1509    #[test]
1510    fn temporal_mode_propagates_radius() {
1511        let opts = DenoiserOptions::builder()
1512            .mode(DenoisingMode::Temporal { radius: 3 })
1513            .build();
1514        let params = opts.to_nlm_params();
1515
1516        assert_eq!(params.temporal_radius, 3);
1517    }
1518
1519    #[test]
1520    fn prefilter_passthrough() {
1521        let opts = DenoiserOptions::builder()
1522            .algorithm(Algorithm::Nlmeans(NlmeansOptions {
1523                prefilter: PrefilterMode::Bilateral {
1524                    sigma_s: 3.0,
1525                    sigma_r: 0.02,
1526                },
1527                ..NlmeansOptions::default()
1528            }))
1529            .build();
1530        let params = opts.to_nlm_params();
1531
1532        assert!(matches!(params.prefilter, PrefilterMode::Bilateral { .. }));
1533    }
1534
1535    #[test]
1536    fn hq_unset_prefilter_defaults_to_none() {
1537        let opts = DenoiserOptions::builder()
1538            .algorithm(hq(HqParams::default()))
1539            .build();
1540        let params = opts.to_nlm_params();
1541
1542        assert!(matches!(params.prefilter, PrefilterMode::None));
1543    }
1544
1545    #[test]
1546    fn fast_unset_prefilter_defaults_to_none() {
1547        let opts = DenoiserOptions::builder()
1548            .algorithm(Algorithm::Nlmeans(NlmeansOptions::default()))
1549            .build();
1550        let params = opts.to_nlm_params();
1551
1552        assert!(matches!(params.prefilter, PrefilterMode::None));
1553    }
1554
1555    #[test]
1556    fn hq_unset_strength_defaults_to_hq_default_strength() {
1557        // Default channel_mode is Yuv, default mode is Spacial (radius 0).
1558        let opts = DenoiserOptions::builder()
1559            .algorithm(hq(HqParams::default()))
1560            .build();
1561        let params = opts.to_nlm_params();
1562
1563        let expected = hq_default_strength(ChannelMode::Yuv, 0);
1564        assert!((params.strength - expected).abs() < f32::EPSILON);
1565    }
1566
1567    #[test]
1568    fn hq_no_auto_strength_falls_back_to_the_legacy_absolute_default() {
1569        // `effective_strength_with` only reads `strength` as a
1570        // multiplier on the measured sigma when `auto_strength` is true.
1571        // With it false, `strength` is an FFmpeg-style absolute value,
1572        // so the fallback has to be the fast path's absolute default
1573        // rather than a calibrated multiplier from
1574        // `hq_default_strength`.
1575        let opts = DenoiserOptions::builder()
1576            .algorithm(hq(HqParams {
1577                auto_strength: false,
1578                ..HqParams::default()
1579            }))
1580            .build();
1581        let params = opts.to_nlm_params();
1582
1583        let expected = NlmParams::default().strength;
1584        assert!(
1585            (params.strength - expected).abs() < f32::EPSILON,
1586            "expected the legacy absolute default {expected}, got {}, which looks like the \
1587             auto-strength multiplier table leaking through",
1588            params.strength
1589        );
1590    }
1591
1592    #[test]
1593    fn hq_luma_r4_uses_measured_table_value() {
1594        let opts = DenoiserOptions::builder()
1595            .channel_mode(ChannelMode::Luma)
1596            .mode(DenoisingMode::Temporal { radius: 4 })
1597            .algorithm(hq(HqParams::default()))
1598            .build();
1599        let params = opts.to_nlm_params();
1600
1601        assert!((params.strength - 0.35).abs() < f32::EPSILON);
1602    }
1603
1604    #[test]
1605    fn hq_chroma_r4_uses_measured_table_value() {
1606        let opts = DenoiserOptions::builder()
1607            .channel_mode(ChannelMode::Chroma)
1608            .mode(DenoisingMode::Temporal { radius: 4 })
1609            .algorithm(hq(HqParams::default()))
1610            .build();
1611        let params = opts.to_nlm_params();
1612
1613        assert!((params.strength - 0.70).abs() < f32::EPSILON);
1614    }
1615
1616    #[test]
1617    fn hq_yuv_r8_uses_measured_table_value() {
1618        let opts = DenoiserOptions::builder()
1619            .channel_mode(ChannelMode::Yuv)
1620            .mode(DenoisingMode::Temporal { radius: 8 })
1621            .algorithm(hq(HqParams::default()))
1622            .build();
1623        let params = opts.to_nlm_params();
1624
1625        assert!((params.strength - 0.30).abs() < f32::EPSILON);
1626    }
1627
1628    #[test]
1629    fn hq_spacial_mode_uses_radius_zero_table_values() {
1630        for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1631            let opts = DenoiserOptions::builder()
1632                .channel_mode(channels)
1633                .mode(DenoisingMode::Spacial)
1634                .algorithm(hq(HqParams::default()))
1635                .build();
1636            let params = opts.to_nlm_params();
1637
1638            let expected = hq_default_strength(channels, 0);
1639            assert!(
1640                (params.strength - expected).abs() < f32::EPSILON,
1641                "for channels {channels:?} expected {expected}, got {}",
1642                params.strength
1643            );
1644        }
1645    }
1646
1647    #[test]
1648    fn hq_explicit_strength_wins_over_the_table_for_every_plane() {
1649        for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1650            let opts = DenoiserOptions::builder()
1651                .channel_mode(channels)
1652                .mode(DenoisingMode::Temporal { radius: 4 })
1653                .algorithm(Algorithm::NlmeansHq(NlmeansHqOptions {
1654                    nlm: NlmeansOptions {
1655                        tuning: NlmTuning {
1656                            strength: Some(0.99),
1657                            ..NlmTuning::default()
1658                        },
1659                        ..NlmeansOptions::default()
1660                    },
1661                    hq: HqParams::default(),
1662                }))
1663                .build();
1664            let params = opts.to_nlm_params();
1665
1666            assert!(
1667                (params.strength - 0.99).abs() < f32::EPSILON,
1668                "for channels {channels:?} the explicit strength was overridden by the table"
1669            );
1670        }
1671    }
1672
1673    #[test]
1674    fn fast_unset_strength_defaults_to_legacy_default() {
1675        let opts = DenoiserOptions::builder()
1676            .algorithm(Algorithm::Nlmeans(NlmeansOptions::default()))
1677            .build();
1678        let params = opts.to_nlm_params();
1679
1680        assert!((params.strength - 1.2).abs() < f32::EPSILON);
1681    }
1682
1683    #[test]
1684    fn nl4d_options_default_matches_nl4d_params_default() {
1685        let opts = Nl4dOptions::default();
1686        let params = crate::nl4d::Nl4dParams::default();
1687
1688        assert_eq!(opts.refine, params.refine);
1689        assert_eq!(opts.spatial_radius, params.spatial_radius);
1690        assert!((opts.c_min - params.c_min).abs() < f32::EPSILON);
1691        assert_eq!(opts.confidence_variance, params.confidence_variance);
1692        // The two `lambda_ht` fields hold different things, so they are
1693        // not compared. `opts.lambda_ht` stays `None` and is deferred to
1694        // `nl4d_default_lambda_ht` once the plane is known (see
1695        // `resolve_lambda_ht_unset_uses_the_per_plane_default` above),
1696        // while `params.lambda_ht` is a concrete default mirroring the
1697        // Luma/Yuv value.
1698        assert_eq!(opts.lambda_ht, None);
1699        assert!((params.lambda_ht - nl4d_default_lambda_ht(ChannelMode::Yuv)).abs() < f32::EPSILON);
1700    }
1701
1702    /// nl4d takes its own noise and confidence knobs rather than a whole
1703    /// [`HqParams`], so the three it does take have to reach the front
1704    /// end and the rest have to arrive at their defaults.
1705    #[test]
1706    fn nl4d_builds_the_front_ends_hq_params_from_its_own_fields() {
1707        let opts = DenoiserOptions::builder()
1708            .mode(DenoisingMode::Temporal { radius: 2 })
1709            .algorithm(Algorithm::Nl4d(Nl4dOptions {
1710                sigma: Some(0.02),
1711                sigma_scale: 1.3,
1712                thsad_scale: 0.8,
1713                ..Nl4dOptions::default()
1714            }))
1715            .build();
1716        let params = opts.to_nlm_params();
1717
1718        let hq = params.hq.expect("nl4d always runs the hq front end");
1719        assert_eq!(hq.sigma_override, Some(0.02));
1720        assert!((hq.sigma_scale - 1.3).abs() < f32::EPSILON);
1721        assert!((hq.thsad_scale - 0.8).abs() < f32::EPSILON);
1722        assert!(
1723            hq.temporal_confidence,
1724            "the grouping kernel reads the confidence scores, so this cannot be off"
1725        );
1726    }
1727
1728    /// The temporal radius has one source now, `mode`, so nl4d cannot
1729    /// disagree with the front end's ring about how wide the window is.
1730    #[test]
1731    fn nl4d_reads_its_temporal_radius_from_the_denoising_mode() {
1732        for radius in [1u32, 4, 8] {
1733            let opts = DenoiserOptions::builder()
1734                .mode(DenoisingMode::Temporal { radius })
1735                .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1736                .build();
1737
1738            assert_eq!(opts.to_nlm_params().temporal_radius, radius);
1739        }
1740    }
1741
1742    /// nl4d never runs an NLM weighting pass, so a prefilter would cost
1743    /// a GPU pass per frame producing a reference image nothing reads.
1744    #[test]
1745    fn nl4d_never_builds_a_prefilter() {
1746        let opts = DenoiserOptions::builder()
1747            .mode(DenoisingMode::Temporal { radius: 2 })
1748            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1749            .build();
1750
1751        assert!(matches!(opts.to_nlm_params().prefilter, PrefilterMode::None));
1752    }
1753
1754    /// Nothing in the nl4d path reads `strength`, so it stays at the
1755    /// library default rather than picking up HQ's calibrated table.
1756    #[test]
1757    fn nl4d_leaves_the_nlm_weighting_knobs_at_their_defaults() {
1758        let defaults = NlmParams::default();
1759        let opts = DenoiserOptions::builder()
1760            .channel_mode(ChannelMode::Luma)
1761            .mode(DenoisingMode::Temporal { radius: 4 })
1762            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1763            .build();
1764        let params = opts.to_nlm_params();
1765
1766        assert!((params.strength - defaults.strength).abs() < f32::EPSILON);
1767        assert_eq!(params.search_radius, defaults.search_radius);
1768        assert_eq!(params.patch_radius, defaults.patch_radius);
1769        assert!((params.self_weight - defaults.self_weight).abs() < f32::EPSILON);
1770    }
1771
1772    /// nl4d always tracks motion, so its `MotionSearch` reaches the
1773    /// front end as an active `Mvtools` mode.
1774    #[test]
1775    fn nl4d_motion_search_becomes_an_active_mvtools_mode() {
1776        let opts = DenoiserOptions::builder()
1777            .mode(DenoisingMode::Temporal { radius: 2 })
1778            .algorithm(Algorithm::Nl4d(Nl4dOptions {
1779                motion: MotionSearch {
1780                    blksize: 32,
1781                    overlap: 16,
1782                    search_radius: 6,
1783                    pyramid_levels: 1,
1784                    estimation: MotionEstimation::Direct,
1785                },
1786                ..Nl4dOptions::default()
1787            }))
1788            .build();
1789        let params = opts.to_nlm_params();
1790
1791        assert!(matches!(
1792            params.motion_compensation,
1793            MotionCompensationMode::Mvtools {
1794                blksize: 32,
1795                overlap: 16,
1796                search_radius: 6,
1797                pyramid_levels: 1,
1798                estimation: MotionEstimation::Direct,
1799            }
1800        ));
1801    }
1802
1803    #[test]
1804    fn nl4d_motion_search_defaults_match_the_front_ends_own_defaults() {
1805        let opts = DenoiserOptions::builder()
1806            .mode(DenoisingMode::Temporal { radius: 2 })
1807            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1808            .build();
1809        let params = opts.to_nlm_params();
1810
1811        assert_eq!(
1812            params.motion_compensation,
1813            crate::nl4d::Nl4dParams::default().nlm.motion_compensation
1814        );
1815    }
1816
1817    #[test]
1818    fn motion_compensation_passthrough() {
1819        let opts = DenoiserOptions::builder()
1820            .mode(DenoisingMode::Temporal { radius: 1 })
1821            .algorithm(Algorithm::Nlmeans(NlmeansOptions {
1822                motion_compensation: MotionCompensationMode::Mvtools {
1823                    blksize: 16,
1824                    overlap: 8,
1825                    search_radius: 4,
1826                    pyramid_levels: 2,
1827                    estimation: MotionEstimation::Direct,
1828                },
1829                ..NlmeansOptions::default()
1830            }))
1831            .build();
1832        let params = opts.to_nlm_params();
1833
1834        assert!(matches!(
1835            params.motion_compensation,
1836            MotionCompensationMode::Mvtools {
1837                blksize: 16,
1838                overlap: 8,
1839                search_radius: 4,
1840                pyramid_levels: 2,
1841                ..
1842            }
1843        ));
1844    }
1845
1846    #[test]
1847    fn motion_compensation_defaults_to_none() {
1848        let opts = DenoiserOptions::builder().build();
1849        let params = opts.to_nlm_params();
1850        assert!(matches!(params.motion_compensation, MotionCompensationMode::None));
1851    }
1852
1853    #[test]
1854    fn nlm_tuning_overrides_individual_fields() {
1855        let defaults = NlmParams::default();
1856        let opts = DenoiserOptions::builder()
1857            .algorithm(fast_tuned(NlmTuning {
1858                search_radius: Some(7),
1859                patch_radius: None,
1860                strength: Some(2.5),
1861                self_weight: None,
1862            }))
1863            .build();
1864        let params = opts.to_nlm_params();
1865
1866        assert_eq!(params.search_radius, 7);
1867        assert_eq!(params.patch_radius, defaults.patch_radius);
1868        assert!((params.strength - 2.5).abs() < f32::EPSILON);
1869        assert!((params.self_weight - defaults.self_weight).abs() < f32::EPSILON);
1870    }
1871}
1872
1873#[cfg(all(test, feature = "vulkan"))]
1874mod tests {
1875    use super::*;
1876
1877    fn opts(mode: DenoisingMode) -> DenoiserOptions {
1878        DenoiserOptions::builder()
1879            .channel_mode(ChannelMode::Luma)
1880            .mode(mode)
1881            .build()
1882    }
1883
1884    fn frame(w: u32, h: u32) -> Vec<f32> {
1885        vec![0.5f32; (w * h) as usize]
1886    }
1887
1888    fn f32_out(out: FrameOutput) -> Vec<f32> {
1889        out.into_f32().expect("f32 output")
1890    }
1891
1892    #[test]
1893    fn spatial_denoise_roundtrip() {
1894        let mut d = Denoiser::create(
1895            &[Accelerator::Vulkan],
1896            &Device::Default,
1897            16,
1898            16,
1899            opts(DenoisingMode::Spacial),
1900        )
1901        .expect("denoiser construction failed");
1902        assert_eq!(d.selected_accelerator(), Accelerator::Vulkan);
1903
1904        d.push_frame(&frame(16, 16)).expect("push failed");
1905        let out = f32_out(d.recv_frame().expect("recv failed").expect("no frame"));
1906        assert_eq!(out.len(), 16 * 16);
1907    }
1908
1909    #[test]
1910    fn nl4d_algorithm_round_trips_through_the_facade() {
1911        let opts = DenoiserOptions::builder()
1912            .channel_mode(ChannelMode::Luma)
1913            .mode(DenoisingMode::Temporal { radius: 2 })
1914            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1915            .build();
1916        let mut d = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts)
1917            .expect("nl4d denoiser construction failed");
1918        assert_eq!(d.selected_accelerator(), Accelerator::Vulkan);
1919
1920        // temporal_radius is 2, so a single push does not fill the
1921        // window yet, the same convention every temporal algorithm here
1922        // follows.
1923        d.push_frame(&frame(16, 16)).expect("push failed");
1924        assert!(d.recv_frame().expect("recv failed").is_none());
1925
1926        let mut out = Vec::new();
1927        d.flush(|f| out.push(f32_out(f))).expect("flush failed");
1928        assert_eq!(out.len(), 1, "expected exactly one output for one pushed frame");
1929        assert_eq!(out[0].len(), 16 * 16);
1930    }
1931
1932    /// nl4d groups patches across neighbouring frames, so a spatial
1933    /// mode leaves it nothing to do. The temporal radius has one source
1934    /// now, `mode`, so this is the only way to ask for that.
1935    #[test]
1936    fn nl4d_rejects_a_spatial_denoising_mode() {
1937        let opts = DenoiserOptions::builder()
1938            .channel_mode(ChannelMode::Luma)
1939            .mode(DenoisingMode::Spacial)
1940            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1941            .build();
1942        let result = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts);
1943
1944        match result {
1945            Err(DenoiserError::Other(e)) => assert!(
1946                e.to_string().contains("temporal window"),
1947                "unexpected error message: {e}"
1948            ),
1949            Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
1950            Ok(_) => panic!("expected a rejection, got Ok"),
1951        }
1952    }
1953
1954    /// Both NLM algorithms only need a symmetric `2r+1` window, so
1955    /// `window_span` must report the same radius on both sides.
1956    #[test]
1957    fn window_span_is_symmetric_for_nlmeans() {
1958        let opts = DenoiserOptions::builder()
1959            .channel_mode(ChannelMode::Luma)
1960            .mode(DenoisingMode::Temporal { radius: 3 })
1961            .algorithm(Algorithm::Nlmeans(NlmeansOptions::default()))
1962            .build();
1963        let d = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts)
1964            .expect("denoiser construction failed");
1965
1966        let span = d.window_span();
1967        assert_eq!(span.behind, 3, "behind should equal the temporal radius");
1968        assert_eq!(span.ahead, 3, "ahead should equal the temporal radius");
1969    }
1970
1971    /// nl4d's cross-frame accumulator needs the target's own `radius`
1972    /// neighbourhood doubled on both sides, so both `behind` and
1973    /// `ahead` must come out to `2 * radius`.
1974    #[test]
1975    fn window_span_is_doubled_on_both_sides_for_nl4d() {
1976        let opts = DenoiserOptions::builder()
1977            .channel_mode(ChannelMode::Luma)
1978            .mode(DenoisingMode::Temporal { radius: 3 })
1979            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1980            .build();
1981        let d = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts)
1982            .expect("nl4d denoiser construction failed");
1983
1984        let span = d.window_span();
1985        assert_eq!(span.behind, 6, "behind should equal 2 * the temporal radius");
1986        assert_eq!(span.ahead, 6, "ahead should equal 2 * the temporal radius");
1987    }
1988
1989    #[test]
1990    fn invalid_params_surface_as_error() {
1991        let bad = DenoiserOptions::builder()
1992            .algorithm(Algorithm::Nlmeans(NlmeansOptions {
1993                tuning: NlmTuning {
1994                    strength: Some(0.0),
1995                    ..NlmTuning::default()
1996                },
1997                ..NlmeansOptions::default()
1998            }))
1999            .build();
2000        let result = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, bad);
2001
2002        match result {
2003            Err(DenoiserError::Other(_)) => {},
2004            Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
2005            Ok(_) => panic!("expected validation error, got Ok"),
2006        }
2007    }
2008
2009    #[test]
2010    fn tiny_frame_dimensions_surface_as_error() {
2011        let result = Denoiser::create(
2012            &[Accelerator::Vulkan],
2013            &Device::Default,
2014            2,
2015            2,
2016            opts(DenoisingMode::Spacial),
2017        );
2018
2019        match result {
2020            Err(DenoiserError::Other(e)) => {
2021                assert!(
2022                    e.to_string().contains("supported minimum"),
2023                    "unexpected error message: {e}"
2024                );
2025            },
2026            Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
2027            Ok(_) => panic!("expected dimension validation error, got Ok"),
2028        }
2029    }
2030
2031    #[test]
2032    fn push_after_pending_returns_queue_full() {
2033        let mut d = Denoiser::create(
2034            &[Accelerator::Vulkan],
2035            &Device::Default,
2036            16,
2037            16,
2038            opts(DenoisingMode::Spacial),
2039        )
2040        .unwrap();
2041
2042        // The pipeline is two deep because the output handles are
2043        // double-buffered, so the first two pushes both submit. The
2044        // third would overwrite the oldest pending frame's output slot,
2045        // so it is rejected with QueueFull.
2046        d.push_frame(&frame(16, 16)).unwrap();
2047        d.push_frame(&frame(16, 16)).unwrap();
2048        let err = d.push_frame(&frame(16, 16)).expect_err("expected QueueFull");
2049        assert!(matches!(err, DenoiserError::QueueFull));
2050
2051        let out = f32_out(d.recv_frame().unwrap().unwrap());
2052        assert_eq!(out.len(), 16 * 16);
2053
2054        // After draining one slot the next push must succeed.
2055        d.push_frame(&frame(16, 16)).expect("push after drain failed");
2056    }
2057
2058    /// `QueueFull` is the documented retry signal, so it must not leave
2059    /// the denoiser poisoned.
2060    #[test]
2061    fn queue_full_does_not_poison() {
2062        let mut d = Denoiser::create(
2063            &[Accelerator::Vulkan],
2064            &Device::Default,
2065            16,
2066            16,
2067            opts(DenoisingMode::Spacial),
2068        )
2069        .unwrap();
2070
2071        d.push_frame(&frame(16, 16)).unwrap();
2072        d.push_frame(&frame(16, 16)).unwrap();
2073        let err = d.push_frame(&frame(16, 16)).expect_err("expected QueueFull");
2074        assert!(matches!(err, DenoiserError::QueueFull));
2075        assert!(!d.poisoned, "QueueFull must not poison the denoiser");
2076
2077        d.recv_frame().unwrap().expect("recv failed after QueueFull");
2078
2079        // The queue has room again, so the next push must succeed rather
2080        // than being rejected as poisoned.
2081        d.push_frame(&frame(16, 16))
2082            .expect("push after QueueFull drain should succeed, not poison");
2083    }
2084
2085    #[test]
2086    fn poisoned_denoiser_refuses_every_entry_point() {
2087        let mut d = Denoiser::create(
2088            &[Accelerator::Vulkan],
2089            &Device::Default,
2090            16,
2091            16,
2092            opts(DenoisingMode::Spacial),
2093        )
2094        .unwrap();
2095        d.poisoned = true;
2096
2097        assert!(matches!(
2098            d.push_frame(&frame(16, 16)),
2099            Err(DenoiserError::Poisoned)
2100        ));
2101        assert!(matches!(d.recv_frame(), Err(DenoiserError::Poisoned)));
2102        assert!(matches!(d.try_recv_frame(), Err(DenoiserError::Poisoned)));
2103        assert!(matches!(d.flush(|_| {}), Err(DenoiserError::Poisoned)));
2104    }
2105
2106    #[test]
2107    fn reset_stream_clears_poison() {
2108        let mut d = Denoiser::create(
2109            &[Accelerator::Vulkan],
2110            &Device::Default,
2111            16,
2112            16,
2113            opts(DenoisingMode::Spacial),
2114        )
2115        .unwrap();
2116        d.poisoned = true;
2117
2118        d.reset_stream();
2119        assert!(!d.poisoned, "reset_stream must clear the poison flag");
2120
2121        d.push_frame(&frame(16, 16))
2122            .expect("push after reset_stream should succeed");
2123    }
2124
2125    fn frame_filled(w: u32, h: u32, value: f32) -> Vec<f32> {
2126        vec![value; (w * h) as usize]
2127    }
2128
2129    /// Pushes `n` frames of the given value, receiving along the way to
2130    /// keep the in-flight pipeline below `MAX_PENDING`.
2131    fn push_n_with_drain(d: &mut Denoiser, n: usize, value: f32, out: &mut Vec<Vec<f32>>) {
2132        for _ in 0..n {
2133            loop {
2134                match d.push_frame(&frame_filled(16, 16, value)) {
2135                    Ok(()) => break,
2136                    Err(DenoiserError::QueueFull) => {
2137                        let f = d
2138                            .recv_frame()
2139                            .expect("recv ok")
2140                            .expect("queue full but recv yielded none");
2141                        out.push(f32_out(f));
2142                    },
2143                    Err(e) => panic!("unexpected push error: {e:?}"),
2144                }
2145            }
2146        }
2147    }
2148
2149    #[test]
2150    fn flush_leaves_denoiser_reusable_spatial() {
2151        let mut d = Denoiser::create(
2152            &[Accelerator::Vulkan],
2153            &Device::Default,
2154            16,
2155            16,
2156            opts(DenoisingMode::Spacial),
2157        )
2158        .unwrap();
2159
2160        let mut batch_a = Vec::new();
2161        push_n_with_drain(&mut d, 5, 0.25, &mut batch_a);
2162        d.flush(|f| batch_a.push(f32_out(f))).expect("first flush failed");
2163        assert_eq!(batch_a.len(), 5);
2164
2165        // After flush the pipeline must be empty.
2166        assert!(d.recv_frame().unwrap().is_none());
2167
2168        let mut batch_b = Vec::new();
2169        push_n_with_drain(&mut d, 5, 0.75, &mut batch_b);
2170        d.flush(|f| batch_b.push(f32_out(f)))
2171            .expect("second flush failed");
2172        assert_eq!(batch_b.len(), 5);
2173
2174        for v in batch_b.iter().flatten() {
2175            assert!((v - 0.75).abs() < 0.1, "batch_b carried state from batch_a: {v}");
2176        }
2177        for v in batch_a.iter().flatten() {
2178            assert!((v - 0.25).abs() < 0.1, "batch_a value unexpectedly drifted: {v}");
2179        }
2180    }
2181
2182    #[test]
2183    fn flush_leaves_denoiser_reusable_temporal() {
2184        let mut d = Denoiser::create(
2185            &[Accelerator::Vulkan],
2186            &Device::Default,
2187            16,
2188            16,
2189            opts(DenoisingMode::Temporal { radius: 1 }),
2190        )
2191        .unwrap();
2192
2193        let mut batch_a = Vec::new();
2194        push_n_with_drain(&mut d, 5, 0.25, &mut batch_a);
2195        d.flush(|f| batch_a.push(f32_out(f))).expect("first flush failed");
2196        assert_eq!(batch_a.len(), 5, "expected 5 frames from first batch");
2197
2198        // The temporal window must be empty after a flush, so the first
2199        // push of the new stream should not produce a pending frame.
2200        // With r=1 the window needs 3 frames before `denoise_submit`
2201        // fires.
2202        assert!(d.recv_frame().unwrap().is_none());
2203        d.push_frame(&frame_filled(16, 16, 0.75)).unwrap();
2204        assert!(
2205            d.recv_frame().unwrap().is_none(),
2206            "first push of new temporal stream should not produce output yet"
2207        );
2208
2209        // Push 4 more frames (5 total in batch B) with drain.
2210        let mut batch_b = Vec::new();
2211        push_n_with_drain(&mut d, 4, 0.75, &mut batch_b);
2212        d.flush(|f| batch_b.push(f32_out(f)))
2213            .expect("second flush failed");
2214        assert_eq!(batch_b.len(), 5, "expected 5 frames from second batch");
2215
2216        for v in batch_b.iter().flatten() {
2217            assert!((v - 0.75).abs() < 0.1, "batch_b carried state from batch_a: {v}");
2218        }
2219    }
2220
2221    #[test]
2222    fn flush_emits_exactly_n_outputs_for_small_n() {
2223        // With temporal radius R=2 the window is 5 frames. Pushing fewer
2224        // than R+1 frames means the window never fills while pushing, so
2225        // flush must still emit one output per pushed frame rather than
2226        // R+1 of them.
2227        for n in 1..=5usize {
2228            let mut d = Denoiser::create(
2229                &[Accelerator::Vulkan],
2230                &Device::Default,
2231                16,
2232                16,
2233                opts(DenoisingMode::Temporal { radius: 2 }),
2234            )
2235            .unwrap();
2236
2237            let mut out = Vec::new();
2238            push_n_with_drain(&mut d, n, 0.5, &mut out);
2239            d.flush(|f| out.push(f32_out(f))).expect("flush failed");
2240            assert_eq!(
2241                out.len(),
2242                n,
2243                "expected {n} outputs for {n} pushes, got {}",
2244                out.len()
2245            );
2246        }
2247    }
2248}