8000 filmcraft/crates/av1/src/state.rs at main · hmmhmmhm/filmcraft · GitHub
[go: up one dir, main page]
Skip to content

state.rs

Latest commit

 

History

History
294 lines (272 loc) · 10.7 KB

state.rs

File metadata and controls

294 lines (272 loc) · 10.7 KB
Raw
//! Per-frame decoding state shared by the tiles and the post-filters.
//!
//! [`FrameShared`] holds what tiles only read (headers, references, projected motion field);
//! [`TileState`] what a tile writes (reconstruction, mode info, per-block filter parameters,
//! entropy contexts). With one tile the frame's own buffers move into the tile state and back;
//! with several, each tile decodes into private buffers covering its area (tiles never read
//! outside themselves), which are then copied into the frame.
use crate::decoder::RefFrame;
use crate::frame::{FrameBuf, MiInfo};
use crate::header::{FrameHeader, RefInfo, SequenceHeader};
use crate::spec_tables::*;
use std::ops::{Deref, DerefMut};
use std::sync::Arc;
/// Read-only (during tile decoding) frame state.
pub(crate) struct FrameShared {
pub seq: SequenceHeader,
pub fh: FrameHeader,
pub num_planes: usize,
pub ssx: usize,
pub ssy: usize,
pub bit_depth: u32,
pub lr_unit_rows: [usize; 3],
pub lr_unit_cols: [usize; 3],
pub cdef_cols: usize,
/// Reference slots (FrameStore and saved state) used by this frame.
pub refs: [Option<Arc<RefFrame>>; NUM_REF_FRAMES],
pub ref_info: [RefInfo; NUM_REF_FRAMES],
/// MotionFieldMvs[ ref ][ y8 * w8 + x8 ] (7.9); empty when unused.
pub motion_field: [Vec<[i32; 2]>; 8],
/// PrevSegmentIds (mi units).
pub prev_segment_ids: Vec<u8>,
}
impl FrameShared {
pub fn plane_residual_size(&self, subsize: usize, plane: usize) -> usize {
let sx = if plane > 0 { self.ssx } else { 0 };
let sy = if plane > 0 { self.ssy } else { 0 };
SUBSAMPLED_SIZE[subsize][sx][sy] as usize
}
}
/// A byte per 4x4 unit of a plane (LoopfilterTxSizes), possibly for a region only.
#[derive(Clone, Default)]
pub(crate) struct Grid8 {
pub data: Vec<u8>,
pub ox: usize,
pub oy: usize,
pub stride: usize,
pub rows: usize,
}
impl Grid8 {
fn new(ox: usize, oy: usize, stride: usize, rows: usize) -> Grid8 {
Grid8 { data: vec![0; stride * rows], ox, oy, stride, rows }
}
#[inline]
pub fn get(&self, row: usize, col: usize) -> u8 {
self.data[(row - self.oy) * self.stride + col - self.ox]
}
/// Stores outside the grid are dropped (blocks reaching past the allocated margin).
#[inline]
pub fn set(&mut self, row: usize, col: usize, v: u8) {
if col >= self.ox && col < self.ox + self.stride && row >= self.oy && row < self.oy + self.rows {
self.data[(row - self.oy) * self.stride + col - self.ox] = v;
}
}
fn copy_into(&self, dst: &mut Grid8) {
for y in 0..self.rows {
let d = (self.oy + y - dst.oy) * dst.stride + self.ox - dst.ox;
dst.data[d..d + self.stride].copy_from_slice(&self.data[y * self.stride..(y + 1) * self.stride]);
}
}
}
/// Loop restoration parameters per unit.
#[derive(Clone, Default)]
pub(crate) struct LrUnits {
pub lr_type: [Vec<u8>; 3],
pub lr_wiener: [Vec<[[i8; 3]; 2]>; 3],
pub lr_sgr_set: [Vec<u8>; 3],
pub lr_sgr_xqd: [Vec<[i8; 2]>; 3],
}
/// What tile decoding writes.
pub(crate) struct TileState {
pub cur: FrameBuf,
pub mi: MiInfo,
/// cdef_idx per 64x64 block (-1: not coded).
pub cdef: Vec<i8>,
/// LoopfilterTxSizes[ plane ][ row ][ col ] in 4x4 units of each plane.
< 8000 div class="react-code-text react-code-line-contents" style="min-height:auto">
pub lf_tx_size: [Grid8; 3],
pub lr: LrUnits,
/// LR units this tile coded (per plane), for merging.
pub lr_written: [Vec<usize>; 3],
pub above_level: [Vec<u8>; 3],
pub above_dc: [Vec<u8>; 3],
pub above_seg_pred: Vec<u8>,
pub left_level: [Vec<u8>; 3],
pub left_dc: [Vec<u8>; 3],
pub left_seg_pred: Vec<u8>,
/// TileIntraFrameYModeCdf (reset per tile).
pub intra_frame_y_mode_cdf: [[[u16; 14]; 5]; 5],
cdef_cols: usize,
}
impl TileState {
fn with(cur: FrameBuf, mi: MiInfo, cdef: Vec<i8>, lf_tx_size: [Grid8; 3], lr: LrUnits, cdef_cols: usize) -> TileState {
let (mi_cols, mi_rows) = (mi.cols, mi.rows);
TileState {
cur,
mi,
cdef,
lf_tx_size,
lr,
lr_written: Default::default(),
above_level: std::array::from_fn(|_| vec![0; mi_cols + 64]),
above_dc: std::array::from_fn(|_| vec![0; mi_cols + 64]),
above_seg_pred: vec![0; mi_cols + 64],
left_level: std::array::from_fn(|_| vec![0; mi_rows + 64]),
left_dc: std::array::from_fn(|_| vec![0; mi_rows + 64]),
left_seg_pred: vec![0; mi_rows + 64],
intra_frame_y_mode_cdf: DEFAULT_INTRA_FRAME_Y_MODE_CDF,
cdef_cols,
}
}
#[inline]
pub fn cdef_idx(&self, r: usize, c: usize) -> i8 {
self.cdef[(r >> 4) * self.cdef_cols + (c >> 4)]
}
#[inline]
pub fn set_cdef_idx(&mut self, r: usize, c: usize, v: i8) {
let i = (r >> 4) * self.cdef_cols + (c >> 4);
if i < self.cdef.len() {
self.cdef[i] = v;
}
}
#[inline]
pub fn set_lf_tx_size(&mut self, plane: usize, row: usize, col: usize, v: u8) {
self.lf_tx_size[plane].set(row, col, v);
}
}
pub(crate) struct FrameState {
pub sh: FrameShared,
pub cur: FrameBuf,
pub mi: MiInfo,
cdef: Vec<i8>,
lf_tx_size: [Grid8; 3],
pub lr: LrUnits,
}
impl Deref for FrameState {
type Target = FrameShared;
fn deref(&self) -> &FrameShared {
&self.sh
}
}
impl DerefMut for FrameState {
fn deref_mut(&mut self) -> &mut FrameShared {
&mut self.sh
}
}
fn count_units_in_frame(unit_size: usize, frame_size: usize) -> usize {
((frame_size + (unit_size >> 1)) / unit_size).max(1)
}
impl FrameState {
pub fn new(seq: &SequenceHeader, fh: &FrameHeader) -> FrameState {
let c = &seq.color;
let (ssx, ssy) = (c.subsampling_x as usize, c.subsampling_y as usize);
let mi_cols = fh.mi_cols as usize;
let mi_rows = fh.mi_rows as usize;
let cur = FrameBuf::new(fh.frame_width as usize, fh.frame_height as usize, c.num_planes, ssx, ssy, c.bit_depth);
let cdef_cols = mi_cols.div_ceil(16) + 2;
let cdef_rows = mi_rows.div_ceil(16) + 2;
let mut lr_unit_rows = [0; 3];
let mut lr_unit_cols = [0; 3];
let mut lr = LrUnits::default();
for plane in 0..c.num_planes {
if fh.lr.frame_restoration_type[plane] == RESTORE_NONE as u8 {
continue;
}
let sub_x = if plane == 0 { 0 } else { ssx };
let sub_y = if plane == 0 { 0 } else { ssy };
let unit = fh.lr.loop_restoration_size[plane] as usize;
lr_unit_rows[plane] = count_units_in_frame(unit, (fh.frame_height as usize + sub_y) >> sub_y);
lr_unit_cols[plane] = count_units_in_frame(unit, (fh.upscaled_width as usize + sub_x) >> sub_x);
let n = lr_unit_rows[plane] * lr_unit_cols[plane];
lr.lr_type[plane] = vec![RESTORE_NONE as u8; n];
lr.lr_wiener[plane] = vec![[[0; 3]; 2]; n];
lr.lr_sgr_set[plane] = vec![0; n];
lr.lr_sgr_xqd[plane] = vec![[0; 2]; n];
}
FrameState {
sh: FrameShared {
seq: seq.clone(),
fh: fh.clone(),
num_planes: c.num_planes,
ssx,
ssy,
bit_depth: c.bit_depth as u32,
lr_unit_rows,
lr_unit_cols,
cdef_cols,
refs: Default::default(),
ref_info: Default::default(),
motion_field: Default::default(),
prev_segment_ids: vec![0; mi_cols * mi_rows],
},
cur,
mi: MiInfo::new(mi_cols, mi_rows, fh.allow_screen_content_tools),
cdef: vec![-1; cdef_cols * cdef_rows],
lf_tx_size: std::array::from_fn(|_| Grid8::new(0, 0, mi_cols + 32, mi_rows + 32)),
lr,
}
}
#[inline]
pub fn cdef_idx(&self, r: usize, c: usize) -> i8 {
self.cdef[(r >> 4) * self.cdef_cols + (c >> 4)]
}
#[inline]
pub fn lf_tx_size(&self, plane: usize, row: usize, col: usize) -> usize {
self.lf_tx_size[plane].get(row, col) as usize
}
/// Move the frame's buffers into a tile state covering the whole frame.
pub fn take_tile_state(&mut self) -> TileState {
TileState::with(
std::mem::take(&mut self.cur),
std::mem::take(&mut self.mi),
std::mem::take(&mut self.cdef),
std::mem::take(&mut self.lf_tx_size),
std::mem::take(&mut self.lr),
self.cdef_cols,
)
}
/// Return the buffers taken by [`FrameState::take_tile_state`].
pub fn put_tile_state(&mut self, t: TileState) {
self.cur = t.cur;
self.mi = t.mi;
self.cdef = t.cdef;
self.lf_tx_size = t.lf_tx_size;
self.lr = t.lr;
}
/// Private buffers for the tile covering mode info rows [r0, r1) x columns [c0, c1); tiles
/// on the right / bottom frame edge also cover the allocation margins.
pub fn tile_state(&self, r0: usize, r1: usize, c0: usize, c1: usize) -> TileState {
let (mi_cols, mi_rows) = (self.mi.cols, self.mi.rows);
let x1 = if c1 >= mi_cols { usize::MAX } else { c1 * 4 };
let y1 = if r1 >= mi_rows { usize::MAX } else { r1 * 4 };
let cur = self.cur.region_like(c0 * 4, r0 * 4, x1, y1);
let mi = MiInfo::region(mi_cols, mi_rows, c0, r0, c1 - c0, r1 - r0, !self.mi.palette_size[0].is_empty());
let lf = std::array::from_fn(|p| {
let (sx, sy) = if p == 0 { (0, 0) } else { (self.ssx, self.ssy) };
let full = &self.lf_tx_size[p];
let (gx0, gy0) = (c0 >> sx, r0 >> sy);
let gx1 = if c1 >= mi_cols { full.stride } else { c1 >> sx };
let gy1 = if r1 >= mi_rows { full.rows } else { r1 >> sy };
Grid8::new(gx0, gy0, gx1 - gx0, gy1 - gy0)
});
TileState::with(cur, mi, vec![-1; self.cdef.len()], lf, self.lr.clone(), self.cdef_cols)
}
/// Copy a tile's private buffers into the frame.
pub fn merge_tile_state(&mut self, t: &TileState) {
for p in 0..self.num_planes {
t.cur.planes[p].copy_into(&mut self.cur.planes[p]);
t.lf_tx_size[p].copy_into(&mut self.lf_tx_size[p]);
for &u in &t.lr_written[p] {
self.lr.lr_type[p][u] = t.lr.lr_type[p][u];
self.lr.lr_wiener[p][u] = t.lr.lr_wiener[p][u];
self.lr.lr_sgr_set[p][u] = t.lr.lr_sgr_set[p][u];
self.lr.lr_sgr_xqd[p][u] = t.lr.lr_sgr_xqd[p][u];
}
}
t.mi.copy_into(&mut self.mi);
for (d, &s) in self.cdef.iter_mut().zip(&t.cdef) {
if s != -1 {
*d = s;
}
}
}
}
0