av_scenechange/data/
tile.rs

1use std::iter::FusedIterator;
2
3use v_frame::{
4    frame::Frame,
5    math::Fixed,
6    pixel::Pixel,
7    plane::{Plane, PlaneOffset},
8};
9
10use crate::data::{
11    block::BlockOffset,
12    frame::{FrameState, MAX_PLANES},
13    motion::{FrameMEStats, TileMEStatsMut, WriteGuardMEStats},
14    plane::{PlaneBlockOffset, PlaneRegion, Rect},
15    superblock::{PlaneSuperBlockOffset, SuperBlockOffset, MI_SIZE, MI_SIZE_LOG2, SB_SIZE_LOG2},
16};
17
18pub const MAX_TILE_WIDTH: usize = 4096;
19pub const MAX_TILE_AREA: usize = 4096 * 2304;
20pub const MAX_TILE_COLS: usize = 64;
21pub const MAX_TILE_ROWS: usize = 64;
22pub const MAX_TILE_RATE: f64 = 4096f64 * 2176f64 * 60f64 * 1.1;
23
24/// Tiled view of a frame
25#[derive(Debug)]
26pub struct Tile<'a, T: Pixel> {
27    pub planes: [PlaneRegion<'a, T>; MAX_PLANES],
28}
29
30// common impl for Tile and TileMut
31macro_rules! tile_common {
32  // $name: Tile or TileMut
33  // $pr_type: PlaneRegion or PlaneRegionMut
34  // $iter: iter or iter_mut
35  //opt_mut: nothing or mut
36  ($name:ident, $pr_type:ident, $iter:ident $(,$opt_mut:tt)?) => {
37    impl<'a, T: Pixel> $name<'a, T> {
38
39      pub fn new(
40        frame: &'a $($opt_mut)? Frame<T>,
41        luma_rect: TileRect,
42      ) -> Self {
43        let mut planes_iter = frame.planes.$iter();
44        Self {
45          planes: [
46            {
47              let plane = planes_iter.next().unwrap();
48              $pr_type::new(plane, luma_rect.into())
49            },
50            {
51              let plane = planes_iter.next().unwrap();
52              let rect = luma_rect.decimated(plane.cfg.xdec, plane.cfg.ydec);
53              $pr_type::new(plane, rect.into())
54            },
55            {
56              let plane = planes_iter.next().unwrap();
57              let rect = luma_rect.decimated(plane.cfg.xdec, plane.cfg.ydec);
58              $pr_type::new(plane, rect.into())
59            },
60          ],
61        }
62      }
63    }
64  }
65}
66
67tile_common!(Tile, PlaneRegion, iter);
68
69/// Rectangle of a tile, in pixels
70///
71/// This is similar to Rect, but with unsigned (x, y) for convenience.
72#[derive(Debug, Clone, Copy)]
73pub struct TileRect {
74    pub x: usize,
75    pub y: usize,
76    pub width: usize,
77    pub height: usize,
78}
79
80impl TileRect {
81    pub const fn decimated(self, xdec: usize, ydec: usize) -> Self {
82        Self {
83            x: self.x >> xdec,
84            y: self.y >> ydec,
85            width: self.width >> xdec,
86            height: self.height >> ydec,
87        }
88    }
89
90    pub const fn to_frame_plane_offset(self, tile_po: PlaneOffset) -> PlaneOffset {
91        PlaneOffset {
92            x: self.x as isize + tile_po.x,
93            y: self.y as isize + tile_po.y,
94        }
95    }
96}
97
98impl From<TileRect> for Rect {
99    fn from(tile_rect: TileRect) -> Rect {
100        Rect {
101            x: tile_rect.x as isize,
102            y: tile_rect.y as isize,
103            width: tile_rect.width,
104            height: tile_rect.height,
105        }
106    }
107}
108
109/// Tiled view of `FrameState`
110///
111/// Contrary to `PlaneRegionMut` and `TileMut`, there is no const version:
112///  - in practice, we don't need it;
113///  - it would require to instantiate a const version of every of its inner
114///    tiled views recursively.
115///
116/// # `TileState` fields
117///
118/// The way the `FrameState` fields are mapped depend on how they are accessed
119/// tile-wise and frame-wise.
120///
121/// Some fields (like `qc`) are only used during tile-encoding, so they are only
122/// stored in `TileState`.
123///
124/// Some other fields (like `input` or `segmentation`) are not written
125/// tile-wise, so they just reference the matching field in `FrameState`.
126///
127/// Some others (like `rec`) are written tile-wise, but must be accessible
128/// frame-wise once the tile views vanish (e.g. for deblocking).
129#[derive(Debug)]
130pub struct TileStateMut<'a, T: Pixel> {
131    pub sbo: PlaneSuperBlockOffset,
132    pub sb_width: usize,
133    pub sb_height: usize,
134    pub mi_width: usize,
135    pub mi_height: usize,
136    pub width: usize,
137    pub height: usize,
138    pub input_tile: Tile<'a, T>, // the current tile
139    pub input_hres: &'a Plane<T>,
140    pub input_qres: &'a Plane<T>,
141    pub me_stats: Vec<TileMEStatsMut<'a>>,
142}
143
144impl<'a, T: Pixel> TileStateMut<'a, T> {
145    pub fn new(
146        fs: &'a mut FrameState<T>,
147        sbo: PlaneSuperBlockOffset,
148        width: usize,
149        height: usize,
150        frame_me_stats: &'a mut [FrameMEStats],
151    ) -> Self {
152        debug_assert!(
153            width % MI_SIZE == 0,
154            "Tile width must be a multiple of MI_SIZE"
155        );
156        debug_assert!(
157            height % MI_SIZE == 0,
158            "Tile width must be a multiple of MI_SIZE"
159        );
160
161        let sb_rounded_width = width.align_power_of_two(SB_SIZE_LOG2);
162        let sb_rounded_height = height.align_power_of_two(SB_SIZE_LOG2);
163
164        let luma_rect = TileRect {
165            x: sbo.0.x << SB_SIZE_LOG2,
166            y: sbo.0.y << SB_SIZE_LOG2,
167            width: sb_rounded_width,
168            height: sb_rounded_height,
169        };
170        let sb_width = width.align_power_of_two_and_shift(SB_SIZE_LOG2);
171        let sb_height = height.align_power_of_two_and_shift(SB_SIZE_LOG2);
172
173        Self {
174            sbo,
175            sb_width,
176            sb_height,
177            mi_width: width >> MI_SIZE_LOG2,
178            mi_height: height >> MI_SIZE_LOG2,
179            width,
180            height,
181            input_tile: Tile::new(&fs.input, luma_rect),
182            input_hres: &fs.input_hres,
183            input_qres: &fs.input_qres,
184            me_stats: frame_me_stats
185                .iter_mut()
186                .map(|fmvs| {
187                    TileMEStatsMut::new(
188                        fmvs,
189                        sbo.0.x << (SB_SIZE_LOG2 - MI_SIZE_LOG2),
190                        sbo.0.y << (SB_SIZE_LOG2 - MI_SIZE_LOG2),
191                        width >> MI_SIZE_LOG2,
192                        height >> MI_SIZE_LOG2,
193                    )
194                })
195                .collect(),
196        }
197    }
198
199    pub fn to_frame_block_offset(&self, tile_bo: TileBlockOffset) -> PlaneBlockOffset {
200        let bx = self.sbo.0.x << (SB_SIZE_LOG2 - MI_SIZE_LOG2);
201        let by = self.sbo.0.y << (SB_SIZE_LOG2 - MI_SIZE_LOG2);
202        PlaneBlockOffset(BlockOffset {
203            x: bx + tile_bo.0.x,
204            y: by + tile_bo.0.y,
205        })
206    }
207}
208
209/// Absolute offset in blocks inside a tile, where a block is defined
210/// to be an `N*N` square where `N == (1 << BLOCK_TO_PLANE_SHIFT)`.
211#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
212pub struct TileBlockOffset(pub BlockOffset);
213
214impl TileBlockOffset {
215    /// Convert to plane offset without decimation.
216    pub const fn to_luma_plane_offset(self) -> PlaneOffset {
217        self.0.to_luma_plane_offset()
218    }
219
220    pub fn with_offset(self, col_offset: isize, row_offset: isize) -> TileBlockOffset {
221        Self(self.0.with_offset(col_offset, row_offset))
222    }
223}
224
225/// Tiling information
226///
227/// This stores everything necessary to split a frame into tiles, and write
228/// headers fields into the bitstream.
229///
230/// The method `tile_iter_mut()` actually provides tiled views of `FrameState`
231/// and `FrameBlocks`.
232#[derive(Debug, Clone, Copy)]
233pub struct TilingInfo {
234    pub frame_width: usize,
235    pub frame_height: usize,
236    pub tile_width_sb: usize,
237    pub tile_height_sb: usize,
238    pub cols: usize, // number of columns of tiles within the whole frame
239    pub rows: usize, // number of rows of tiles within the whole frame
240}
241
242impl TilingInfo {
243    /// # Panics
244    ///
245    /// Panics if the resulting tile sizes would be too large.
246    pub fn from_target_tiles(
247        frame_width: usize,
248        frame_height: usize,
249        frame_rate: f64,
250        tile_cols_log2: usize,
251        tile_rows_log2: usize,
252        is_422_p: bool,
253    ) -> Self {
254        // <https://aomediacodec.github.io/av1-spec/#tile-info-syntax>
255
256        // Frame::new() aligns to the next multiple of 8
257        let frame_width = frame_width.align_power_of_two(3);
258        let frame_height = frame_height.align_power_of_two(3);
259        let frame_width_sb = frame_width.align_power_of_two_and_shift(SB_SIZE_LOG2);
260        let frame_height_sb = frame_height.align_power_of_two_and_shift(SB_SIZE_LOG2);
261        let sb_cols = frame_width.align_power_of_two_and_shift(SB_SIZE_LOG2);
262        let sb_rows = frame_height.align_power_of_two_and_shift(SB_SIZE_LOG2);
263
264        // these are bitstream-defined values and must not be changed
265        let max_tile_width_sb = MAX_TILE_WIDTH >> SB_SIZE_LOG2;
266        let max_tile_area_sb = MAX_TILE_AREA >> (2 * SB_SIZE_LOG2);
267        let min_tile_cols_log2 = Self::tile_log2(max_tile_width_sb, sb_cols).unwrap();
268        let max_tile_cols_log2 = Self::tile_log2(1, sb_cols.min(MAX_TILE_COLS)).unwrap();
269        let max_tile_rows_log2 = Self::tile_log2(1, sb_rows.min(MAX_TILE_ROWS)).unwrap();
270        let min_tiles_log2 =
271            min_tile_cols_log2.max(Self::tile_log2(max_tile_area_sb, sb_cols * sb_rows).unwrap());
272
273        // Implements restriction in Annex A of the spec.
274        // Unlike the other restrictions, this one does not change
275        // the header coding of the tile rows/cols.
276        let min_tiles_ratelimit_log2 = min_tiles_log2.max(
277            ((frame_width * frame_height) as f64 * frame_rate / MAX_TILE_RATE)
278                .ceil()
279                .log2()
280                .ceil() as usize,
281        );
282
283        let tile_cols_log2 = tile_cols_log2.clamp(min_tile_cols_log2, max_tile_cols_log2);
284        let tile_width_sb_pre = sb_cols.align_power_of_two_and_shift(tile_cols_log2);
285
286        // If this is 4:2:2, our UV horizontal is subsampled but not our
287        // vertical.  Loop Restoration Units must be square, so they
288        // will always have an even number of horizontal superblocks. For
289        // tiles and LRUs to align, tile_width_sb must be even in 4:2:2
290        // video.
291
292        // This is only relevant when doing loop restoration RDO inline
293        // with block/superblock encoding, that is, where tiles are
294        // relevant.  If (when) we introduce optionally delaying loop-filter
295        // encode to after the partitioning loop, we won't need to make
296        // any 4:2:2 adjustment.
297
298        let tile_width_sb = if is_422_p {
299            (tile_width_sb_pre + 1) >> 1 << 1
300        } else {
301            tile_width_sb_pre
302        };
303
304        let cols = frame_width_sb.div_ceil(tile_width_sb);
305
306        // Adjust tile_cols_log2 in case of rounding tile_width_sb to even.
307        let tile_cols_log2 = Self::tile_log2(1, cols).unwrap();
308        assert!(tile_cols_log2 >= min_tile_cols_log2);
309
310        let min_tile_rows_log2 = min_tiles_log2.saturating_sub(tile_cols_log2);
311        let min_tile_rows_ratelimit_log2 = min_tiles_ratelimit_log2.saturating_sub(tile_cols_log2);
312        let tile_rows_log2 = tile_rows_log2
313            .max(min_tile_rows_log2)
314            .clamp(min_tile_rows_ratelimit_log2, max_tile_rows_log2);
315        let tile_height_sb = sb_rows.align_power_of_two_and_shift(tile_rows_log2);
316
317        let rows = frame_height_sb.div_ceil(tile_height_sb);
318
319        Self {
320            frame_width,
321            frame_height,
322            tile_width_sb,
323            tile_height_sb,
324            cols,
325            rows,
326        }
327    }
328
329    /// Return the smallest value for `k` such that `blkSize << k` is greater
330    /// than or equal to `target`.
331    ///
332    /// <https://aomediacodec.github.io/av1-spec/#tile-size-calculation-function>
333    pub fn tile_log2(blk_size: usize, target: usize) -> Option<usize> {
334        let mut k = 0;
335        while (blk_size.checked_shl(k)?) < target {
336            k += 1;
337        }
338        Some(k as usize)
339    }
340
341    /// Split frame-level structures into tiles
342    ///
343    /// Provide mutable tiled views of frame-level structures.
344    pub fn tile_iter_mut<'a, T: Pixel>(
345        &self,
346        fs: &'a mut FrameState<T>,
347    ) -> TileContextIterMut<'a, T> {
348        let afs = fs as *mut _;
349        let frame_me_stats = fs.frame_me_stats.write().expect("poisoned lock");
350        TileContextIterMut {
351            ti: *self,
352            fs: afs,
353            next: 0,
354            frame_me_stats,
355        }
356    }
357}
358
359/// Iterator over tiled views
360pub struct TileContextIterMut<'a, T: Pixel> {
361    ti: TilingInfo,
362    fs: *mut FrameState<T>,
363    frame_me_stats: WriteGuardMEStats<'a>,
364    next: usize,
365}
366
367impl<'a, T: Pixel> Iterator for TileContextIterMut<'a, T> {
368    type Item = TileContextMut<'a, T>;
369
370    fn next(&mut self) -> Option<Self::Item> {
371        if self.next < self.ti.rows * self.ti.cols {
372            let tile_col = self.next % self.ti.cols;
373            let tile_row = self.next / self.ti.cols;
374            let ctx = TileContextMut {
375                ts: {
376                    // SAFETY: Multiple tiles mutably access this struct.
377                    // The dimensions must be configured correctly to ensure
378                    // the tiles do not overlap.
379                    let fs = unsafe { &mut *self.fs };
380                    // SAFETY: ditto
381                    let frame_me_stats = unsafe {
382                        let len = self.frame_me_stats.len();
383                        let ptr = self.frame_me_stats.as_mut_ptr();
384                        std::slice::from_raw_parts_mut(ptr, len)
385                    };
386                    let sbo = PlaneSuperBlockOffset(SuperBlockOffset {
387                        x: tile_col * self.ti.tile_width_sb,
388                        y: tile_row * self.ti.tile_height_sb,
389                    });
390                    let x = sbo.0.x << SB_SIZE_LOG2;
391                    let y = sbo.0.y << SB_SIZE_LOG2;
392                    let tile_width = self.ti.tile_width_sb << SB_SIZE_LOG2;
393                    let tile_height = self.ti.tile_height_sb << SB_SIZE_LOG2;
394                    let width = tile_width.min(self.ti.frame_width - x);
395                    let height = tile_height.min(self.ti.frame_height - y);
396                    TileStateMut::new(fs, sbo, width, height, frame_me_stats)
397                },
398            };
399            self.next += 1;
400            Some(ctx)
401        } else {
402            None
403        }
404    }
405
406    fn size_hint(&self) -> (usize, Option<usize>) {
407        let remaining = self.ti.cols * self.ti.rows - self.next;
408        (remaining, Some(remaining))
409    }
410}
411
412impl<T: Pixel> ExactSizeIterator for TileContextIterMut<'_, T> {
413}
414impl<T: Pixel> FusedIterator for TileContextIterMut<'_, T> {
415}
416
417/// Container for all tiled views
418pub struct TileContextMut<'a, T: Pixel> {
419    pub ts: TileStateMut<'a, T>,
420}
av_scenechange/data/tile.rs

av_scenechange/data/
tile.rs
