rav1e/tiling/
tiler.rs

1// Copyright (c) 2019-2022, The rav1e contributors. All rights reserved
2//
3// This source code is subject to the terms of the BSD 2 Clause License and
4// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
5// was not distributed with this source code in the LICENSE file, you can
6// obtain it at www.aomedia.org/license/software. If the Alliance for Open
7// Media Patent License 1.0 was not distributed with this source code in the
8// PATENTS file, you can obtain it at www.aomedia.org/license/patent.
9
10use super::*;
11
12use crate::context::*;
13use crate::encoder::*;
14use crate::me::WriteGuardMEStats;
15use crate::util::*;
16
17use std::iter::FusedIterator;
18use std::marker::PhantomData;
19use std::ops::DerefMut;
20
21pub const MAX_TILE_WIDTH: usize = 4096;
22pub const MAX_TILE_AREA: usize = 4096 * 2304;
23pub const MAX_TILE_COLS: usize = 64;
24pub const MAX_TILE_ROWS: usize = 64;
25pub const MAX_TILE_RATE: f64 = 4096f64 * 2176f64 * 60f64 * 1.1;
26
27/// Tiling information
28///
29/// This stores everything necessary to split a frame into tiles, and write
30/// headers fields into the bitstream.
31///
32/// The method `tile_iter_mut()` actually provides tiled views of `FrameState`
33/// and `FrameBlocks`.
34#[derive(Debug, Clone, Copy)]
35pub struct TilingInfo {
36  pub frame_width: usize,
37  pub frame_height: usize,
38  pub tile_width_sb: usize,
39  pub tile_height_sb: usize,
40  pub cols: usize, // number of columns of tiles within the whole frame
41  pub rows: usize, // number of rows of tiles within the whole frame
42  pub tile_cols_log2: usize,
43  pub tile_rows_log2: usize,
44  pub min_tile_cols_log2: usize,
45  pub max_tile_cols_log2: usize,
46  pub min_tile_rows_log2: usize,
47  pub max_tile_rows_log2: usize,
48  pub sb_size_log2: usize,
49  pub min_tiles_log2: usize,
50}
51
52impl TilingInfo {
53  /// # Panics
54  ///
55  /// Panics if the resulting tile sizes would be too large.
56  pub fn from_target_tiles(
57    sb_size_log2: usize, frame_width: usize, frame_height: usize,
58    frame_rate: f64, tile_cols_log2: usize, tile_rows_log2: usize,
59    is_422_p: bool,
60  ) -> Self {
61    // <https://aomediacodec.github.io/av1-spec/#tile-info-syntax>
62
63    // Frame::new() aligns to the next multiple of 8
64    let frame_width = frame_width.align_power_of_two(3);
65    let frame_height = frame_height.align_power_of_two(3);
66    let frame_width_sb =
67      frame_width.align_power_of_two_and_shift(sb_size_log2);
68    let frame_height_sb =
69      frame_height.align_power_of_two_and_shift(sb_size_log2);
70    let sb_cols = frame_width.align_power_of_two_and_shift(sb_size_log2);
71    let sb_rows = frame_height.align_power_of_two_and_shift(sb_size_log2);
72
73    // these are bitstream-defined values and must not be changed
74    let max_tile_width_sb = MAX_TILE_WIDTH >> sb_size_log2;
75    let max_tile_area_sb = MAX_TILE_AREA >> (2 * sb_size_log2);
76    let min_tile_cols_log2 =
77      Self::tile_log2(max_tile_width_sb, sb_cols).unwrap();
78    let max_tile_cols_log2 =
79      Self::tile_log2(1, sb_cols.min(MAX_TILE_COLS)).unwrap();
80    let max_tile_rows_log2 =
81      Self::tile_log2(1, sb_rows.min(MAX_TILE_ROWS)).unwrap();
82    let min_tiles_log2 = min_tile_cols_log2
83      .max(Self::tile_log2(max_tile_area_sb, sb_cols * sb_rows).unwrap());
84
85    // Implements restriction in Annex A of the spec.
86    // Unlike the other restrictions, this one does not change
87    // the header coding of the tile rows/cols.
88    let min_tiles_ratelimit_log2 = min_tiles_log2.max(
89      ((frame_width * frame_height) as f64 * frame_rate / MAX_TILE_RATE)
90        .ceil()
91        .log2()
92        .ceil() as usize,
93    );
94
95    let tile_cols_log2 =
96      tile_cols_log2.clamp(min_tile_cols_log2, max_tile_cols_log2);
97    let tile_width_sb_pre =
98      sb_cols.align_power_of_two_and_shift(tile_cols_log2);
99
100    // If this is 4:2:2, our UV horizontal is subsampled but not our
101    // vertical.  Loop Restoration Units must be square, so they
102    // will always have an even number of horizontal superblocks. For
103    // tiles and LRUs to align, tile_width_sb must be even in 4:2:2
104    // video.
105
106    // This is only relevant when doing loop restoration RDO inline
107    // with block/superblock encoding, that is, where tiles are
108    // relevant.  If (when) we introduce optionally delaying loop-filter
109    // encode to after the partitioning loop, we won't need to make
110    // any 4:2:2 adjustment.
111
112    let tile_width_sb = if is_422_p {
113      (tile_width_sb_pre + 1) >> 1 << 1
114    } else {
115      tile_width_sb_pre
116    };
117
118    let cols = (frame_width_sb + tile_width_sb - 1) / tile_width_sb;
119
120    // Adjust tile_cols_log2 in case of rounding tile_width_sb to even.
121    let tile_cols_log2 = Self::tile_log2(1, cols).unwrap();
122    assert!(tile_cols_log2 >= min_tile_cols_log2);
123
124    let min_tile_rows_log2 = if min_tiles_log2 > tile_cols_log2 {
125      min_tiles_log2 - tile_cols_log2
126    } else {
127      0
128    };
129    let min_tile_rows_ratelimit_log2 =
130      if min_tiles_ratelimit_log2 > tile_cols_log2 {
131        min_tiles_ratelimit_log2 - tile_cols_log2
132      } else {
133        0
134      };
135    let tile_rows_log2 = tile_rows_log2
136      .max(min_tile_rows_log2)
137      .clamp(min_tile_rows_ratelimit_log2, max_tile_rows_log2);
138    let tile_height_sb = sb_rows.align_power_of_two_and_shift(tile_rows_log2);
139
140    let rows = (frame_height_sb + tile_height_sb - 1) / tile_height_sb;
141
142    Self {
143      frame_width,
144      frame_height,
145      tile_width_sb,
146      tile_height_sb,
147      cols,
148      rows,
149      tile_cols_log2,
150      tile_rows_log2,
151      min_tile_cols_log2,
152      max_tile_cols_log2,
153      min_tile_rows_log2,
154      max_tile_rows_log2,
155      sb_size_log2,
156      min_tiles_log2,
157    }
158  }
159
160  /// Return the smallest value for `k` such that `blkSize << k` is greater than
161  /// or equal to `target`.
162  ///
163  /// <https://aomediacodec.github.io/av1-spec/#tile-size-calculation-function>
164  pub fn tile_log2(blk_size: usize, target: usize) -> Option<usize> {
165    let mut k = 0;
166    while (blk_size.checked_shl(k)?) < target {
167      k += 1;
168    }
169    Some(k as usize)
170  }
171
172  #[inline(always)]
173  pub const fn tile_count(&self) -> usize {
174    self.cols * self.rows
175  }
176
177  /// Split frame-level structures into tiles
178  ///
179  /// Provide mutable tiled views of frame-level structures.
180  pub fn tile_iter_mut<'a, T: Pixel>(
181    &self, fs: &'a mut FrameState<T>, fb: &'a mut FrameBlocks,
182  ) -> TileContextIterMut<'a, T> {
183    let afs = fs as *mut _;
184    let afb = fb as *mut _;
185    let frame_me_stats = fs.frame_me_stats.write().expect("poisoned lock");
186    TileContextIterMut { ti: *self, fs: afs, fb: afb, next: 0, frame_me_stats }
187  }
188}
189
190/// Container for all tiled views
191pub struct TileContextMut<'a, T: Pixel> {
192  pub ts: TileStateMut<'a, T>,
193  pub tb: TileBlocksMut<'a>,
194}
195
196/// Iterator over tiled views
197pub struct TileContextIterMut<'a, T: Pixel> {
198  ti: TilingInfo,
199  fs: *mut FrameState<T>,
200  fb: *mut FrameBlocks,
201  frame_me_stats: WriteGuardMEStats<'a>,
202  next: usize,
203}
204
205impl<'a, T: Pixel> Iterator for TileContextIterMut<'a, T> {
206  type Item = TileContextMut<'a, T>;
207
208  fn next(&mut self) -> Option<Self::Item> {
209    if self.next < self.ti.rows * self.ti.cols {
210      let tile_col = self.next % self.ti.cols;
211      let tile_row = self.next / self.ti.cols;
212      let ctx = TileContextMut {
213        ts: {
214          // SAFETY: Multiple tiles mutably access this struct.
215          // The dimensions must be configured correctly to ensure
216          // the tiles do not overlap.
217          let fs = unsafe { &mut *self.fs };
218          // SAFETY: ditto
219          let frame_me_stats = unsafe {
220            let len = self.frame_me_stats.len();
221            let ptr = self.frame_me_stats.as_mut_ptr();
222            std::slice::from_raw_parts_mut(ptr, len)
223          };
224          let sbo = PlaneSuperBlockOffset(SuperBlockOffset {
225            x: tile_col * self.ti.tile_width_sb,
226            y: tile_row * self.ti.tile_height_sb,
227          });
228          let x = sbo.0.x << self.ti.sb_size_log2;
229          let y = sbo.0.y << self.ti.sb_size_log2;
230          let tile_width = self.ti.tile_width_sb << self.ti.sb_size_log2;
231          let tile_height = self.ti.tile_height_sb << self.ti.sb_size_log2;
232          let width = tile_width.min(self.ti.frame_width - x);
233          let height = tile_height.min(self.ti.frame_height - y);
234          TileStateMut::new(
235            fs,
236            sbo,
237            self.ti.sb_size_log2,
238            width,
239            height,
240            frame_me_stats,
241          )
242        },
243        tb: {
244          // SAFETY: Multiple tiles mutably access this struct.
245          // The dimensions must be configured correctly to ensure
246          // the tiles do not overlap.
247          let fb = unsafe { &mut *self.fb };
248          let tile_width_mi =
249            self.ti.tile_width_sb << (self.ti.sb_size_log2 - MI_SIZE_LOG2);
250          let tile_height_mi =
251            self.ti.tile_height_sb << (self.ti.sb_size_log2 - MI_SIZE_LOG2);
252          let x = tile_col * tile_width_mi;
253          let y = tile_row * tile_height_mi;
254          let cols = tile_width_mi.min(fb.cols - x);
255          let rows = tile_height_mi.min(fb.rows - y);
256          TileBlocksMut::new(fb, x, y, cols, rows)
257        },
258      };
259      self.next += 1;
260      Some(ctx)
261    } else {
262      None
263    }
264  }
265
266  fn size_hint(&self) -> (usize, Option<usize>) {
267    let remaining = self.ti.cols * self.ti.rows - self.next;
268    (remaining, Some(remaining))
269  }
270}
271
272impl<T: Pixel> ExactSizeIterator for TileContextIterMut<'_, T> {}
273impl<T: Pixel> FusedIterator for TileContextIterMut<'_, T> {}
274
275#[cfg(test)]
276pub mod test {
277  use super::*;
278  use crate::api::*;
279  use crate::lrf::*;
280  use crate::mc::MotionVector;
281  use crate::predict::PredictionMode;
282  use std::sync::Arc;
283
284  #[test]
285  fn test_tiling_info_from_tile_count() {
286    let sb_size_log2 = 6;
287    let (width, height) = (160, 144);
288    let frame_rate = 25f64;
289
290    let ti = TilingInfo::from_target_tiles(
291      sb_size_log2,
292      width,
293      height,
294      frame_rate,
295      0,
296      0,
297      false,
298    );
299    assert_eq!(1, ti.cols);
300    assert_eq!(1, ti.rows);
301    assert_eq!(3, ti.tile_width_sb);
302    assert_eq!(3, ti.tile_height_sb);
303
304    let ti = TilingInfo::from_target_tiles(
305      sb_size_log2,
306      width,
307      height,
308      frame_rate,
309      1,
310      1,
311      false,
312    );
313    assert_eq!(2, ti.cols);
314    assert_eq!(2, ti.rows);
315    assert_eq!(2, ti.tile_width_sb);
316    assert_eq!(2, ti.tile_height_sb);
317
318    let ti = TilingInfo::from_target_tiles(
319      sb_size_log2,
320      width,
321      height,
322      frame_rate,
323      2,
324      2,
325      false,
326    );
327    assert_eq!(3, ti.cols);
328    assert_eq!(3, ti.rows);
329    assert_eq!(1, ti.tile_width_sb);
330    assert_eq!(1, ti.tile_height_sb);
331
332    // cannot split more than superblocks
333    let ti = TilingInfo::from_target_tiles(
334      sb_size_log2,
335      width,
336      height,
337      frame_rate,
338      10,
339      8,
340      false,
341    );
342    assert_eq!(3, ti.cols);
343    assert_eq!(3, ti.rows);
344    assert_eq!(1, ti.tile_width_sb);
345    assert_eq!(1, ti.tile_height_sb);
346
347    let ti = TilingInfo::from_target_tiles(
348      sb_size_log2,
349      1024,
350      1024,
351      frame_rate,
352      0,
353      0,
354      false,
355    );
356    assert_eq!(1, ti.cols);
357    assert_eq!(1, ti.rows);
358    assert_eq!(16, ti.tile_width_sb);
359    assert_eq!(16, ti.tile_height_sb);
360  }
361
362  fn setup(
363    width: usize, height: usize,
364  ) -> (FrameInvariants<u16>, FrameState<u16>, FrameBlocks, f64) {
365    // FrameInvariants aligns to the next multiple of 8, so using other values could make tests confusing
366    assert!(width & 7 == 0);
367    assert!(height & 7 == 0);
368    // We test only for 420 for now
369    let chroma_sampling = ChromaSampling::Cs420;
370    let config = Arc::new(EncoderConfig {
371      width,
372      height,
373      bit_depth: 8,
374      chroma_sampling,
375      ..Default::default()
376    });
377    let mut sequence = Sequence::new(&config);
378    // These tests are all assuming SB-sized LRUs, so set that.
379    sequence.enable_large_lru = false;
380    let frame_rate = config.frame_rate();
381    let fi = FrameInvariants::new(config, Arc::new(sequence));
382    let fs = FrameState::new(&fi);
383    let fb = FrameBlocks::new(fi.w_in_b, fi.h_in_b);
384
385    (fi, fs, fb, frame_rate)
386  }
387
388  #[test]
389  fn test_tile_iter_len() {
390    // frame size 160x144, 40x36 in 4x4-blocks
391    let (fi, mut fs, mut fb, frame_rate) = setup(160, 144);
392
393    {
394      // 2x2 tiles
395      let ti = TilingInfo::from_target_tiles(
396        fi.sb_size_log2(),
397        fi.width,
398        fi.height,
399        frame_rate,
400        1,
401        1,
402        false,
403      );
404      let mut iter = ti.tile_iter_mut(&mut fs, &mut fb);
405      assert_eq!(4, iter.len());
406      assert!(iter.next().is_some());
407      assert_eq!(3, iter.len());
408      assert!(iter.next().is_some());
409      assert_eq!(2, iter.len());
410      assert!(iter.next().is_some());
411      assert_eq!(1, iter.len());
412      assert!(iter.next().is_some());
413      assert_eq!(0, iter.len());
414      assert!(iter.next().is_none());
415    }
416
417    {
418      // 4x4 tiles requested, will actually get 3x3 tiles
419      let ti = TilingInfo::from_target_tiles(
420        fi.sb_size_log2(),
421        fi.width,
422        fi.height,
423        frame_rate,
424        2,
425        2,
426        false,
427      );
428      let mut iter = ti.tile_iter_mut(&mut fs, &mut fb);
429      assert_eq!(9, iter.len());
430      assert!(iter.next().is_some());
431      assert_eq!(8, iter.len());
432      assert!(iter.next().is_some());
433      assert_eq!(7, iter.len());
434      assert!(iter.next().is_some());
435      assert_eq!(6, iter.len());
436      assert!(iter.next().is_some());
437      assert_eq!(5, iter.len());
438      assert!(iter.next().is_some());
439      assert_eq!(4, iter.len());
440      assert!(iter.next().is_some());
441      assert_eq!(3, iter.len());
442      assert!(iter.next().is_some());
443      assert_eq!(2, iter.len());
444      assert!(iter.next().is_some());
445      assert_eq!(1, iter.len());
446      assert!(iter.next().is_some());
447      assert_eq!(0, iter.len());
448      assert!(iter.next().is_none());
449    }
450  }
451
452  #[inline]
453  fn rect<T: Pixel>(
454    region: &PlaneRegionMut<'_, T>,
455  ) -> (isize, isize, usize, usize) {
456    let &Rect { x, y, width, height } = region.rect();
457    (x, y, width, height)
458  }
459
460  #[test]
461  fn test_tile_area() {
462    let (fi, mut fs, mut fb, frame_rate) = setup(160, 144);
463
464    // 4x4 tiles requested, will actually get 3x3 tiles
465    let ti = TilingInfo::from_target_tiles(
466      fi.sb_size_log2(),
467      fi.width,
468      fi.height,
469      frame_rate,
470      2,
471      2,
472      false,
473    );
474    let iter = ti.tile_iter_mut(&mut fs, &mut fb);
475    let tile_states = iter.map(|ctx| ctx.ts).collect::<Vec<_>>();
476
477    // the frame must be split into 9 tiles:
478    //
479    //       luma (Y)             chroma (U)            chroma (V)
480    //   64x64 64x64 32x64     32x32 32x32 16x32     32x32 32x32 16x32
481    //   64x64 64x64 32x64     32x32 32x32 16x32     32x32 32x32 16x32
482    //   64x16 64x16 32x16     32x 8 32x 8 16x 8     32x 8 32x 8 16x 8
483
484    assert_eq!(9, tile_states.len());
485
486    let tile = &tile_states[0].rec; // the top-left tile
487    assert_eq!((0, 0, 64, 64), rect(&tile.planes[0]));
488    assert_eq!((0, 0, 32, 32), rect(&tile.planes[1]));
489    assert_eq!((0, 0, 32, 32), rect(&tile.planes[2]));
490
491    let tile = &tile_states[1].rec; // the top-middle tile
492    assert_eq!((64, 0, 64, 64), rect(&tile.planes[0]));
493    assert_eq!((32, 0, 32, 32), rect(&tile.planes[1]));
494    assert_eq!((32, 0, 32, 32), rect(&tile.planes[2]));
495
496    let tile = &tile_states[2].rec; // the top-right tile
497    assert_eq!((128, 0, 64, 64), rect(&tile.planes[0]));
498    assert_eq!((64, 0, 32, 32), rect(&tile.planes[1]));
499    assert_eq!((64, 0, 32, 32), rect(&tile.planes[2]));
500
501    let tile = &tile_states[3].rec; // the middle-left tile
502    assert_eq!((0, 64, 64, 64), rect(&tile.planes[0]));
503    assert_eq!((0, 32, 32, 32), rect(&tile.planes[1]));
504    assert_eq!((0, 32, 32, 32), rect(&tile.planes[2]));
505
506    let tile = &tile_states[4].rec; // the center tile
507    assert_eq!((64, 64, 64, 64), rect(&tile.planes[0]));
508    assert_eq!((32, 32, 32, 32), rect(&tile.planes[1]));
509    assert_eq!((32, 32, 32, 32), rect(&tile.planes[2]));
510
511    let tile = &tile_states[5].rec; // the middle-right tile
512    assert_eq!((128, 64, 64, 64), rect(&tile.planes[0]));
513    assert_eq!((64, 32, 32, 32), rect(&tile.planes[1]));
514    assert_eq!((64, 32, 32, 32), rect(&tile.planes[2]));
515
516    let tile = &tile_states[6].rec; // the bottom-left tile
517    assert_eq!((0, 128, 64, 64), rect(&tile.planes[0]));
518    assert_eq!((0, 64, 32, 32), rect(&tile.planes[1]));
519    assert_eq!((0, 64, 32, 32), rect(&tile.planes[2]));
520
521    let tile = &tile_states[7].rec; // the bottom-middle tile
522    assert_eq!((64, 128, 64, 64), rect(&tile.planes[0]));
523    assert_eq!((32, 64, 32, 32), rect(&tile.planes[1]));
524    assert_eq!((32, 64, 32, 32), rect(&tile.planes[2]));
525
526    let tile = &tile_states[8].rec; // the bottom-right tile
527    assert_eq!((128, 128, 64, 64), rect(&tile.planes[0]));
528    assert_eq!((64, 64, 32, 32), rect(&tile.planes[1]));
529    assert_eq!((64, 64, 32, 32), rect(&tile.planes[2]));
530  }
531
532  #[inline]
533  const fn b_area(region: &TileBlocksMut<'_>) -> (usize, usize, usize, usize) {
534    (region.x(), region.y(), region.cols(), region.rows())
535  }
536
537  #[test]
538  fn test_tile_blocks_area() {
539    let (fi, mut fs, mut fb, frame_rate) = setup(160, 144);
540
541    // 4x4 tiles requested, will actually get 3x3 tiles
542    let ti = TilingInfo::from_target_tiles(
543      fi.sb_size_log2(),
544      fi.width,
545      fi.height,
546      frame_rate,
547      2,
548      2,
549      false,
550    );
551    let iter = ti.tile_iter_mut(&mut fs, &mut fb);
552    let tbs = iter.map(|ctx| ctx.tb).collect::<Vec<_>>();
553
554    // the FrameBlocks must be split into 9 TileBlocks:
555    //
556    //   16x16 16x16  8x16
557    //   16x16 16x16  8x16
558    //   16x 4 16x4   8x 4
559
560    assert_eq!(9, tbs.len());
561
562    assert_eq!((0, 0, 16, 16), b_area(&tbs[0]));
563    assert_eq!((16, 0, 16, 16), b_area(&tbs[1]));
564    assert_eq!((32, 0, 8, 16), b_area(&tbs[2]));
565
566    assert_eq!((0, 16, 16, 16), b_area(&tbs[3]));
567    assert_eq!((16, 16, 16, 16), b_area(&tbs[4]));
568    assert_eq!((32, 16, 8, 16), b_area(&tbs[5]));
569
570    assert_eq!((0, 32, 16, 4), b_area(&tbs[6]));
571    assert_eq!((16, 32, 16, 4), b_area(&tbs[7]));
572    assert_eq!((32, 32, 8, 4), b_area(&tbs[8]));
573  }
574
575  #[test]
576  fn test_tile_write() {
577    let (fi, mut fs, mut fb, frame_rate) = setup(160, 144);
578
579    {
580      // 4x4 tiles requested, will actually get 3x3 tiles
581      let ti = TilingInfo::from_target_tiles(
582        fi.sb_size_log2(),
583        fi.width,
584        fi.height,
585        frame_rate,
586        2,
587        2,
588        false,
589      );
590      let iter = ti.tile_iter_mut(&mut fs, &mut fb);
591      let mut tile_states = iter.map(|ctx| ctx.ts).collect::<Vec<_>>();
592
593      {
594        // row 12 of Y-plane of the top-left tile
595        let tile_plane = &mut tile_states[0].rec.planes[0];
596        let row = &mut tile_plane[12];
597        assert_eq!(64, row.len());
598        row[35..41].copy_from_slice(&[4, 42, 12, 18, 15, 31]);
599      }
600
601      {
602        // row 8 of U-plane of the middle-right tile
603        let tile_plane = &mut tile_states[5].rec.planes[1];
604        let row = &mut tile_plane[8];
605        assert_eq!(32, row.len());
606        row[..4].copy_from_slice(&[14, 121, 1, 3]);
607      }
608
609      {
610        // row 1 of V-plane of the bottom-middle tile
611        let tile_plane = &mut tile_states[7].rec.planes[2];
612        let row = &mut tile_plane[1];
613        assert_eq!(32, row.len());
614        row[11..16].copy_from_slice(&[6, 5, 2, 11, 8]);
615      }
616    }
617
618    // check that writes on tiles correctly affected the underlying frame
619
620    let plane = &fs.rec.planes[0];
621    let y = plane.cfg.yorigin + 12;
622    let x = plane.cfg.xorigin + 35;
623    let idx = y * plane.cfg.stride + x;
624    assert_eq!(&[4, 42, 12, 18, 15, 31], &plane.data[idx..idx + 6]);
625
626    let plane = &fs.rec.planes[1];
627    let offset = (64, 32); // middle-right tile, chroma plane
628    let y = plane.cfg.yorigin + offset.1 + 8;
629    let x = plane.cfg.xorigin + offset.0;
630    let idx = y * plane.cfg.stride + x;
631    assert_eq!(&[14, 121, 1, 3], &plane.data[idx..idx + 4]);
632
633    let plane = &fs.rec.planes[2];
634    let offset = (32, 64); // bottom-middle tile, chroma plane
635    let y = plane.cfg.yorigin + offset.1 + 1;
636    let x = plane.cfg.xorigin + offset.0 + 11;
637    let idx = y * plane.cfg.stride + x;
638    assert_eq!(&[6, 5, 2, 11, 8], &plane.data[idx..idx + 5]);
639  }
640
641  #[test]
642  fn test_tile_restoration_edges() {
643    let (fi, mut fs, mut fb, frame_rate) = setup(64, 80);
644
645    let ti = TilingInfo::from_target_tiles(
646      fi.sb_size_log2(),
647      fi.width,
648      fi.height,
649      frame_rate,
650      2,
651      2,
652      false,
653    );
654    let iter = ti.tile_iter_mut(&mut fs, &mut fb);
655    let mut tile_states = iter.map(|ctx| ctx.ts).collect::<Vec<_>>();
656
657    assert_eq!(tile_states.len(), 2);
658
659    {
660      let trs = &mut tile_states[0].restoration;
661      let units = &trs.planes[0].units;
662      assert_eq!(units.x(), 0);
663      assert_eq!(units.y(), 0);
664      assert_eq!(units.cols(), 1);
665      assert_eq!(units.rows(), 1);
666    }
667
668    {
669      let trs = &mut tile_states[1].restoration;
670      let units = &trs.planes[0].units;
671      assert_eq!(units.x(), 0);
672      assert_eq!(units.y(), 1);
673      // no units, the tile is too small (less than 1/2 super-block)
674      assert_eq!(units.cols() * units.rows(), 0);
675    }
676  }
677
678  #[test]
679  fn test_tile_restoration_write() {
680    let (fi, mut fs, mut fb, frame_rate) = setup(256, 256);
681
682    {
683      // 2x2 tiles, each one containing 2×2 restoration units (1 super-block per restoration unit)
684      let ti = TilingInfo::from_target_tiles(
685        fi.sb_size_log2(),
686        fi.width,
687        fi.height,
688        frame_rate,
689        1,
690        1,
691        false,
692      );
693      let iter = ti.tile_iter_mut(&mut fs, &mut fb);
694      let mut tile_states = iter.map(|ctx| ctx.ts).collect::<Vec<_>>();
695
696      {
697        // unit (1, 0) of Y-plane of the top-left tile
698        let units = &mut tile_states[0].restoration.planes[0].units;
699        units[0][1].filter =
700          RestorationFilter::Wiener { coeffs: [[1, 2, 3], [4, 5, 6]] };
701      }
702
703      {
704        // unit (0, 1) of U-plane of the bottom-right tile
705        let units = &mut tile_states[3].restoration.planes[1].units;
706        units[1][0].filter =
707          RestorationFilter::Sgrproj { set: 42, xqd: [10, 20] };
708      }
709
710      {
711        // unit (1, 1) of V-plane of the bottom-left tile
712        let units = &mut tile_states[2].restoration.planes[2].units;
713        units[1][1].filter =
714          RestorationFilter::Sgrproj { set: 5, xqd: [1, 2] };
715      }
716    }
717
718    // check that writes on tiles correctly affected the underlying restoration units
719
720    let units = &mut fs.restoration.planes[0].units;
721    assert_eq!(
722      units[0][1].filter,
723      RestorationFilter::Wiener { coeffs: [[1, 2, 3], [4, 5, 6]] }
724    );
725
726    let units = &mut fs.restoration.planes[1].units;
727    assert_eq!(
728      units[3][2].filter,
729      RestorationFilter::Sgrproj { set: 42, xqd: [10, 20] }
730    );
731
732    let units = &mut fs.restoration.planes[2].units;
733    assert_eq!(
734      units[3][1].filter,
735      RestorationFilter::Sgrproj { set: 5, xqd: [1, 2] }
736    );
737  }
738
739  #[test]
740  fn test_tile_motion_vectors_write() {
741    let (fi, mut fs, mut fb, frame_rate) = setup(160, 144);
742
743    {
744      // 4x4 tiles requested, will actually get 3x3 tiles
745      let ti = TilingInfo::from_target_tiles(
746        fi.sb_size_log2(),
747        fi.width,
748        fi.height,
749        frame_rate,
750        2,
751        2,
752        false,
753      );
754      let iter = ti.tile_iter_mut(&mut fs, &mut fb);
755      let mut tile_states = iter.map(|ctx| ctx.ts).collect::<Vec<_>>();
756
757      {
758        // block (8, 5) of the top-left tile (of the first ref frame)
759        let me_stats = &mut tile_states[0].me_stats[0];
760        me_stats[5][8].mv = MotionVector { col: 42, row: 38 };
761        println!("{:?}", me_stats[5][8].mv);
762      }
763
764      {
765        // block (4, 2) of the middle-right tile (of ref frame 2)
766        let me_stats = &mut tile_states[5].me_stats[2];
767        me_stats[2][3].mv = MotionVector { col: 2, row: 14 };
768      }
769    }
770
771    // check that writes on tiled views affected the underlying motion vectors
772
773    let me_stats = &fs.frame_me_stats.read().unwrap()[0];
774    assert_eq!(MotionVector { col: 42, row: 38 }, me_stats[5][8].mv);
775
776    let me_stats = &fs.frame_me_stats.read().unwrap()[2];
777    let mix = (128 >> MI_SIZE_LOG2) + 3;
778    let miy = (64 >> MI_SIZE_LOG2) + 2;
779    assert_eq!(MotionVector { col: 2, row: 14 }, me_stats[miy][mix].mv);
780  }
781
782  #[test]
783  fn test_tile_blocks_write() {
784    let (fi, mut fs, mut fb, frame_rate) = setup(160, 144);
785
786    {
787      // 4x4 tiles requested, will actually get 3x3 tiles
788      let ti = TilingInfo::from_target_tiles(
789        fi.sb_size_log2(),
790        fi.width,
791        fi.height,
792        frame_rate,
793        2,
794        2,
795        false,
796      );
797      let iter = ti.tile_iter_mut(&mut fs, &mut fb);
798      let mut tbs = iter.map(|ctx| ctx.tb).collect::<Vec<_>>();
799
800      {
801        // top-left tile
802        let tb = &mut tbs[0];
803        // block (4, 3)
804        tb[3][4].n4_w = 42;
805        // block (8, 5)
806        tb[5][8].segmentation_idx = 14;
807      }
808
809      {
810        // middle-right tile
811        let tb = &mut tbs[5];
812        // block (0, 1)
813        tb[1][0].n4_h = 11;
814        // block (7, 5)
815        tb[5][7].cdef_index = 3;
816      }
817
818      {
819        // bottom-middle tile
820        let tb = &mut tbs[7];
821        // block (3, 2)
822        tb[2][3].mode = PredictionMode::PAETH_PRED;
823        // block (1, 1)
824        tb[1][1].n4_w = 8;
825      }
826    }
827
828    // check that writes on tiles correctly affected the underlying blocks
829
830    assert_eq!(42, fb[3][4].n4_w);
831    assert_eq!(14, fb[5][8].segmentation_idx);
832
833    assert_eq!(11, fb[17][32].n4_h);
834    assert_eq!(3, fb[21][39].cdef_index);
835
836    assert_eq!(PredictionMode::PAETH_PRED, fb[34][19].mode);
837    assert_eq!(8, fb[33][17].n4_w);
838  }
839
840  #[test]
841  fn tile_log2_overflow() {
842    assert_eq!(TilingInfo::tile_log2(1, usize::MAX), None);
843  }
844
845  #[test]
846  fn from_target_tiles_422() {
847    let sb_size_log2 = 6;
848    let is_422_p = true;
849    let frame_rate = 60.;
850    let sb_size = 1 << sb_size_log2;
851
852    for frame_height in (sb_size..4352).step_by(sb_size) {
853      for tile_rows_log2 in
854        0..=TilingInfo::tile_log2(1, frame_height >> sb_size_log2).unwrap()
855      {
856        for frame_width in (sb_size..7680).step_by(sb_size) {
857          for tile_cols_log2 in
858            0..=TilingInfo::tile_log2(1, frame_width >> sb_size_log2).unwrap()
859          {
860            let ti = TilingInfo::from_target_tiles(
861              sb_size_log2,
862              frame_width,
863              frame_height,
864              frame_rate,
865              tile_cols_log2,
866              tile_rows_log2,
867              is_422_p,
868            );
869            assert_eq!(
870              ti.tile_cols_log2,
871              TilingInfo::tile_log2(1, ti.cols).unwrap()
872            );
873            assert_eq!(
874              ti.tile_rows_log2,
875              TilingInfo::tile_log2(1, ti.rows).unwrap()
876            );
877          }
878        }
879      }
880    }
881  }
882}
rav1e/tiling/tiler.rs

rav1e/tiling/
tiler.rs
