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