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 | |
10 | use super::*; |
11 | |
12 | use crate::context::*; |
13 | use crate::encoder::*; |
14 | use crate::frame::*; |
15 | use crate::lrf::{IntegralImageBuffer, SOLVE_IMAGE_SIZE}; |
16 | use crate::mc::MotionVector; |
17 | use crate::me::FrameMEStats; |
18 | use crate::me::WriteGuardMEStats; |
19 | use crate::partition::{RefType, REF_FRAMES}; |
20 | use crate::predict::{InterCompoundBuffers, PredictionMode}; |
21 | use crate::quantize::*; |
22 | use crate::rdo::*; |
23 | use crate::stats::EncoderStats; |
24 | use crate::util::*; |
25 | use std::ops::{Index, IndexMut}; |
26 | use 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)] |
49 | pub 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)] |
77 | pub struct CodedBlockInfo { |
78 | pub luma_mode: PredictionMode, |
79 | pub chroma_mode: PredictionMode, |
80 | pub reference_types: [RefType; 2], |
81 | } |
82 | |
83 | impl 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)] |
94 | pub struct MiTileState { |
95 | mi_width: usize, |
96 | mi_height: usize, |
97 | mi_block_info: Vec<CodedBlockInfo>, |
98 | } |
99 | |
100 | impl 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 | |
110 | impl 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 | |
119 | impl 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 | |
126 | impl<'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 | } |
266 | |