work in progress

Author: mcroomp

src/structs/lepton_decoder.rs

@@ -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,
@@ -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,
@@ -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,
@@ -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;
 
@@ -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;

src/structs/lepton_encoder.rs

@@ -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,
@@ -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,
@@ -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,
@@ -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;
 
@@ -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
@@ -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;

src/structs/model.rs

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

src/structs/probability_tables.rs

@@ -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,
@@ -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>(

src/structs/quantization_tables.rs

@@ -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::*;
 
@@ -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
@@ -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 {
@@ -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 {
@@ -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] {