rav1e/tiling/
tile_state.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::frame::*;
15use crate::lrf::{IntegralImageBuffer, SOLVE_IMAGE_SIZE};
16use crate::mc::MotionVector;
17use crate::me::FrameMEStats;
18use crate::me::WriteGuardMEStats;
19use crate::partition::{RefType, REF_FRAMES};
20use crate::predict::{InterCompoundBuffers, PredictionMode};
21use crate::quantize::*;
22use crate::rdo::*;
23use crate::stats::EncoderStats;
24use crate::util::*;
25use std::ops::{Index, IndexMut};
26use std::sync::Arc;
27
28/// Tiled view of `FrameState`
29///
30/// Contrary to `PlaneRegionMut` and `TileMut`, there is no const version:
31///  - in practice, we don't need it;
32///  - it would require to instantiate a const version of every of its inner
33///    tiled views recursively.
34///
35/// # `TileState` fields
36///
37/// The way the `FrameState` fields are mapped depend on how they are accessed
38/// tile-wise and frame-wise.
39///
40/// Some fields (like `qc`) are only used during tile-encoding, so they are only
41/// stored in `TileState`.
42///
43/// Some other fields (like `input` or `segmentation`) are not written
44/// tile-wise, so they just reference the matching field in `FrameState`.
45///
46/// Some others (like `rec`) are written tile-wise, but must be accessible
47/// frame-wise once the tile views vanish (e.g. for deblocking).
48#[derive(Debug)]
49pub struct TileStateMut<'a, T: Pixel> {
50  pub sbo: PlaneSuperBlockOffset,
51  pub sb_size_log2: usize,
52  pub sb_width: usize,
53  pub sb_height: usize,
54  pub mi_width: usize,
55  pub mi_height: usize,
56  pub width: usize,
57  pub height: usize,
58  pub input: &'a Frame<T>,     // the whole frame
59  pub input_tile: Tile<'a, T>, // the current tile
60  pub input_hres: &'a Plane<T>,
61  pub input_qres: &'a Plane<T>,
62  pub deblock: &'a DeblockState,
63  pub rec: TileMut<'a, T>,
64  pub qc: QuantizationContext,
65  pub segmentation: &'a SegmentationState,
66  pub restoration: TileRestorationStateMut<'a>,
67  pub me_stats: Vec<TileMEStatsMut<'a>>,
68  pub coded_block_info: MiTileState,
69  pub integral_buffer: IntegralImageBuffer,
70  pub inter_compound_buffers: InterCompoundBuffers,
71}
72
73/// Contains information for a coded block that is
74/// useful to persist. For example, the intra edge
75/// filter requires surrounding coded block information.
76#[derive(Debug, Clone, Copy)]
77pub struct CodedBlockInfo {
78  pub luma_mode: PredictionMode,
79  pub chroma_mode: PredictionMode,
80  pub reference_types: [RefType; 2],
81}
82
83impl Default for CodedBlockInfo {
84  fn default() -> Self {
85    CodedBlockInfo {
86      luma_mode: PredictionMode::DC_PRED,
87      chroma_mode: PredictionMode::DC_PRED,
88      reference_types: [RefType::INTRA_FRAME, RefType::NONE_FRAME],
89    }
90  }
91}
92
93#[derive(Debug, Clone)]
94pub struct MiTileState {
95  mi_width: usize,
96  mi_height: usize,
97  mi_block_info: Vec<CodedBlockInfo>,
98}
99
100impl MiTileState {
101  pub fn new(mi_width: usize, mi_height: usize) -> Self {
102    MiTileState {
103      mi_width,
104      mi_height,
105      mi_block_info: vec![CodedBlockInfo::default(); mi_width * mi_height],
106    }
107  }
108}
109
110impl Index<usize> for MiTileState {
111  type Output = [CodedBlockInfo];
112
113  #[inline(always)]
114  fn index(&self, index: usize) -> &Self::Output {
115    &self.mi_block_info[index * self.mi_width..(index + 1) * self.mi_width]
116  }
117}
118
119impl IndexMut<usize> for MiTileState {
120  #[inline(always)]
121  fn index_mut(&mut self, index: usize) -> &mut Self::Output {
122    &mut self.mi_block_info[index * self.mi_width..(index + 1) * self.mi_width]
123  }
124}
125
126impl<'a, T: Pixel> TileStateMut<'a, T> {
127  pub fn new(
128    fs: &'a mut FrameState<T>, sbo: PlaneSuperBlockOffset,
129    sb_size_log2: usize, width: usize, height: usize,
130    frame_me_stats: &'a mut [FrameMEStats],
131  ) -> Self {
132    debug_assert!(
133      width % MI_SIZE == 0,
134      "Tile width must be a multiple of MI_SIZE"
135    );
136    debug_assert!(
137      height % MI_SIZE == 0,
138      "Tile width must be a multiple of MI_SIZE"
139    );
140
141    let sb_rounded_width = width.align_power_of_two(sb_size_log2);
142    let sb_rounded_height = height.align_power_of_two(sb_size_log2);
143
144    let luma_rect = TileRect {
145      x: sbo.0.x << sb_size_log2,
146      y: sbo.0.y << sb_size_log2,
147      width: sb_rounded_width,
148      height: sb_rounded_height,
149    };
150    let sb_width = width.align_power_of_two_and_shift(sb_size_log2);
151    let sb_height = height.align_power_of_two_and_shift(sb_size_log2);
152
153    Self {
154      sbo,
155      sb_size_log2,
156      sb_width,
157      sb_height,
158      mi_width: width >> MI_SIZE_LOG2,
159      mi_height: height >> MI_SIZE_LOG2,
160      width,
161      height,
162      input: &fs.input,
163      input_tile: Tile::new(&fs.input, luma_rect),
164      input_hres: &fs.input_hres,
165      input_qres: &fs.input_qres,
166      deblock: &fs.deblock,
167      rec: TileMut::new(Arc::make_mut(&mut fs.rec), luma_rect),
168      qc: Default::default(),
169      segmentation: &fs.segmentation,
170      restoration: TileRestorationStateMut::new(
171        &mut fs.restoration,
172        sbo,
173        sb_width,
174        sb_height,
175      ),
176      me_stats: frame_me_stats
177        .iter_mut()
178        .map(|fmvs| {
179          TileMEStatsMut::new(
180            fmvs,
181            sbo.0.x << (sb_size_log2 - MI_SIZE_LOG2),
182            sbo.0.y << (sb_size_log2 - MI_SIZE_LOG2),
183            width >> MI_SIZE_LOG2,
184            height >> MI_SIZE_LOG2,
185          )
186        })
187        .collect(),
188      coded_block_info: MiTileState::new(
189        width >> MI_SIZE_LOG2,
190        height >> MI_SIZE_LOG2,
191      ),
192      integral_buffer: IntegralImageBuffer::zeroed(SOLVE_IMAGE_SIZE),
193      inter_compound_buffers: InterCompoundBuffers::default(),
194    }
195  }
196
197  #[inline(always)]
198  pub fn tile_rect(&self) -> TileRect {
199    TileRect {
200      x: self.sbo.0.x << self.sb_size_log2,
201      y: self.sbo.0.y << self.sb_size_log2,
202      width: self.width,
203      height: self.height,
204    }
205  }
206
207  #[inline(always)]
208  pub fn to_frame_block_offset(
209    &self, tile_bo: TileBlockOffset,
210  ) -> PlaneBlockOffset {
211    let bx = self.sbo.0.x << (self.sb_size_log2 - MI_SIZE_LOG2);
212    let by = self.sbo.0.y << (self.sb_size_log2 - MI_SIZE_LOG2);
213    PlaneBlockOffset(BlockOffset { x: bx + tile_bo.0.x, y: by + tile_bo.0.y })
214  }
215
216  #[inline(always)]
217  pub fn to_frame_super_block_offset(
218    &self, tile_sbo: TileSuperBlockOffset,
219  ) -> PlaneSuperBlockOffset {
220    PlaneSuperBlockOffset(SuperBlockOffset {
221      x: self.sbo.0.x + tile_sbo.0.x,
222      y: self.sbo.0.y + tile_sbo.0.y,
223    })
224  }
225
226  /// Returns above block information for context during prediction.
227  /// If there is no above block, returns `None`.
228  /// `xdec` and `ydec` are the decimation factors of the targeted plane.
229  pub fn above_block_info(
230    &self, bo: TileBlockOffset, xdec: usize, ydec: usize,
231  ) -> Option<CodedBlockInfo> {
232    let (mut bo_x, mut bo_y) = (bo.0.x, bo.0.y);
233    if bo_x & 1 == 0 {
234      bo_x += xdec
235    };
236    if bo_y & 1 == 1 {
237      bo_y -= ydec
238    };
239    if bo_y == 0 {
240      None
241    } else {
242      Some(self.coded_block_info[bo_y - 1][bo_x])
243    }
244  }
245
246  /// Returns left block information for context during prediction.
247  /// If there is no left block, returns `None`.
248  /// `xdec` and `ydec` are the decimation factors of the targeted plane.
249  pub fn left_block_info(
250    &self, bo: TileBlockOffset, xdec: usize, ydec: usize,
251  ) -> Option<CodedBlockInfo> {
252    let (mut bo_x, mut bo_y) = (bo.0.x, bo.0.y);
253    if bo_x & 1 == 1 {
254      bo_x -= xdec
255    };
256    if bo_y & 1 == 0 {
257      bo_y += ydec
258    };
259    if bo_x == 0 {
260      None
261    } else {
262      Some(self.coded_block_info[bo_y][bo_x - 1])
263    }
264  }
265}