Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

struct Rav1dTaskContext::scratch: Make into zerocopy-based "unions" #969

Merged
merged 2 commits into from
Apr 25, 2024
Merged

struct Rav1dTaskContext::scratch: Make into zerocopy-based "unions" #969

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f048bb1
Make `Rav1dTaskContext::scratch` into a safe "union"
randomPoison Apr 9, 2024
3909136
Bump alignment of scratch structs
randomPoison Apr 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 10 additions & 8 deletions src/decode.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,9 +65,9 @@ use crate::src::internal::Rav1dContextTaskType;
use crate::src::internal::Rav1dDSPContext;
use crate::src::internal::Rav1dFrameData;
use crate::src::internal::Rav1dTaskContext;
use crate::src::internal::Rav1dTaskContext_scratch_pal;
use crate::src::internal::Rav1dTileState;
use crate::src::internal::ScalableMotionParams;
use crate::src::internal::ScratchPal;
use crate::src::internal::TileStateRef;
use crate::src::intra_edge::EdgeFlags;
use crate::src::intra_edge::EdgeIndex;
Expand Down Expand Up @@ -701,7 +701,7 @@ fn order_palette(
unsafe fn read_pal_indices(
ts: &mut Rav1dTileState,
pal_dsp: &Rav1dPalDSPContext,
scratch_pal: &mut Rav1dTaskContext_scratch_pal,
scratch_pal: &mut ScratchPal,
pal_tmp: &mut [u8],
pal_idx: Option<&mut [u8]>, // if None, use pal_tmp instead of pal_idx
b: &Av1Block,
Expand All @@ -718,7 +718,7 @@ unsafe fn read_pal_indices(
let stride = bw4 * 4;
pal_tmp[0] = rav1d_msac_decode_uniform(&mut ts.msac, pal_sz as c_uint) as u8;
let color_map_cdf = &mut ts.cdf.m.color_map[pli][pal_sz - 2];
let Rav1dTaskContext_scratch_pal {
let ScratchPal {
pal_order: order,
pal_ctx: ctx,
} = scratch_pal;
Expand Down Expand Up @@ -1771,6 +1771,7 @@ unsafe fn decode_b(

if b.pal_sz()[0] != 0 {
let mut pal_idx_guard;
let scratch = t.scratch.inter_intra_mut();
let pal_idx = if t.frame_thread.pass != 0 {
let p = t.frame_thread.pass & 1;
let frame_thread = &mut ts.frame_thread[p as usize];
Expand All @@ -1782,13 +1783,13 @@ unsafe fn decode_b(
frame_thread.pal_idx += len;
&mut *pal_idx_guard
} else {
&mut t.scratch.c2rust_unnamed_0.pal_idx_y
&mut scratch.pal_idx_y
};
read_pal_indices(
ts,
&c.pal_dsp,
&mut t.scratch.c2rust_unnamed_0.c2rust_unnamed.c2rust_unnamed,
&mut t.scratch.c2rust_unnamed_0.pal_idx_uv,
scratch.levels_pal.pal_mut(),
&mut scratch.pal_idx_uv,
Some(pal_idx),
b,
false,
Expand All @@ -1804,6 +1805,7 @@ unsafe fn decode_b(

if has_chroma && b.pal_sz()[1] != 0 {
let mut pal_idx_guard;
let scratch = t.scratch.inter_intra_mut();
let pal_idx = if t.frame_thread.pass != 0 {
let p = t.frame_thread.pass & 1;
let frame_thread = &mut ts.frame_thread[p as usize];
Expand All @@ -1820,8 +1822,8 @@ unsafe fn decode_b(
read_pal_indices(
ts,
&c.pal_dsp,
&mut t.scratch.c2rust_unnamed_0.c2rust_unnamed.c2rust_unnamed,
&mut t.scratch.c2rust_unnamed_0.pal_idx_uv,
scratch.levels_pal.pal_mut(),
&mut scratch.pal_idx_uv,
pal_idx,
b,
true,
Expand Down
178 changes: 126 additions & 52 deletions src/internal.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,9 @@ use std::sync::Condvar;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::thread::JoinHandle;
use zerocopy::AsBytes;
use zerocopy::FromBytes;
use zerocopy::FromZeroes;

#[repr(C)]
pub(crate) struct Rav1dDSPContext {
Expand Down Expand Up @@ -930,90 +933,161 @@ impl BitDepthDependentType for AlPal {
type T<BD: BitDepth> = [[[[BD::Pixel; 8]; 3]; 32]; 2]; /* [2 a/l][32 bx/y4][3 plane][8 palette_idx] */
}

#[derive(Clone, Copy)]
#[repr(C)]
pub struct Rav1dTaskContext_scratch_compinter_seg_mask {
#[derive(FromZeroes, FromBytes, AsBytes)]
#[repr(C, align(64))]
pub struct ScratchCompinter {
pub compinter: [[i16; 16384]; 2],
pub seg_mask: [u8; 16384],
}

pub struct Lap;
// Larger of the two between `ScratchCompinter` and `[BD::Pixel; 128 * 32]`.
const SCRATCH_COMPINTER_SIZE: usize = mem::size_of::<ScratchCompinter>();
randomPoison marked this conversation as resolved.
Show resolved Hide resolved

impl BitDepthDependentType for Lap {
type T<BD: BitDepth> = [BD::Pixel; 128 * 32];
}
#[derive(FromZeroes, FromBytes, AsBytes)]
#[repr(C, align(64))]
pub struct ScratchLapInter([u8; SCRATCH_COMPINTER_SIZE]);

#[derive(Clone, Copy)]
#[repr(C)]
pub union Rav1dTaskContext_scratch_lap {
pub lap: BitDepthUnion<Lap>,
pub c2rust_unnamed: Rav1dTaskContext_scratch_compinter_seg_mask,
impl ScratchLapInter {
pub fn lap_mut<BD: BitDepth>(&mut self) -> &mut [BD::Pixel; 128 * 32] {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}
kkysen marked this conversation as resolved.
Show resolved Hide resolved

pub fn inter_mut(&mut self) -> &mut ScratchCompinter {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}
}

const EMU_EDGE_LEN: usize = 320 * (256 + 7);
// stride=192 for non-SVC, or 320 for SVC
pub struct EmuEdge;
#[derive(FromZeroes, FromBytes, AsBytes)]
#[repr(C, align(64))]
pub struct ScratchEmuEdge([u8; EMU_EDGE_LEN * 2]);

impl BitDepthDependentType for EmuEdge {
type T<BD: BitDepth> = [BD::Pixel; 320 * (256 + 7)];
impl ScratchEmuEdge {
pub fn buf_mut<BD: BitDepth>(&mut self) -> &mut [BD::Pixel; EMU_EDGE_LEN] {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}
}

#[derive(Clone, Copy)]
#[derive(FromZeroes, FromBytes, AsBytes)]
#[repr(C)]
pub struct Rav1dTaskContext_scratch_lap_emu_edge {
pub c2rust_unnamed: Rav1dTaskContext_scratch_lap,
pub emu_edge: BitDepthUnion<EmuEdge>,
pub struct ScratchInter {
randomPoison marked this conversation as resolved.
Show resolved Hide resolved
pub lap_inter: ScratchLapInter,
pub emu_edge: ScratchEmuEdge,
}

#[derive(Clone, Copy)]
#[repr(C)]
pub struct Rav1dTaskContext_scratch_pal {
#[derive(FromZeroes, FromBytes, AsBytes)]
#[repr(C, align(64))]
pub struct ScratchPal {
pub pal_order: [[u8; 8]; 64],
pub pal_ctx: [u8; 64],
}

#[derive(Clone, Copy)]
#[repr(C)]
pub union Rav1dTaskContext_scratch_levels_pal {
pub levels: [u8; 1088],
pub c2rust_unnamed: Rav1dTaskContext_scratch_pal,
#[derive(FromZeroes, FromBytes, AsBytes)]
#[repr(C, align(64))]
pub struct ScratchLevelsPal([u8; 1088]);

impl ScratchLevelsPal {
pub fn levels_mut(&mut self) -> &mut [u8; 1088] {
randomPoison marked this conversation as resolved.
Show resolved Hide resolved
&mut self.0
}

pub fn pal_mut(&mut self) -> &mut ScratchPal {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}
}

#[derive(Clone, Copy)]
#[repr(C)]
pub struct InterIntraEdgePalBD<BD: BitDepth> {
pub interintra: [BD::Pixel; 64 * 64],
pub edge: [BD::Pixel; 257],
_align: [BD::AlignPixelX8; 0],
pub pal: [[BD::Pixel; 8]; 3], /* [3 plane][8 palette_idx] */
#[derive(Clone, Copy, FromZeroes, FromBytes, AsBytes)]
#[repr(C, align(64))]
pub struct ScratchInterintraBuf([u8; 64 * 64 * 2]);
randomPoison marked this conversation as resolved.
Show resolved Hide resolved

impl ScratchInterintraBuf {
pub fn buf_mut<BD: BitDepth>(&mut self) -> &mut [BD::Pixel; 64 * 64] {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}
}

pub struct InterIntraEdgePal;
#[derive(Clone, Copy, FromZeroes, FromBytes, AsBytes)]
#[repr(C, align(64))]
pub struct ScratchEdgeBuf([u8; 257 * 2 + 62]); // 257 Pixel elements + 62 padding bytes

impl BitDepthDependentType for InterIntraEdgePal {
type T<BD: BitDepth> = InterIntraEdgePalBD<BD>;
impl ScratchEdgeBuf {
pub fn buf_mut<BD: BitDepth>(&mut self) -> &mut [BD::Pixel; 257] {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}
}

#[derive(Clone, Copy)]
#[derive(Clone, Copy, FromZeroes, FromBytes, AsBytes)]
#[repr(C, align(16))] // Over-aligned for 8bpc (needs to be `align(8)` for 8bpc, `align(16)` for 16bpc).
pub struct ScratchPalBuf([u8; 8 * 3 * 2]); /* [3 plane][8 palette_idx] */

impl ScratchPalBuf {
pub fn buf<BD: BitDepth>(&self) -> &[[BD::Pixel; 8]; 3] {
FromBytes::ref_from_prefix(&self.0).unwrap()
}

pub fn buf_mut<BD: BitDepth>(&mut self) -> &mut [[BD::Pixel; 8]; 3] {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}
}

#[derive(Clone, Copy, FromZeroes, FromBytes, AsBytes)]
#[repr(C)]
pub union Rav1dTaskContext_scratch_ac_txtp_map {
pub ac: [i16; 1024], // intra only
pub txtp_map: [TxfmType; 1024], // inter only
pub struct ScratchInterintraEdgePal {
pub interintra: ScratchInterintraBuf,
pub edge: ScratchEdgeBuf,
pub pal: ScratchPalBuf,

/// For `AsBytes`, so there's no implicit padding.
_padding: [u8; 16],
}

#[derive(Clone, Copy)]
#[derive(Clone, Copy, FromZeroes, FromBytes, AsBytes)]
#[repr(C, align(64))]
pub struct ScratchAcTxtpMap([u8; 1024 * 2]);

impl ScratchAcTxtpMap {
pub fn ac_mut(&mut self) -> &mut [i16; 1024] {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}

pub fn txtp_map(&self) -> &[TxfmType; 1024] {
FromBytes::ref_from_prefix(&self.0).unwrap()
}

pub fn txtp_map_mut(&mut self) -> &mut [TxfmType; 1024] {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}
}

#[derive(FromZeroes, FromBytes, AsBytes)]
#[repr(C)]
pub struct Rav1dTaskContext_scratch_levels_pal_ac_interintra_edge {
pub c2rust_unnamed: Rav1dTaskContext_scratch_levels_pal,
pub ac_txtp_map: Rav1dTaskContext_scratch_ac_txtp_map,
pub struct ScratchInterIntra {
pub levels_pal: ScratchLevelsPal,
pub ac_txtp_map: ScratchAcTxtpMap,
pub pal_idx_y: [u8; 32 * 64],
pub pal_idx_uv: [u8; 64 * 64], // also used as pre-pack scratch buffer
pub interintra_edge_pal: BitDepthUnion<InterIntraEdgePal>,
pub interintra_edge_pal: ScratchInterintraEdgePal,
}

#[repr(C, align(64))]
pub union Rav1dTaskContext_scratch {
pub c2rust_unnamed: Rav1dTaskContext_scratch_lap_emu_edge,
pub c2rust_unnamed_0: Rav1dTaskContext_scratch_levels_pal_ac_interintra_edge,
// Larger of the two between `ScratchInter` and `ScratchInterIntra`.
const SCRATCH_SIZE: usize = mem::size_of::<ScratchInter>();
#[derive(FromZeroes)]
#[repr(C)]
pub struct TaskContextScratch([u8; SCRATCH_SIZE]);

impl TaskContextScratch {
pub fn inter_mut(&mut self) -> &mut ScratchInter {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}

pub fn inter_intra(&self) -> &ScratchInterIntra {
FromBytes::ref_from_prefix(&self.0).unwrap()
}

pub fn inter_intra_mut(&mut self) -> &mut ScratchInterIntra {
FromBytes::mut_from_prefix(&mut self.0).unwrap()
}
}

#[repr(C)]
Expand Down Expand Up @@ -1056,7 +1130,7 @@ pub(crate) struct Rav1dTaskContext {
pub cf: BitDepthUnion<Cf>,
pub al_pal: BitDepthUnion<AlPal>,
pub pal_sz_uv: [[u8; 32]; 2], /* [2 a/l][32 bx4/by4] */
pub scratch: Rav1dTaskContext_scratch,
pub scratch: TaskContextScratch,

pub warpmv: Rav1dWarpedMotionParams,
/// Index into the relevant `Rav1dFrameContext::lf.mask` array.
Expand All @@ -1082,7 +1156,7 @@ impl Rav1dTaskContext {
cf: Default::default(),
al_pal: Default::default(),
pal_sz_uv: Default::default(),
scratch: unsafe { mem::zeroed() },
scratch: FromZeroes::new_zeroed(),
warpmv: Default::default(),
lf_mask: None,
top_pre_cdef_toggle: 0,
Expand Down
Loading
Loading