Skip to content

Commit

Permalink
work in progress
Browse files Browse the repository at this point in the history
  • Loading branch information
@mcroomp
mcroomp committed May 27, 2024
1 parent db3fe91 commit 2f8c324
Show file tree
Hide file tree
Showing 5 changed files with 92 additions and 72 deletions.
30 changes: 17 additions & 13 deletions src/structs/lepton_decoder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -475,7 +475,7 @@ fn decode_edge<R: Read, const ALL_PRESENT: bool>(
decode_one_edge::<R, ALL_PRESENT, true>(
model_per_color,
bool_reader,
&curr_horiz_pred.to_array(),
&curr_horiz_pred,
here_mut,
qt,
pt,
Expand All @@ -486,7 +486,7 @@ fn decode_edge<R: Read, const ALL_PRESENT: bool>(
decode_one_edge::<R, ALL_PRESENT, false>(
model_per_color,
bool_reader,
&curr_vert_pred.to_array(),
&curr_vert_pred,
here_mut,
qt,
pt,
Expand All @@ -504,7 +504,7 @@ fn decode_edge<R: Read, const ALL_PRESENT: bool>(
fn decode_one_edge<R: Read, const ALL_PRESENT: bool, const HORIZONTAL: bool>(
model_per_color: &mut ModelPerColor,
bool_reader: &mut VPXBoolReader<R>,
pred: &[i32; 8],
pred: &i32x8,
here_mut: &mut AlignedBlock,
qt: &QuantizationTables,
pt: &ProbabilityTables,
Expand All @@ -516,6 +516,10 @@ fn decode_one_edge<R: Read, const ALL_PRESENT: bool, const HORIZONTAL: bool>(
.read_non_zero_edge_count::<R, HORIZONTAL>(bool_reader, est_eob, num_non_zeros_bin)
.context(here!())?;

if num_non_zeros_edge == 0 {
return Ok(());
}

let delta;
let mut zig15offset;

Expand All @@ -529,28 +533,28 @@ fn decode_one_edge<R: Read, const ALL_PRESENT: bool, const HORIZONTAL: bool>(

let mut coord_tr = delta;

for _lane in 0..7 {
if num_non_zeros_edge == 0 {
break;
}

let (best_prior_sign_index, best_prior_abs) =
pt.calc_coefficient_context8_lak::<ALL_PRESENT, HORIZONTAL>(qt, coord_tr, pred);
let (best_prior_sign, best_prior_abs) =
pt.calc_coefficient_context8_lak::<ALL_PRESENT, HORIZONTAL>(qt, pred);

for lane in 0..7 {
let coef = model_per_color.read_edge_coefficient(
bool_reader,
qt,
zig15offset,
num_non_zeros_edge,
best_prior_sign_index,
best_prior_abs,
best_prior_sign.as_array_ref()[lane + 1],
best_prior_abs.as_array_ref()[lane + 1],
)?;

if coef != 0 {
num_non_zeros_edge -= 1;
here_mut.set_coefficient(coord_tr, coef);
raster[coord_tr as usize] =
i32::from(coef) * i32::from(qt.get_quantization_table_transposed()[coord_tr]);

num_non_zeros_edge -= 1;
if num_non_zeros_edge == 0 {
break;
}
}

coord_tr += delta;
Expand Down
27 changes: 15 additions & 12 deletions src/structs/lepton_encoder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -500,7 +500,7 @@ fn encode_edge<W: Write, const ALL_PRESENT: bool>(
here_tr,
model_per_color,
bool_writer,
&curr_horiz_pred.to_array(),
&curr_horiz_pred,
qt,
pt,
num_non_zeros_bin,
Expand All @@ -512,7 +512,7 @@ fn encode_edge<W: Write, const ALL_PRESENT: bool>(
here_tr,
model_per_color,
bool_writer,
&curr_vert_pred.to_array(),
&curr_vert_pred,
qt,
pt,
num_non_zeros_bin,
Expand All @@ -538,7 +538,7 @@ fn encode_one_edge<W: Write, const ALL_PRESENT: bool, const HORIZONTAL: bool>(
block: &AlignedBlock,
model_per_color: &mut ModelPerColor,
bool_writer: &mut VPXBoolWriter<W>,
pred: &[i32; 8],
pred: &i32x8,
qt: &QuantizationTables,
pt: &ProbabilityTables,
num_non_zeros_bin: u8,
Expand Down Expand Up @@ -573,6 +573,10 @@ fn encode_one_edge<W: Write, const ALL_PRESENT: bool, const HORIZONTAL: bool>(
)
.context(here!())?;

if num_non_zeros_edge == 0 {
return Ok(());
}

let delta;
let mut zig15offset;

Expand All @@ -586,14 +590,10 @@ fn encode_one_edge<W: Write, const ALL_PRESENT: bool, const HORIZONTAL: bool>(

let mut coord_tr = delta;

for _lane in 0..7 {
if num_non_zeros_edge == 0 {
break;
}

let (best_prior_sign_index, best_prior_abs) =
pt.calc_coefficient_context8_lak::<ALL_PRESENT, HORIZONTAL>(qt, coord_tr, pred);
let (best_prior_sign, best_prior_abs) =
pt.calc_coefficient_context8_lak::<ALL_PRESENT, HORIZONTAL>(qt, pred);

for lane in 0..7 {
let coef = block.get_coefficient(coord_tr);

model_per_color
Expand All @@ -603,13 +603,16 @@ fn encode_one_edge<W: Write, const ALL_PRESENT: bool, const HORIZONTAL: bool>(
coef,
zig15offset,
num_non_zeros_edge,
best_prior_sign_index,
best_prior_abs,
best_prior_sign.as_array_ref()[lane + 1],
best_prior_abs.as_array_ref()[lane + 1],
)
.context(here!())?;

if coef != 0 {
num_non_zeros_edge -= 1;
if num_non_zeros_edge == 0 {
break;
}
}

coord_tr += delta;
Expand Down
6 changes: 3 additions & 3 deletions src/structs/model.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use crate::metrics::{ModelComponent, ModelSubComponent};
use crate::structs::branch::Branch;
use default_boxed::DefaultBoxed;

use super::probability_tables::{EdgeSign, ProbabilityTables};
use super::probability_tables::ProbabilityTables;
use super::quantization_tables::QuantizationTables;
use super::vpx_bool_reader::VPXBoolReader;
use super::vpx_bool_writer::VPXBoolWriter;
Expand Down Expand Up @@ -333,7 +333,7 @@ impl ModelPerColor {
qt: &QuantizationTables,
zig15offset: usize,
num_non_zeros_edge: u8,
best_prior_sign_index: EdgeSign,
best_prior_sign_index: i32,
best_prior_abs: u32,
) -> Result<i16> {
let num_non_zeros_edge_bin = usize::from(num_non_zeros_edge) - 1;
Expand Down Expand Up @@ -428,7 +428,7 @@ impl ModelPerColor {
coef: i16,
zig15offset: usize,
num_non_zeros_edge: u8,
best_prior_sign_index: EdgeSign,
best_prior_sign_index: i32,
best_prior_abs: u32,
) -> Result<()> {
let num_non_zeros_edge_bin = usize::from(num_non_zeros_edge) - 1;
Expand Down
64 changes: 35 additions & 29 deletions src/structs/probability_tables.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,11 @@ use crate::structs::quantization_tables::*;
use super::block_based_image::AlignedBlock;
use super::block_context::NeighborData;

use bytemuck::cast;
use wide::i16x8;
use wide::i32x8;

pub enum EdgeSign {
Zero = 0,
Positive = 1,
Negative = 2,
}
use wide::u32x8;
use wide::CmpLt;

pub struct ProbabilityTables {
left_present: bool,
Expand Down Expand Up @@ -214,34 +211,43 @@ impl ProbabilityTables {
pub fn calc_coefficient_context8_lak<const ALL_PRESENT: bool, const HORIZONTAL: bool>(
&self,
qt: &QuantizationTables,
coefficient_tr: usize,
pred: &[i32; 8],
) -> (EdgeSign, u32) {
if !ALL_PRESENT
&& ((HORIZONTAL && !self.above_present) || (!HORIZONTAL && !self.left_present))
{
return (EdgeSign::Zero, 0);
pred: &i32x8,
) -> (i32x8, u32x8) {
#[inline]
fn mul_hi(lhs: u32x8, rhs: u32x8) -> u32x8 {
let a: [u32; 8] = cast(lhs);
let b: [u32; 8] = cast(rhs);
cast([
((u64::from(a[0]) * u64::from(b[0])) >> 32) as u32,
((u64::from(a[1]) * u64::from(b[1])) >> 32) as u32,
((u64::from(a[2]) * u64::from(b[2])) >> 32) as u32,
((u64::from(a[3]) * u64::from(b[3])) >> 32) as u32,
((u64::from(a[4]) * u64::from(b[4])) >> 32) as u32,
((u64::from(a[5]) * u64::from(b[5])) >> 32) as u32,
((u64::from(a[6]) * u64::from(b[6])) >> 32) as u32,
((u64::from(a[7]) * u64::from(b[7])) >> 32) as u32,
])
}

let best_prior: i32 = pred[if HORIZONTAL {
coefficient_tr >> 3
if !ALL_PRESENT
&& ((HORIZONTAL && !self.is_above_present())
|| (!HORIZONTAL && !self.is_left_present()))
{
(i32x8::ZERO, u32x8::ZERO)
} else {
coefficient_tr
}];
let recip = qt.quantization_table_transposed_recip::<HORIZONTAL>();

let best_prior_abs =
qt.div_quantization_table_transposed_recip(best_prior.unsigned_abs(), coefficient_tr);
let pred_abs = pred.unsigned_abs();

let sign = if best_prior_abs == 0 {
EdgeSign::Zero
} else {
if best_prior > 0 {
EdgeSign::Positive
} else {
EdgeSign::Negative
}
};
(sign, best_prior_abs)
let best_prior_abs: u32x8 = mul_hi(pred_abs, recip) >> 12;

let orig_lt_zero = 1 - pred.cmp_lt(i32x8::ZERO);
let best_prior_is_zero: i32x8 = cast(!best_prior_abs.cmp_eq(u32x8::ZERO));

let best_prior_sign = best_prior_is_zero & orig_lt_zero;

(best_prior_sign, best_prior_abs)
}
}

pub fn adv_predict_dc_pix<const ALL_PRESENT: bool>(
Expand Down
37 changes: 22 additions & 15 deletions src/structs/quantization_tables.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
* This software incorporates material from third parties. See NOTICE.txt for details.
*--------------------------------------------------------------------------------------------*/

use wide::u32x8;

use crate::consts::*;
use crate::helpers::*;

Expand All @@ -13,7 +15,9 @@ pub struct QuantizationTables {
quantization_table: [u16; 64],
quantization_table_transposed: [u16; 64],

quantization_table_transposed_recip: [u32; 64],
quantization_table_transposed_recip_vert: u32x8,
quantization_table_transposed_recip_horiz: u32x8,

// Values for discrimination between "regular" and "noise" part of
// edge AC coefficients, used in `read/write_edge_coefficient`.
// Calculated using approximate maximal magnitudes
Expand All @@ -33,7 +37,8 @@ impl QuantizationTables {
quantization_table: [0; 64],
quantization_table_transposed: [0; 64],
min_noise_threshold: [0; 14],
quantization_table_transposed_recip: [0; 64],
quantization_table_transposed_recip_vert: u32x8::default(),
quantization_table_transposed_recip_horiz: u32x8::default(),
};

for pixel_row in 0..8 {
Expand All @@ -44,10 +49,18 @@ impl QuantizationTables {

retval.quantization_table[coord] = q;
retval.quantization_table_transposed[coord_tr] = q;
retval.quantization_table_transposed_recip[coord_tr] = QuantizationTables::recip(q);
}
}

for i in 0..8 {
retval
.quantization_table_transposed_recip_horiz
.as_array_mut()[i] = Self::recip(retval.quantization_table[i]);
retval
.quantization_table_transposed_recip_vert
.as_array_mut()[i] = Self::recip(retval.quantization_table_transposed[i]);
}

for i in 0..14 {
let coord = if i < 7 { i + 1 } else { (i - 6) * 8 };
if retval.quantization_table[coord] < 9 {
Expand Down Expand Up @@ -78,18 +91,12 @@ impl QuantizationTables {
return retval;
}

pub fn div_quantization_table_transposed_recip(&self, v: u32, index: usize) -> u32 {
// no need to round negative towards zero since this is unsigened division
let result = ((u64::from(v) * u64::from(self.quantization_table_transposed_recip[index]))
>> 31 + 13) as u32;

debug_assert_eq!(
result,
(v / (u32::from(self.quantization_table_transposed[index])) / (1 << 13)),
"right shift"
);

result
pub fn quantization_table_transposed_recip<const HORIZONTAL: bool>(&self) -> u32x8 {
if HORIZONTAL {
self.quantization_table_transposed_recip_horiz
} else {
self.quantization_table_transposed_recip_vert
}
}

pub fn get_quantization_table(&self) -> &[u16; 64] {
Expand Down

0 comments on commit 2f8c324

Please sign in to comment.