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}
883