M13.4: av1enc frame-to-frame CDF carry-over, intra mode and transform type RD

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

crates/av1enc/README.md

@@ -22,8 +22,8 @@ reconstruction, so prediction never drifts.
   `Encoder::av1c()` returns the `AV1CodecConfigurationRecord` (also the WebM `CodecPrivate`).
 - **Frames**: a `KEY_FRAME` every `keyint` frames, `INTER_FRAME`s otherwise. Inter frames use
   slot 0 (the previous reconstruction) for all seven references and refresh slot 0 only, so
-  `LAST_FRAME` is always the previous frame. `primary_ref_frame = PRIMARY_REF_NONE` (default
-  CDFs every frame), symbol-level CDF adaptation on, frame-end CDF update off. One tile,
+  `LAST_FRAME` is always the previous frame. CDFs adapt per symbol and are saved at the end of
+  every frame; inter frames load them from `LAST_FRAME` (`primary_ref_frame` = 0). One tile,
   `TX_MODE_LARGEST`, `reduced_tx_set`, no segmentation, no delta q / lf.
 - **Deblocking**: the normative loop filter (7.14) runs on the reconstruction; the frame level
   (the same for both luma directions and both chroma planes, no deltas, sharpness 0) is searched
@@ -40,8 +40,9 @@ reconstruction, so prediction never drifts.
   from neighbour / co-located / parent candidates, a whole-sample diamond search and half- then
   quarter-sample refinement (`allow_high_precision_mv` = 0); the normative 8-tap `EIGHTTAP`
   prediction for luma and chroma; skip blocks; per-block intra / inter choice.
-- **Residual**: one transform per plane and block (4x4 ... 32x32): `DCT_DCT` for luma, the
-  mode-dependent `DCT` / `ADST` combination for intra chroma; dead-zone quantisation with the
+- **Residual**: one transform per plane and block (4x4 ... 32x32): for intra luma `DCT_DCT` or
+  the mode's preferred `ADST` combination (rate-distortion choice over the two best modes), the
+  mode-dependent type for intra chroma, `DCT_DCT` for inter blocks; dead-zone quantisation with the
   spec's `Dc_Qlookup` / `Ac_Qlookup`; the normative dequantisation and inverse transforms in the
   reconstruction loop; full coefficient coding (all_zero, tx type, eob_pt, eob_extra,
   coeff_base_eob, coeff_base, coeff_br, dc_sign, Golomb) with the exact contexts.

crates/av1enc/src/cdf.rs

@@ -95,3 +95,70 @@ impl Cdfs {
         })
     }
 }
+
+impl Cdfs {
+    /// The counter reset of load_cdfs( ): the last entry of every CDF array is set to 0.
+    pub fn clear_counts(&mut self) {
+        macro_rules! clr {
+            ($($f:ident),*) => { $( self.$f.clear_count(); )* };
+        }
+        clr!(
+            y_mode,
+            uv_mode_cfl_not_allowed,
+            uv_mode_cfl_allowed,
+            angle_delta,
+            partition_w8,
+            partition_w16,
+            partition_w32,
+            partition_w64,
+            skip,
+            is_inter,
+            single_ref,
+            new_mv,
+            zero_mv,
+            ref_mv,
+            drl_mode,
+            mv_joint,
+            mv_class,
+            mv_class0_bit,
+            mv_fr,
+            mv_class0_fr,
+            mv_sign,
+            mv_bit,
+            intra_tx_type_set2,
+            inter_tx_type_set3,
+            txb_skip,
+            eob_pt_16,
+            eob_pt_32,
+            eob_pt_64,
+            eob_pt_128,
+            eob_pt_256,
+            eob_pt_512,
+            eob_pt_1024,
+            eob_extra,
+            dc_sign,
+            coeff_base_eob,
+            coeff_base,
+            coeff_br
+        );
+    }
+}
+
+/// Zeroes the adaptation counter (last element) of every innermost CDF array.
+trait ClearCount {
+    fn clear_count(&mut self);
+}
+
+impl<const N: usize> ClearCount for [u16; N] {
+    fn clear_count(&mut self) {
+        self[N - 1] = 0;
+    }
+}
+
+impl<T: ClearCount, const M: usize> ClearCount for [T; M] {
+    fn clear_count(&mut self) {
+        for e in self.iter_mut() {
+            e.clear_count();
+        }
+    }
+}

crates/av1enc/src/decide.rs

@@ -35,6 +35,8 @@ pub(crate) struct RefFrame {
     pub mvs: Vec<[i32; 2]>,
     /// The frame's loop filter level.
     pub lf_level: u32,
+    /// The CDFs saved at the end of the frame (loaded by the next frame).
+    pub cdfs: Box<crate::cdf::Cdfs>,
 }
 
 pub(crate) struct FrameGeom {
@@ -339,19 +341,45 @@ impl<'a> Decider<'a> {
             bit_depth: bd,
         };
         let ed = edges(&self.rec[0], &e);
-        let mut best = (f64::MAX, DC_PRED);
-        for &mode in INTRA_MODES_USED.iter() {
+        // the two best modes by SAD, then a rate-distortion choice of mode and transform type
+        let mut ranked: Vec<(f64, usize)> = INTRA_MODES_USED
+            .iter()
+            .map(|&mode| {
+                predict_intra(&ed, &e, mode, &mut pred);
+                let sad = self.sad(0, c * 4, r * 4, n, &pred[..n * n]) as f64;
+                (sad + self.lambda_sad * if mode == DC_PRED { 1.0 } else { MODE_BITS_INTRA }, mode)
+            })
+            .collect();
+        ranked.sort_by(|a, b| a.0.total_cmp(&b.0));
+        let tx_y = tx_size_for(log2);
+        let mut best: Option<(f64, usize, usize, TxResult, Option<Vec<u16>>)> = None;
+        for &(_, mode) in ranked.iter().take(2) {
             predict_intra(&ed, &e, mode, &mut pred);
-            let sad = self.sad(0, c * 4, r * 4, n, &pred[..n * n]) as f64;
-            let cost = sad + self.lambda_sad * if mode == DC_PRED { 1.0 } else { MODE_BITS_INTRA };
-            if cost < best.0 {
-                best = (cost, mode);
-                best_pred[..n * n].copy_from_slice(&pred[..n * n]);
+            let mut types = vec![DCT_DCT];
+            let alt = MODE_TO_TXFM[mode] as usize;
+            if alt != DCT_DCT && tx_in_set(tx_set(tx_y, false), false, alt) {
+                types.push(alt);
+            }
+            for &t in &types {
+                let (res, recon) = self.try_tx(0, c * 4, r * 4, log2, t, &pred[..n * n], true);
+                let type_bits = if !res.nonzero || tx_set(tx_y, false) == TX_SET_DCTONLY {
+                    0.0
+                } else if t == DCT_DCT {
+                    1.0
+                } else {
+                    2.0
+                };
+                let mode_bits = if mode == DC_PRED { 1.0 } else { MODE_BITS_INTRA };
+                let cost = res.dist + self.lambda * (res.bits + type_bits + mode_bits);
+                if best.as_ref().is_none_or(|b| cost < b.0) {
+                    best_pred[..n * n].copy_from_slice(&pred[..n * n]);
+                    best = Some((cost, mode, t, res, recon));
+                }
             }
         }
-        let y_mode = best.1;
-        let tx_y = tx_size_for(log2);
-        let ty = self.code_tx(0, c * 4, r * 4, log2, DCT_DCT, &best_pred[..n * n], true);
+        let (_, y_mode, tx_type_y, ty, recon) = best.expect("an intra mode");
+        self.put(0, c * 4, r * 4, n, recon.as_deref().unwrap_or(&best_pred[..n * n]));
+        let tx_type_y = if ty.nonzero { tx_type_y } else { DCT_DCT };
         self.mark_decoded(0, r, c, n4);
         // chroma (one mode for both planes)
         let cn = n >> 1;
@@ -391,12 +419,11 @@ impl<'a> Decider<'a> {
             tuv.push(self.code_tx(1 + i, c * 2, r * 2, clog2, uv_type, &pred[..cn * cn], true));
             self.mark_decoded(1 + i, r, c, n4);
         }
-        let _ = tx_y;
         let skip = !ty.nonzero && !tuv[0].nonzero && !tuv[1].nonzero;
         let dist = ty.dist + tuv[0].dist + tuv[1].dist;
         let bits = if skip { 1.0 } else { ty.bits + tuv[0].bits + tuv[1].bits + 1.0 } + MODE_BITS_INTRA + 2.0;
         let [u, v]: [TxResult; 2] = tuv.try_into().ok().expect("two chroma planes");
-        let leaf = Leaf { bsize, is_inter: false, y_mode, uv_mode, tx_type_y: DCT_DCT, mv: [0, 0], skip, coefs: [ty.levels, u.levels, v.levels] };
+        let leaf = Leaf { bsize, is_inter: false, y_mode, uv_mode, tx_type_y, mv: [0, 0], skip, coefs: [ty.levels, u.levels, v.levels] };
         (leaf, dist + self.lambda * bits)
     }
 
@@ -688,6 +715,15 @@ impl<'a> Decider<'a> {
     /// Transform, quantise and reconstruct one square transform block of `plane` at (x, y)
     /// over the prediction `pred`; the reconstruction is written to `rec`.
     fn code_tx(&mut self, p: usize, x: usize, y: usize, log2: u32, tx_type: usize, pred: &[u16], intra: bool) -> TxResult {
+        let (res, recon) = self.try_tx(p, x, y, log2, tx_type, pred, intra);
+        let n = 1usize << log2;
+        self.put(p, x, y, n, recon.as_deref().unwrap_or(pred));
+        res
+    }
+
+    /// Transform and quantise one block; returns the result and its reconstruction (None: the
+    /// prediction) without writing it.
+    fn try_tx(&self, p: usize, x: usize, y: usize, log2: u32, tx_type: usize, pred: &[u16], intra: bool) -> (TxResult, Option<Vec<u16>>) {
         let n = 1usize << log2;
         let tx_sz = tx_size_for(log2);
         let bd = self.bd;
@@ -715,19 +751,16 @@ impl<'a> Decider<'a> {
         }
         let zero_dist = self.sse(p, x, y, n, pred);
         if !nonzero {
-            self.put(p, x, y, n, pred);
-            return TxResult { levels, nonzero, dist: zero_dist, bits: 1.0 };
+            return (TxResult { levels, nonzero, dist: zero_dist, bits: 1.0 }, None);
         }
         let recon = self.reconstruct(&levels, tx_sz, tx_type, pred, n, bd);
         let dist = self.sse(p, x, y, n, &recon);
         let bits = estimate_bits(&levels, tx_sz);
         if dist + self.lambda * bits >= zero_dist + self.lambda {
             levels.iter_mut().for_each(|l| *l = 0);
-            self.put(p, x, y, n, pred);
-            return TxResult { levels, nonzero: false, dist: zero_dist, bits: 1.0 };
+            return (TxResult { levels, nonzero: false, dist: zero_dist, bits: 1.0 }, None);
         }
-        self.put(p, x, y, n, &recon);
-        TxResult { levels, nonzero: true, dist, bits }
+        (TxResult { levels, nonzero: true, dist, bits }, Some(recon))
     }
 
     /// Dequantisation (7.12.3) and the inverse transform, added to the prediction.

crates/av1enc/src/lib.rs

@@ -275,7 +275,7 @@ impl Encoder {
         w.f(1, 0); // disable_cdf_update
         w.f(1, 0); // frame_size_override_flag
         if !key {
-            w.f(3, 7); // primary_ref_frame = PRIMARY_REF_NONE
+            w.f(3, 0); // primary_ref_frame: LAST_FRAME's saved CDFs
             w.f(8, 1); // refresh_frame_flags: slot 0 becomes LAST for the next frame
             for _ in 0..7 {
                 w.f(3, 0); // ref_frame_idx[ i ]
@@ -288,7 +288,7 @@ impl Encoder {
         } else {
             w.f(1, 0); // render_and_frame_size_different
         }
-        w.f(1, 1); // disable_frame_end_update_cdf
+        w.f(1, 0); // disable_frame_end_update_cdf
         // tile_info( ): uniform spacing, one tile
         let g = &self.geom;
         let tile_log2 = |blk: usize, target: usize| {
@@ -364,7 +364,7 @@ impl Encoder {
         let g = &self.geom;
         let refr = if key { None } else { self.refr.as_ref() };
         let mut dec = Decider::new(g, self.cfg.bit_depth as u32, qidx, key, src, refr);
-        let mut tw = TileWriter::new(g, key, qidx);
+        let mut tw = TileWriter::new(g, key, qidx, refr.map(|r| &*r.cdfs));
         for sbr in 0..g.sb_rows {
             tw.start_sb_row();
             for sbc in 0..g.sb_cols {
@@ -374,7 +374,7 @@ impl Encoder {
         }
         let lf_level = self.pick_filter_level(&dec.rec, &tw.lf_info(), src, qidx);
         loop_filter(&mut dec.rec, &tw.lf_info(), [lf_level; 4], self.cfg.bit_depth as u32, false);
-        let tile = tw.finish();
+        let (tile, cdfs) = tw.finish();
         let mut w = BitWriter::new();
         self.frame_header(&mut w, key, qidx, lf_level);
         w.byte_align();
@@ -383,7 +383,7 @@ impl Encoder {
         payload.extend_from_slice(&tile);
         let mut out = Vec::with_capacity(payload.len() + 8);
         obu(OBU_FRAME, &payload, &mut out);
-        let rf = RefFrame { planes: dec.rec, mvs: dec.mvs, lf_level };
+        let rf = RefFrame { planes: dec.rec, mvs: dec.mvs, lf_level, cdfs };
         (out, rf)
     }
 

crates/av1enc/src/writer.rs

@@ -80,14 +80,23 @@ fn is_directional(mode: usize) -> bool {
 }
 
 impl<'a> TileWriter<'a> {
-    pub fn new(g: &'a FrameGeom, intra_only: bool, qidx: u8) -> Self {
+    /// `init`: the CDFs saved by the primary reference frame (load_cdfs), or None for the
+    /// defaults (init_non_coeff_cdfs / init_coeff_cdfs).
+    pub fn new(g: &'a FrameGeom, intra_only: bool, qidx: u8, init: Option<&Cdfs>) -> Self {
         let n = g.mi_cols * g.mi_rows;
         TileWriter {
             g,
             intra_only,
             qidx,
             sw: SymbolWriter::new(false),
-            cdf: Cdfs::new(qidx as u32),
+            cdf: match init {
+                Some(c) => {
+                    let mut c = Box::new(c.clone());
+                    c.clear_counts();
+                    c
+                }
+                None => Cdfs::new(qidx as u32),
+            },
             kf_y_mode: DEFAULT_INTRA_FRAME_Y_MODE_CDF,
             mi: Mi {
                 cols: g.mi_cols,
@@ -114,8 +123,9 @@ impl<'a> TileWriter<'a> {
         }
     }
 
-    pub fn finish(self) -> Vec<u8> {
-        self.sw.finish()
+    /// The tile data and the final CDFs (saved for the frame-end CDF update).
+    pub fn finish(self) -> (Vec<u8>, Box<Cdfs>) {
+        (self.sw.finish(), self.cdf)
     }
 
     /// The mode info the loop filter needs.

crates/av1enc/tests/ffmpeg_oracle.rs

@@ -117,6 +117,11 @@ fn libdav1d_decodes_bit_exact() {
     // all-intra at a fine quantizer (large coefficients, Golomb codes)
     let pics: Vec<Pic> = (0..2).map(|t| noise(64, 48, 8, t as u32 + 3)).collect();
     check("noise-intra", cfg(64, 48, 8, RateControl::ConstantQ(4), 1), &pics);
+    // the finest quantizer: the largest coefficients, 8 and 10 bit, intra and inter
+    let pics: Vec<Pic> = (0..3).map(|t| noise(72, 40, 8, t as u32 + 11)).collect();
+    check("noise-q1", cfg(72, 40, 8, RateControl::ConstantQ(1), 2), &pics);
+    let pics: Vec<Pic> = (0..3).map(|t| noise(40, 72, 10, t as u32 + 21)).collect();
+    check("noise-q1-10", cfg(40, 72, 10, RateControl::ConstantQ(1), 2), &pics);
     // odd sizes, 8 and 10 bit
     let pics: Vec<Pic> = (0..4).map(|t| synth(33, 17, 8, t)).collect();
     check("odd8", cfg(33, 17, 8, RateControl::ConstantQ(60), 3), &pics);