av_denoise/nlmeans/motion/
mod.rs

1mod analyse;
2mod compensate;
3mod pyramid;
4
5pub(crate) use analyse::run_analyse;
6pub(crate) use compensate::run_compensate;
7use cubecl::prelude::*;
8use cubecl::server::Handle;
9pub(crate) use pyramid::{pyramid_pixels_per_frame, run_pyramid_build};
10
11/// How motion compensation is configured for a denoise pass.
12///
13/// `None` disables motion compensation entirely (zero-cost; no extra
14/// buffers are allocated). `Mvtools` enables an MVTools-inspired
15/// per-block estimator and warps neighbours toward the centre at
16/// denoise time.
17#[non_exhaustive]
18#[derive(Debug, Default, Clone, Copy, PartialEq)]
19pub enum MotionCompensationMode {
20    #[default]
21    None,
22    Mvtools {
23        /// Side length of each motion-estimation block in pixels at
24        /// the finest pyramid level.
25        blksize: u32,
26        /// Overlap between neighbouring blocks in pixels. Must be
27        /// strictly less than `blksize` so the step (`blksize - overlap`)
28        /// stays positive. Values > 0 reserve room for raised-cosine
29        /// blending in the compensate step (v1 uses a winner-block rule).
30        overlap: u32,
31        /// Pixel search radius at the *finest* pyramid level. The
32        /// coarse pass uses the same radius on the `/2` image so its
33        /// effective reach is doubled.
34        search_radius: u32,
35        /// Number of pyramid levels. `1` disables the hierarchical
36        /// coarse pass; `2` adds a `/2` coarse pass that seeds the
37        /// fine pass. Bounded by [`MAX_PYRAMID_LEVELS`].
38        pyramid_levels: u32,
39    },
40}
41
42/// Default block size used when callers don't override it. Matches the
43/// MVTools default and lines up well with NLM's typical patch sizes.
44pub const DEFAULT_BLKSIZE: u32 = 16;
45/// Default block overlap (= `blksize / 2`).
46pub const DEFAULT_OVERLAP: u32 = 8;
47/// Default finest-level search radius. With a 2-level pyramid this
48/// reaches motion up to roughly ±12 pixels at the finest scale.
49pub const DEFAULT_SEARCH_RADIUS: u32 = 4;
50/// Default number of pyramid levels. `2` gives a single `/2` coarse
51/// pass, enough to handle most heavy-motion anime while keeping the
52/// kernel count manageable.
53pub const DEFAULT_PYRAMID_LEVELS: u32 = 2;
54
55/// Hard ceiling on `pyramid_levels`. Each extra level halves the
56/// resolution and adds an analyse-kernel launch per neighbour; 3 is
57/// already overkill for 1080p content.
58pub const MAX_PYRAMID_LEVELS: u32 = 3;
59/// Hard ceiling on `search_radius`. The analyse kernel SAD-sweeps a
60/// `(2·r + 1)²` window per block, so the cost is quadratic.
61pub const MAX_SEARCH_RADIUS: u32 = 8;
62/// Hard ceiling on `blksize`. Above this the per-block SMEM tile is
63/// uncomfortably large on RDNA-class GPUs.
64pub const MAX_BLKSIZE: u32 = 32;
65
66impl MotionCompensationMode {
67    /// Convenience constructor for `Mvtools` with library defaults.
68    pub fn mvtools_default() -> Self {
69        Self::Mvtools {
70            blksize: DEFAULT_BLKSIZE,
71            overlap: DEFAULT_OVERLAP,
72            search_radius: DEFAULT_SEARCH_RADIUS,
73            pyramid_levels: DEFAULT_PYRAMID_LEVELS,
74        }
75    }
76
77    /// Whether motion compensation is active at all.
78    pub(crate) fn is_active(self) -> bool {
79        !matches!(self, Self::None)
80    }
81
82    /// Reject parameter combinations that the kernels can't honour.
83    pub fn validate(&self) -> Result<(), anyhow::Error> {
84        let Self::Mvtools {
85            blksize,
86            overlap,
87            search_radius,
88            pyramid_levels,
89        } = *self
90        else {
91            return Ok(());
92        };
93
94        if blksize < 4 {
95            anyhow::bail!("motion-compensation blksize={blksize} is too small; minimum is 4 pixels per side");
96        }
97        if blksize > MAX_BLKSIZE {
98            anyhow::bail!(
99                "motion-compensation blksize={blksize} exceeds the supported maximum ({MAX_BLKSIZE})"
100            );
101        }
102        if blksize % 2 != 0 {
103            anyhow::bail!(
104                "motion-compensation blksize={blksize} must be even so the /2 coarse level is well-defined"
105            );
106        }
107        if overlap >= blksize {
108            anyhow::bail!(
109                "motion-compensation overlap={overlap} must be strictly less than blksize ({blksize}) so step > 0"
110            );
111        }
112        if search_radius == 0 || search_radius > MAX_SEARCH_RADIUS {
113            anyhow::bail!(
114                "motion-compensation search_radius={search_radius} must be in 1..={MAX_SEARCH_RADIUS}"
115            );
116        }
117        if pyramid_levels == 0 || pyramid_levels > MAX_PYRAMID_LEVELS {
118            anyhow::bail!(
119                "motion-compensation pyramid_levels={pyramid_levels} must be in 1..={MAX_PYRAMID_LEVELS}"
120            );
121        }
122
123        Ok(())
124    }
125}
126
127/// Per-denoiser MC state, owned by `NlmDenoiser` when MC is active.
128///
129/// Lives next to (not inside) the optional buffer handles so the hot
130/// dispatch path can fish out comptime-relevant scalars without
131/// pattern-matching the enum every call.
132/// Per-denoiser MC state cached at construction time so the hot
133/// dispatch path doesn't re-pattern-match the enum on every call.
134/// Holds only the fields actually read by analyse / compensate
135/// dispatchers; the full configuration lives on
136/// [`MotionCompensationMode`].
137#[derive(Debug, Clone)]
138pub(crate) struct MotionCtx {
139    pub blksize: u32,
140    pub step: u32,
141    pub search_radius: u32,
142    pub pyramid_levels: u32,
143    pub blocks_x: u32,
144    pub blocks_y: u32,
145}
146
147impl MotionCtx {
148    pub fn new(mode: MotionCompensationMode, width: u32, height: u32) -> Option<Self> {
149        let MotionCompensationMode::Mvtools {
150            blksize,
151            overlap,
152            search_radius,
153            pyramid_levels,
154        } = mode
155        else {
156            return None;
157        };
158
159        let step = blksize - overlap;
160        let blocks_x = width.div_ceil(step).max(1);
161        let blocks_y = height.div_ceil(step).max(1);
162
163        Some(Self {
164            blksize,
165            step,
166            search_radius,
167            pyramid_levels,
168            blocks_x,
169            blocks_y,
170        })
171    }
172
173    /// MV-field slot count per neighbour. One i16x2 per block.
174    pub fn mv_slots_per_neighbour(&self) -> usize {
175        (self.blocks_x * self.blocks_y) as usize
176    }
177}
178
179/// Build the per-frame pyramid for the slot just uploaded by
180/// `push_frame`. Cheap no-op if `pyramid_levels == 1`.
181#[allow(clippy::too_many_arguments)]
182pub(crate) fn build_pyramid_for_slot<R: Runtime>(
183    client: &ComputeClient<R>,
184    mc: &MotionCtx,
185    width: u32,
186    height: u32,
187    frame_count: u32,
188    slot: u32,
189    full_res: &Handle,
190    pyramid: &Handle,
191    stored_ch: u32,
192) -> Result<(), anyhow::Error> {
193    if mc.pyramid_levels <= 1 {
194        return Ok(());
195    }
196    run_pyramid_build::<R>(
197        client,
198        mc,
199        width,
200        height,
201        frame_count,
202        slot,
203        full_res,
204        pyramid,
205        stored_ch,
206    )
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212
213    #[test]
214    fn none_is_inactive() {
215        let m = MotionCompensationMode::None;
216        assert!(!m.is_active());
217        m.validate().unwrap();
218    }
219
220    #[test]
221    fn mvtools_default_is_active() {
222        let m = MotionCompensationMode::mvtools_default();
223        assert!(m.is_active());
224        m.validate().unwrap();
225    }
226
227    #[test]
228    fn validate_rejects_tiny_blksize() {
229        let m = MotionCompensationMode::Mvtools {
230            blksize: 2,
231            overlap: 0,
232            search_radius: 4,
233            pyramid_levels: 2,
234        };
235        assert!(m.validate().is_err());
236    }
237
238    #[test]
239    fn validate_rejects_odd_blksize() {
240        let m = MotionCompensationMode::Mvtools {
241            blksize: 9,
242            overlap: 0,
243            search_radius: 4,
244            pyramid_levels: 2,
245        };
246        assert!(m.validate().is_err());
247    }
248
249    #[test]
250    fn validate_rejects_overlap_equal_to_blksize() {
251        let m = MotionCompensationMode::Mvtools {
252            blksize: 16,
253            overlap: 16,
254            search_radius: 4,
255            pyramid_levels: 2,
256        };
257        // overlap == blksize would give step=0.
258        assert!(m.validate().is_err());
259    }
260
261    #[test]
262    fn validate_accepts_half_overlap() {
263        let m = MotionCompensationMode::Mvtools {
264            blksize: 16,
265            overlap: 8,
266            search_radius: 4,
267            pyramid_levels: 2,
268        };
269        m.validate().unwrap();
270    }
271
272    #[test]
273    fn validate_rejects_zero_search_radius() {
274        let m = MotionCompensationMode::Mvtools {
275            blksize: 16,
276            overlap: 4,
277            search_radius: 0,
278            pyramid_levels: 2,
279        };
280        assert!(m.validate().is_err());
281    }
282
283    #[test]
284    fn validate_rejects_zero_pyramid_levels() {
285        let m = MotionCompensationMode::Mvtools {
286            blksize: 16,
287            overlap: 4,
288            search_radius: 4,
289            pyramid_levels: 0,
290        };
291        assert!(m.validate().is_err());
292    }
293
294    #[test]
295    fn motion_ctx_blocks_match_step() {
296        let mode = MotionCompensationMode::Mvtools {
297            blksize: 16,
298            overlap: 8,
299            search_radius: 4,
300            pyramid_levels: 2,
301        };
302        let ctx = MotionCtx::new(mode, 1920, 1080).unwrap();
303        assert_eq!(ctx.step, 8);
304        assert_eq!(ctx.blocks_x, 1920u32.div_ceil(8));
305        assert_eq!(ctx.blocks_y, 1080u32.div_ceil(8));
306    }
307}
av_denoise/nlmeans/motion/mod.rs

av_denoise/nlmeans/motion/
mod.rs
