refactor(av1): split inverse transform into stages and passes (S3776)

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

crates/av1/src/decoder.rs

@@ -405,35 +405,49 @@ impl Decoder {
         r
     }
 
-    /// Whether an OBU belongs to the selected operating point (spec 7.5: OBUs of other layers are dropped).
-    fn in_operating_point(&self, obu: &ObuUnit) -> bool {
-        if obu.obu_type != OBU_SEQUENCE_HEADER_T
-            && obu.obu_type != OBU_TEMPORAL_DELIMITER_T
-            && obu.ext
-            && let Some(seq) = &self.seq
-            && seq.op_idc != 0
-        {
-            let idc = seq.op_idc;
-            let in_t = (idc >> obu.temporal_id) & 1;
-            let in_s = (idc >> (obu.spatial_id + 8)) & 1;
-            if in_t == 0 || in_s == 0 {
-                return false;
-            }
-        }
-        true
-    }
-
     fn parse_obus(&mut self, data: &[u8], out: &mut Vec<(Output, usize, u32)>) -> Result<()> {
         let mut tu = 0usize;
         let mut pos = 0;
         while pos < data.len() {
-            let obu = parse_obu_unit(data, pos)?;
-            pos = obu.end;
-            if !self.in_operating_point(&obu) {
-                continue;
+            let h = data[pos];
+            if h & 0x80 != 0 {
+                return Err(Error::Invalid("obu_forbidden_bit"));
             }
-            let (payload, spatial_id) = (obu.payload, obu.spatial_id);
-            match obu.obu_type {
+            let obu_type = ((h >> 3) & 0xf) as u32;
+            let ext = h & 4 != 0;
+            let has_size = h & 2 != 0;
+            let mut p = pos + 1;
+            let (mut temporal_id, mut spatial_id) = (0, 0);
+            if ext {
+                let e = *data.get(p).ok_or(Error::Truncated)?;
+                temporal_id = (e >> 5) as u32;
+                spatial_id = ((e >> 3) & 3) as u32;
+                p += 1;
+            }
+            let size = if has_size {
+                let (v, n) = leb128(data.get(p..).ok_or(Error::Truncated)?)?;
+                p += n;
+                v as usize
+            } else {
+                data.len().checked_sub(p).ok_or(Error::Truncated)?
+            };
+            let end = p.checked_add(size).filter(|&e| e <= data.len()).ok_or(Error::Truncated)?;
+            let payload = &data[p..end];
+            pos = end;
+            if obu_type != OBU_SEQUENCE_HEADER_T
+                && obu_type != OBU_TEMPORAL_DELIMITER_T
+                && ext
+                && let Some(seq) = &self.seq
+                && seq.op_idc != 0
+            {
+                let idc = seq.op_idc;
+                let in_t = (idc >> temporal_id) & 1;
+                let in_s = (idc >> (spatial_id + 8)) & 1;
+                if in_t == 0 || in_s == 0 {
+                    continue;
+                }
+            }
+            match obu_type {
                 OBU_SEQUENCE_HEADER_T => {
                     let s = SequenceHeader::parse(payload)?;
                     self.seq = Some(s);
@@ -443,8 +457,24 @@ impl Decoder {
                     tu += 1;
                 }
                 OBU_FRAME_HEADER_T | OBU_REDUNDANT_FRAME_HEADER_T | OBU_FRAME_T => {
-                    if let Some(o) = self.frame_obu(&obu)? {
+                    if self.seen_frame_header {
+                        // frame_header_copy(): identical to the active header
+                        if obu_type != OBU_FRAME_T {
+                            continue;
+                        }
+                    }
+                    let mut r = BitReader::new(payload);
+                    let shown = self.frame_header(&mut r, temporal_id, spatial_id)?;
+                    if let Some(o) = shown {
                         out.push((o, tu, spatial_id));
+                        continue;
+                    }
+                    if obu_type == OBU_FRAME_T {
+                        r.byte_align();
+                        let rest = payload.get(r.byte_pos()..).ok_or(Error::Truncated)?;
+                        if let Some(o) = self.tile_group(rest)? {
+                            out.push((o, tu, spatial_id));
+                        }
                     }
                 }
                 OBU_TILE_GROUP_T => {
@@ -459,22 +489,6 @@ impl Decoder {
         Ok(())
     }
 
-    /// A frame header or frame OBU: returns the picture it completes, if any.
-    fn frame_obu(&mut self, obu: &ObuUnit) -> Result<Option<Output>> {
-        // frame_header_copy(): identical to the active header
-        if self.seen_frame_header && obu.obu_type != OBU_FRAME_T {
-            return Ok(None);
-        }
-        let mut r = BitReader::new(obu.payload);
-        let shown = self.frame_header(&mut r, obu.temporal_id, obu.spatial_id)?;
-        if shown.is_some() || obu.obu_type != OBU_FRAME_T {
-            return Ok(shown);
-        }
-        r.byte_align();
-        let rest = obu.payload.get(r.byte_pos()..).ok_or(Error::Truncated)?;
-        self.tile_group(rest)
-    }
-
     /// frame_header_obu( ): returns the output for show_existing_frame.
     fn frame_header(&mut self, r: &mut BitReader, temporal_id: u32, spatial_id: u32) -> Result<Option<Output>> {
         let seq = self.seq.clone().ok_or(Error::Invalid("frame before sequence header"))?;
@@ -1000,43 +1014,6 @@ mod tests {
     }
 }
 
-/// One OBU: its header fields, payload and the position just past it.
-struct ObuUnit<'a> {
-    obu_type: u32,
-    ext: bool,
-    temporal_id: u32,
-    spatial_id: u32,
-    payload: &'a [u8],
-    end: usize,
-}
-
-fn parse_obu_unit(data: &[u8], pos: usize) -> Result<ObuUnit<'_>> {
-    let h = data[pos];
-    if h & 0x80 != 0 {
-        return Err(Error::Invalid("obu_forbidden_bit"));
-    }
-    let obu_type = ((h >> 3) & 0xf) as u32;
-    let ext = h & 4 != 0;
-    let has_size = h & 2 != 0;
-    let mut p = pos + 1;
-    let (mut temporal_id, mut spatial_id) = (0, 0);
-    if ext {
-        let e = *data.get(p).ok_or(Error::Truncated)?;
-        temporal_id = (e >> 5) as u32;
-        spatial_id = ((e >> 3) & 3) as u32;
-        p += 1;
-    }
-    let size = if has_size {
-        let (v, n) = leb128(data.get(p..).ok_or(Error::Truncated)?)?;
-        p += n;
-        v as usize
-    } else {
-        data.len().checked_sub(p).ok_or(Error::Truncated)?
-    };
-    let end = p.checked_add(size).filter(|&e| e <= data.len()).ok_or(Error::Truncated)?;
-    Ok(ObuUnit { obu_type, ext, temporal_id, spatial_id, payload: &data[p..end], end })
-}
-
 #[cfg(test)]
 mod guard_tests {
     use super::*;

crates/av1/src/transform.rs

@@ -165,6 +165,15 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
     for i in 0..len {
         t[i] = copy[brev(n, i)];
     }
+    inverse_dct_stage1(t, n, r);
+    inverse_dct_stage2(t, n, r);
+    inverse_dct_stage3(t, n, r);
+    inverse_dct_stage4(t, n, r);
+    inverse_dct_stage5(t, n, r);
+    inverse_dct_stage6(t, n, r);
+}
+
+fn inverse_dct_stage1<T: Lane>(t: &mut [T], n: u32, r: u32) {
     if n == 6 {
         for i in 0..16 {
             bf(t, 32 + i, 63 - i, 63 - 4 * brev(4, i) as i32, false);
@@ -190,6 +199,9 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
             hd(t, 16 + 2 * i, 17 + 2 * i, i & 1 == 1, r);
         }
     }
+}
+
+fn inverse_dct_stage2<T: Lane>(t: &mut [T], n: u32, r: u32) {
     if n == 6 {
         for i in 0..4 {
             for j in 0..2 {
@@ -221,6 +233,9 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
             }
         }
     }
+}
+
+fn inverse_dct_stage3<T: Lane>(t: &mut [T], n: u32, r: u32) {
     for i in 0..2 {
         bf(t, 2 * i, 2 * i + 1, 32 + 16 * i as i32, i == 0);
     }
@@ -248,6 +263,9 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
             }
         }
     }
+}
+
+fn inverse_dct_stage4<T: Lane>(t: &mut [T], n: u32, r: u32) {
     for i in 0..2 {
         hd(t, i, 3 - i, false, r);
     }
@@ -273,6 +291,9 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
             }
         }
     }
+}
+
+fn inverse_dct_stage5<T: Lane>(t: &mut [T], n: u32, r: u32) {
     if n >= 3 {
         for i in 0..4 {
             hd(t, i, 7 - i, false, r);
@@ -300,6 +321,9 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
             hd(t, i, 15 - i, false, r);
         }
     }
+}
+
+fn inverse_dct_stage6<T: Lane>(t: &mut [T], n: u32, r: u32) {
     if n >= 5 {
         for i in 0..4 {
             bf(t, 27 - i, 20 + i, 32, true);
@@ -531,10 +555,9 @@ pub(crate) fn inverse_transform_2d(dequant: &[i32], residual: &mut [i32], tx_sz:
     let col_clamp = (bit_depth + 6).max(16);
     let rk = row_kind(tx_type);
     let ck = col_kind(tx_type);
-    let mut t = [0i32; 64];
     let rect2 = log2w.abs_diff(log2h) == 1;
+    let p = Tx2d { w, h, log2w, log2h, row_shift, col_shift, row_clamp, col_clamp, rk, ck, rect2 };
     if !lossless && bit_depth <= 10 {
-        let p = Tx2d { w, h, log2w, log2h, row_shift, col_shift, row_clamp, col_clamp, rk, ck, rect2 };
         if w >= 8 && h >= 8 {
             transform_lanes::<8>(dequant, residual, &p);
         } else {
@@ -542,6 +565,13 @@ pub(crate) fn inverse_transform_2d(dequant: &[i32], residual: &mut [i32], tx_sz:
         }
         return;
     }
+    scalar_row_pass(dequant, residual, &p, lossless);
+    scalar_col_pass(residual, &p, lossless);
+}
+
+fn scalar_row_pass(dequant: &[i32], residual: &mut [i32], p: &Tx2d, lossless: bool) {
+    let (w, h, log2w, row_shift, row_clamp, col_clamp, rk, rect2) = (p.w, p.h, p.log2w, p.row_shift, p.row_clamp, p.col_clamp, p.rk, p.rect2);
+    let mut t = [0i32; 64];
     let in_w = w.min(32);
     for i in 0..h {
         // Every 1D transform maps an all-zero row to zeros (each step is linear with
@@ -572,6 +602,11 @@ pub(crate) fn inverse_transform_2d(dequant: &[i32], residual: &mut [i32], tx_sz:
             residual[i * w + j] = round2(t[j], row_shift).clamp(lo, hi);
         }
     }
+}
+
+fn scalar_col_pass(residual: &mut [i32], p: &Tx2d, lossless: bool) {
+    let (w, h, log2h, col_shift, col_clamp, ck) = (p.w, p.h, p.log2h, p.col_shift, p.col_clamp, p.ck);
+    let mut t = [0i32; 64];
     for j in 0..w {
         for i in 0..h {
             t[i] = residual[i * w + j];
@@ -617,11 +652,16 @@ fn transform_1d<T: Lane>(t: &mut [T], kind: Kind1d, n: u32, r: u32) {
 /// The 2D inverse transform with `L` rows (row pass) or columns (column pass) at a time;
 /// bit-identical to the one-at-a-time process for bit depths up to 10.
 fn transform_lanes<const L: usize>(dequant: &[i32], residual: &mut [i32], p: &Tx2d) {
+    let mut t = [[0i32; L]; 64];
+    lanes_row_pass(dequant, residual, p, &mut t);
+    lanes_col_pass(residual, p, &mut t);
+}
+
+fn lanes_row_pass<const L: usize>(dequant: &[i32], residual: &mut [i32], p: &Tx2d, t: &mut [[i32; L]; 64]) {
     let (w, h) = (p.w, p.h);
     let in_w = w.min(32);
     let lo = -(1i32 << (p.col_clamp - 1));
     let hi = (1i32 << (p.col_clamp - 1)) - 1;
-    let mut t = [[0i32; L]; 64];
     let mut i0 = 0;
     while i0 < h {
         // Rows without coefficients (and rows 32..63) transform to zeros.
@@ -649,6 +689,10 @@ fn transform_lanes<const L: usize>(dequant: &[i32], residual: &mut [i32], p: &Tx
         }
         i0 += L;
     }
+}
+
+fn lanes_col_pass<const L: usize>(residual: &mut [i32], p: &Tx2d, t: &mut [[i32; L]; 64]) {
+    let (w, h) = (p.w, p.h);
     let mut j0 = 0;
     while j0 < w {
         for i in 0..h {