From 5afd7205f0d71e4d79c1845f4d6d9c6333ff1a40 Mon Sep 17 00:00:00 2001 From: "Dylan M. Taylor" Date: Sat, 24 Oct 2026 08:53:14 +0000 Subject: [PATCH] refactor(av1): reduce complexity in cdef and loop restoration (S3776) Co-Authored-By: Claude Sonnet 5.5 --- crates/av1/src/cdef.rs | 189 ++++++++++++++++++++------------- crates/av1/src/restoration.rs | 191 ++++++++++++++++++++-------------- 2 files changed, 228 insertions(+), 152 deletions(-) diff --git a/crates/av1/src/cdef.rs b/crates/av1/src/cdef.rs index 0a919351..3ea3ee9b 100644 --- a/crates/av1/src/cdef.rs +++ b/crates/av1/src/cdef.rs @@ -45,7 +45,7 @@ fn cdef_block(fs: &FrameState, out: &mut FrameBuf, r: usize, c: usize, idx: usiz let var_str = if (var >> 6) != 0 { (31 - ((var >> 6) as u32).leading_zeros()).min(12) as i32 } else { 0 }; pri = if var != 0 { (pri * (4 + var_str) + 8) >> 4 } else { 0 }; let damping = cp.damping as i32 + coeff_shift as i32; - filter(fs, out, 0, r, c, pri, sec, damping, dir); + filter(fs, out, 0, (r, c), (pri, sec), damping, dir); if fs.num_planes == 1 { return true; } @@ -53,8 +53,8 @@ fn cdef_block(fs: &FrameState, out: &mut FrameBuf, r: usize, c: usize, idx: usiz let sec = (cp.uv_sec[idx] << coeff_shift) as i32; let dir = if pri == 0 { 0 } else { CDEF_UV_DIR[fs.ssx][fs.ssy][y_dir] as usize }; let damping = cp.damping as i32 + coeff_shift as i32 - 1; - filter(fs, out, 1, r, c, pri, sec, damping, dir); - filter(fs, out, 2, r, c, pri, sec, damping, dir); + filter(fs, out, 1, (r, c), (pri, sec), damping, dir); + filter(fs, out, 2, (r, c), (pri, sec), damping, dir); true } @@ -124,33 +124,74 @@ fn constrain(diff: i32, threshold: i32, adj: u32) -> i32 { if diff < 0 { -mag } else { mag } } -#[allow(clippy::too_many_arguments)] -fn filter(fs: &FrameState, out: &mut FrameBuf, plane: usize, r: usize, c: usize, pri: i32, sec: i32, damping: i32, dir: usize) { - if pri == 0 && sec == 0 { - // Every tap contributes 0 and the clamp to [min, max] keeps x: CdefFrame = CurrFrame. - return; +/// Tap offsets, weights and strengths of one block. +struct Taps { + pri: i32, + sec: i32, + pri_adj: u32, + sec_adj: u32, + pri_off: [isize; 2], + sec_off: [[isize; 2]; 2], + pri_taps: [i32; 2], + sec_taps: [i32; 2], +} + +impl Taps { + fn new(coeff_shift: u32, (pri, sec): (i32, i32), damping: i32, dir: usize) -> Taps { + let pri_tap = ((pri >> coeff_shift) & 1) as usize; + let off = |d: usize, k: usize| CDEF_DIRECTIONS[d][k][0] as isize * PAD_STRIDE as isize + CDEF_DIRECTIONS[d][k][1] as isize; + let pri_off = [off(dir, 0), off(dir, 1)]; + let d_lo = (dir + 6) & 7; + let d_hi = (dir + 2) & 7; + let sec_off = [[off(d_lo, 0), off(d_hi, 0)], [off(d_lo, 1), off(d_hi, 1)]]; + let pri_taps = [CDEF_PRI_TAPS[pri_tap][0] as i32, CDEF_PRI_TAPS[pri_tap][1] as i32]; + let sec_taps = [CDEF_SEC_TAPS[pri_tap][0] as i32, CDEF_SEC_TAPS[pri_tap][1] as i32]; + let adj = |t: i32| if t == 0 { 0 } else { (damping - (31 - (t as u32).leading_zeros()) as i32).max(0) as u32 }; + Taps { pri, sec, pri_adj: adj(pri), sec_adj: adj(sec), pri_off, sec_off, pri_taps, sec_taps } } - let coeff_shift = fs.bit_depth - 8; + + /// The 12 (offset, weight, strength, adjustment) taps in filtering order. + fn tap_list(&self) -> [(isize, i32, i32, u32); 12] { + let mut taps = [(0isize, 0i32, 0i32, 0u32); 12]; + let mut n = 0; + for k in 0..2 { + for sign in [-1isize, 1] { + taps[n] = (sign * self.pri_off[k], self.pri_taps[k], self.pri, self.pri_adj); + taps[n + 1] = (sign * self.sec_off[k][0], self.sec_taps[k], self.sec, self.sec_adj); + taps[n + 2] = (sign * self.sec_off[k][1], self.sec_taps[k], self.sec, self.sec_adj); + n += 3; + } + } + taps + } +} + +/// Position and size of a block within its plane. +#[derive(Clone, Copy)] +struct BlockRect { + x0: usize, + y0: usize, + w: usize, + h: usize, +} + +/// The block with two samples of padding on every side, plus whether every sample is available +/// (CdefAvailable: its 4x4 block is inside the frame). +fn padded_block(fs: &FrameState, plane: usize, src: &Plane, rect: BlockRect) -> ([i32; PAD_STRIDE * PAD_STRIDE], bool) { let (sub_x, sub_y) = if plane > 0 { (fs.ssx, fs.ssy) } else { (0, 0) }; - let x0 = (c * 4) >> sub_x; - let y0 = (r * 4) >> sub_y; - let w = 8 >> sub_x; - let h = 8 >> sub_y; - let src = &fs.cur.planes[plane]; let (mi_rows, mi_cols) = (fs.fh.mi_rows as isize, fs.fh.mi_cols as isize); - // CdefAvailable per sample: its 4x4 block is inside the frame. let mut buf = [CDEF_NA; PAD_STRIDE * PAD_STRIDE]; let mut all_avail = true; - for i in 0..h + 4 { - let y = (y0 + i) as isize - 2; + for i in 0..rect.h + 4 { + let y = (rect.y0 + i) as isize - 2; let cand_r = (y << sub_y) >> 2; if y < 0 || cand_r >= mi_rows { all_avail = false; continue; } let row = src.row(y as usize); - for j in 0..w + 4 { - let x = (x0 + j) as isize - 2; + for j in 0..rect.w + 4 { + let x = (rect.x0 + j) as isize - 2; let cand_c = (x << sub_x) >> 2; if x >= 0 && cand_c < mi_cols { buf[i * PAD_STRIDE + j] = row[x as usize] as i32; @@ -159,70 +200,70 @@ fn filter(fs: &FrameState, out: &mut FrameBuf, plane: usize, r: usize, c: usize, } } } - let pri_tap = ((pri >> coeff_shift) & 1) as usize; - let off = |d: usize, k: usize| CDEF_DIRECTIONS[d][k][0] as isize * PAD_STRIDE as isize + CDEF_DIRECTIONS[d][k][1] as isize; - let pri_off = [off(dir, 0), off(dir, 1)]; - let d_lo = (dir + 6) & 7; - let d_hi = (dir + 2) & 7; - let sec_off = [[off(d_lo, 0), off(d_hi, 0)], [off(d_lo, 1), off(d_hi, 1)]]; - let pri_taps = [CDEF_PRI_TAPS[pri_tap][0] as i32, CDEF_PRI_TAPS[pri_tap][1] as i32]; - let sec_taps = [CDEF_SEC_TAPS[pri_tap][0] as i32, CDEF_SEC_TAPS[pri_tap][1] as i32]; - let adj = |t: i32| if t == 0 { 0 } else { (damping - (31 - (t as u32).leading_zeros()) as i32).max(0) as u32 }; - let (pri_adj, sec_adj) = (adj(pri), adj(sec)); + (buf, all_avail) +} + +/// Filters one sample whose taps may be unavailable (CDEF_NA). +fn filter_sample(buf: &[i32; PAD_STRIDE * PAD_STRIDE], taps: &Taps, ci: isize, x: i32) -> u16 { + let mut sum = 0i32; + let mut max = x; + let mut min = x; + for k in 0..2 { + for sign in [-1isize, 1] { + let p = buf[(ci + sign * taps.pri_off[k]) as usize]; + if p != CDEF_NA { + if taps.pri != 0 { + sum += taps.pri_taps[k] * constrain(p - x, taps.pri, taps.pri_adj); + } + max = max.max(p); + min = min.min(p); + } + for s in 0..2 { + let q = buf[(ci + sign * taps.sec_off[k][s]) as usize]; + if q != CDEF_NA { + if taps.sec != 0 { + sum += taps.sec_taps[k] * constrain(q - x, taps.sec, taps.sec_adj); + } + max = max.max(q); + min = min.min(q); + } + } + } + } + (x + ((8 + sum - (sum < 0) as i32) >> 4)).clamp(min, max) as u16 +} + +fn filter(fs: &FrameState, out: &mut FrameBuf, plane: usize, (r, c): (usize, usize), strengths: (i32, i32), damping: i32, dir: usize) { + let (pri, sec) = strengths; + if pri == 0 && sec == 0 { + // Every tap contributes 0 and the clamp to [min, max] keeps x: CdefFrame = CurrFrame. + return; + } + let coeff_shift = fs.bit_depth - 8; + let (sub_x, sub_y) = if plane > 0 { (fs.ssx, fs.ssy) } else { (0, 0) }; + let rect = BlockRect { x0: (c * 4) >> sub_x, y0: (r * 4) >> sub_y, w: 8 >> sub_x, h: 8 >> sub_y }; + let src = &fs.cur.planes[plane]; + let (buf, all_avail) = padded_block(fs, plane, src, rect); + let taps = Taps::new(coeff_shift, strengths, damping, dir); let dst = &mut out.planes[plane]; if all_avail { // Every tap is available: whole rows at a time (constrain() with threshold 0 is 0). - let mut taps = [(0isize, 0i32, 0i32, 0u32); 12]; - let mut n = 0; - for k in 0..2 { - for sign in [-1isize, 1] { - taps[n] = (sign * pri_off[k], pri_taps[k], pri, pri_adj); - taps[n + 1] = (sign * sec_off[k][0], sec_taps[k], sec, sec_adj); - taps[n + 2] = (sign * sec_off[k][1], sec_taps[k], sec, sec_adj); - n += 3; - } - } - let rows = (y0, x0, h); - match w { - 8 => filter_rows::<8>(&buf, &taps, dst, rows), - _ => filter_rows::<4>(&buf, &taps, dst, rows), + let list = taps.tap_list(); + let rows = (rect.y0, rect.x0, rect.h); + match rect.w { + 8 => filter_rows::<8>(&buf, &list, dst, rows), + _ => filter_rows::<4>(&buf, &list, dst, rows), } return; } - for i in 0..h { + for i in 0..rect.h { // The centre sample is read even where CdefAvailable is 0 (8x8 blocks straddling an odd // MiRows / MiCols edge). - let srow = &src.row(y0 + i)[x0..x0 + w]; - let drow = &mut dst.row_mut(y0 + i)[x0..x0 + w]; - for j in 0..w { + let srow = &src.row(rect.y0 + i)[rect.x0..rect.x0 + rect.w]; + let drow = &mut dst.row_mut(rect.y0 + i)[rect.x0..rect.x0 + rect.w]; + for j in 0..rect.w { let ci = ((i + 2) * PAD_STRIDE + j + 2) as isize; - let x = srow[j] as i32; - let mut sum = 0i32; - let mut max = x; - let mut min = x; - for k in 0..2 { - for sign in [-1isize, 1] { - let p = buf[(ci + sign * pri_off[k]) as usize]; - if p != CDEF_NA { - if pri != 0 { - sum += pri_taps[k] * constrain(p - x, pri, pri_adj); - } - max = max.max(p); - min = min.min(p); - } - for s in 0..2 { - let q = buf[(ci + sign * sec_off[k][s]) as usize]; - if q != CDEF_NA { - if sec != 0 { - sum += sec_taps[k] * constrain(q - x, sec, sec_adj); - } - max = max.max(q); - min = min.min(q); - } - } - } - } - drow[j] = (x + ((8 + sum - (sum < 0) as i32) >> 4)).clamp(min, max) as u16; + drow[j] = filter_sample(&buf, &taps, ci, srow[j] as i32); } } } diff --git a/crates/av1/src/restoration.rs b/crates/av1/src/restoration.rs index 0b829733..b31f982b 100644 --- a/crates/av1/src/restoration.rs +++ b/crates/av1/src/restoration.rs @@ -73,57 +73,70 @@ impl LrCtx<'_> { } } +/// Position and size of a (at most 4x4) restoration block within its plane. +#[derive(Clone, Copy)] +struct LrBlock { + x: i32, + y: i32, + w: i32, + h: i32, +} + /// Loop restoration (7.17): returns LrFrame given UpscaledCurrFrame and UpscaledCdefFrame. pub(crate) fn loop_restoration(fs: &FrameState, cur: &FrameBuf, cdef: &FrameBuf) -> FrameBuf { - let fh = &fs.fh; let mut lr = cdef.clone(); - let up_w = fh.upscaled_width as i32; - let fr_h = fh.frame_height as i32; - let bd = fs.bit_depth; for plane in 0..fs.num_planes { - let rtype_frame = fh.lr.frame_restoration_type[plane]; - if rtype_frame == RESTORE_NONE as u8 { + if fs.fh.lr.frame_restoration_type[plane] == RESTORE_NONE as u8 { continue; } - let (sub_x, sub_y) = if plane > 0 { (fs.ssx as i32, fs.ssy as i32) } else { (0, 0) }; - let unit_size = fh.lr.loop_restoration_size[plane] as i32; - let unit_rows = fs.lr_unit_rows[plane] as i32; - let unit_cols = fs.lr_unit_cols[plane] as i32; - let plane_end_x = ((up_w + sub_x) >> sub_x) - 1; - let plane_end_y = ((fr_h + sub_y) >> sub_y) - 1; - let mut row = 0; - while row * 4 < fr_h { - let luma_y = row * 4; - let stripe_num = (luma_y + 8) / 64; - let stripe_start_y = (-8 + stripe_num * 64) >> sub_y; - let stripe_end_y = stripe_start_y + (64 >> sub_y) - 1; - let ctx = LrCtx { cur: &cur.planes[plane], cdef: &cdef.planes[plane], plane_end_x, plane_end_y, stripe_start_y, stripe_end_y }; - let unit_row = (unit_rows - 1).min(((row * 4 + 8) >> sub_y) / unit_size); - let mut col = 0; - while col * 4 < up_w { - let unit_col = (unit_cols - 1).min(((col * 4) >> sub_x) / unit_size); - let x = (col * 4) >> sub_x; - let y = (row * 4) >> sub_y; - let w = (4 >> sub_x).min(plane_end_x - x + 1); - let h = (4 >> sub_y).min(plane_end_y - y + 1); - if w > 0 && h > 0 { - let ui = (unit_row * unit_cols + unit_col) as usize; - match fs.lr.lr_type[plane][ui] { - t if t == RESTORE_WIENER as u8 => wiener(fs, &ctx, &mut lr.planes[plane], plane, ui, x, y, w, h, bd), - t if t == RESTORE_SGRPROJ as u8 => self_guided(fs, &ctx, &mut lr.planes[plane], plane, ui, x, y, w, h, bd), - _ => {} - } - } - col += 1; - } - row += 1; - } + restore_plane(fs, cur, cdef, &mut lr.planes[plane], plane); } lr } -#[allow(clippy::too_many_arguments)] -fn wiener(fs: &FrameState, ctx: &LrCtx, out: &mut Plane, plane: usize, ui: usize, x: i32, y: i32, w: i32, h: i32, bd: u32) { +fn restore_plane(fs: &FrameState, cur: &FrameBuf, cdef: &FrameBuf, out: &mut Plane, plane: usize) { + let fh = &fs.fh; + let up_w = fh.upscaled_width as i32; + let fr_h = fh.frame_height as i32; + let bd = fs.bit_depth; + let (sub_x, sub_y) = if plane > 0 { (fs.ssx as i32, fs.ssy as i32) } else { (0, 0) }; + let unit_size = fh.lr.loop_restoration_size[plane] as i32; + let unit_rows = fs.lr_unit_rows[plane] as i32; + let unit_cols = fs.lr_unit_cols[plane] as i32; + let plane_end_x = ((up_w + sub_x) >> sub_x) - 1; + let plane_end_y = ((fr_h + sub_y) >> sub_y) - 1; + let mut row = 0; + while row * 4 < fr_h { + let luma_y = row * 4; + let stripe_num = (luma_y + 8) / 64; + let stripe_start_y = (-8 + stripe_num * 64) >> sub_y; + let stripe_end_y = stripe_start_y + (64 >> sub_y) - 1; + let ctx = LrCtx { cur: &cur.planes[plane], cdef: &cdef.planes[plane], plane_end_x, plane_end_y, stripe_start_y, stripe_end_y }; + let unit_row = (unit_rows - 1).min(((row * 4 + 8) >> sub_y) / unit_size); + let mut col = 0; + while col * 4 < up_w { + let unit_col = (unit_cols - 1).min(((col * 4) >> sub_x) / unit_size); + let x = (col * 4) >> sub_x; + let y = (row * 4) >> sub_y; + let w = (4 >> sub_x).min(plane_end_x - x + 1); + let h = (4 >> sub_y).min(plane_end_y - y + 1); + if w > 0 && h > 0 { + let ui = (unit_row * unit_cols + unit_col) as usize; + let blk = LrBlock { x, y, w, h }; + match fs.lr.lr_type[plane][ui] { + t if t == RESTORE_WIENER as u8 => wiener(fs, &ctx, out, (plane, ui), blk, bd), + t if t == RESTORE_SGRPROJ as u8 => self_guided(fs, &ctx, out, (plane, ui), blk, bd), + _ => {} + } + } + col += 1; + } + row += 1; + } +} + +fn wiener(fs: &FrameState, ctx: &LrCtx, out: &mut Plane, (plane, ui): (usize, usize), blk: LrBlock, bd: u32) { + let LrBlock { x, y, w, h } = blk; let round0: u32 = if bd == 12 { 5 } else { 3 }; let round1: u32 = if bd == 12 { 9 } else { 11 }; let coef = &fs.lr.lr_wiener[plane][ui]; @@ -166,13 +179,13 @@ fn wiener(fs: &FrameState, ctx: &LrCtx, out: &mut Plane, plane: usize, ui: usize } } -#[allow(clippy::too_many_arguments)] -fn self_guided(fs: &FrameState, ctx: &LrCtx, out: &mut Plane, plane: usize, ui: usize, x: i32, y: i32, w: i32, h: i32, bd: u32) { +fn self_guided(fs: &FrameState, ctx: &LrCtx, out: &mut Plane, (plane, ui): (usize, usize), blk: LrBlock, bd: u32) { + let LrBlock { x, y, w, h } = blk; let set = fs.lr.lr_sgr_set[plane][ui] as usize; let mut flt0 = [[0i32; 4]; 4]; let mut flt1 = [[0i32; 4]; 4]; - box_filter(ctx, x, y, w, h, set, 0, bd, &mut flt0); - box_filter(ctx, x, y, w, h, set, 1, bd, &mut flt1); + box_filter(ctx, blk, (set, 0), bd, &mut flt0); + box_filter(ctx, blk, (set, 1), bd, &mut flt1); let xqd = fs.lr.lr_sgr_xqd[plane][ui]; let w0 = xqd[0] as i32; let w1 = xqd[1] as i32; @@ -198,8 +211,56 @@ fn round2_64(x: i64, n: u32) -> i64 { if n == 0 { x } else { (x + (1i64 << (n - 1))) >> n } } -#[allow(clippy::too_many_arguments)] -fn box_filter(ctx: &LrCtx, x: i32, y: i32, w: i32, h: i32, set: usize, pass: usize, bd: u32, f: &mut [[i32; 4]; 4]) { +/// Box sums of the samples within `r` of (x, y): (sum of squares, sum). +fn box_sums(ctx: &LrCtx, x: i32, y: i32, r: i32) -> (i64, i64) { + let mut a = 0i64; + let mut b = 0i64; + for dy in -r..=r { + for dx in -r..=r { + let c = ctx.sample(x + dx, y + dy) as i64; + a += c * c; + b += c; + } + } + (a, b) +} + +/// Self-guided filter constants of one pass. +struct SgrConsts { + n: i64, + s: i64, + one_over_n: i64, +} + +/// (A, B) coefficients of one sample from its box sums. +fn sgr_coefficients(a: i64, b: i64, k: &SgrConsts, bd: u32) -> (i64, i64) { + let a = round2_64(a, 2 * (bd - 8)); + let d = round2_64(b, bd - 8); + let p = (a * k.n - d * d).max(0); + let z = round2_64(p * k.s, SGRPROJ_MTABLE_BITS as u32); + let a2 = if z >= 255 { + 256 + } else if z == 0 { + 1 + } else { + ((z << SGRPROJ_SGR_BITS) + z / 2) / (z + 1) + }; + let b2 = ((1i64 << SGRPROJ_SGR_BITS) - a2) * b * k.one_over_n; + (a2, round2_64(b2, SGRPROJ_RECIP_BITS as u32)) +} + +fn box_weight(pass: usize, i: i32, dx: i32, dy: i32) -> i64 { + if pass == 0 { + if ((i + dy) & 1) != 0 { if dx == 0 { 6 } else { 5 } } else { 0 } + } else if dx == 0 || dy == 0 { + 4 + } else { + 3 + } +} + +fn box_filter(ctx: &LrCtx, blk: LrBlock, (set, pass): (usize, usize), bd: u32, f: &mut [[i32; 4]; 4]) { + let LrBlock { x, y, w, h } = blk; let r = SGR_PARAMS[set][pass * 2] as i32; if r == 0 { return; @@ -207,35 +268,15 @@ fn box_filter(ctx: &LrCtx, x: i32, y: i32, w: i32, h: i32, set: usize, pass: usi let eps = SGR_PARAMS[set][pass * 2 + 1] as i64; let n = ((2 * r + 1) * (2 * r + 1)) as i64; let n2e = n * n * eps; - let s = ((1i64 << SGRPROJ_MTABLE_BITS) + n2e / 2) / n2e; - let one_over_n = ((1i64 << SGRPROJ_RECIP_BITS) + n / 2) / n; + let k = SgrConsts { n, s: ((1i64 << SGRPROJ_MTABLE_BITS) + n2e / 2) / n2e, one_over_n: ((1i64 << SGRPROJ_RECIP_BITS) + n / 2) / n }; let mut a_arr = [[0i64; 6]; 6]; let mut b_arr = [[0i64; 6]; 6]; for i in -1..h + 1 { for j in -1..w + 1 { - let mut a = 0i64; - let mut b = 0i64; - for dy in -r..=r { - for dx in -r..=r { - let c = ctx.sample(x + j + dx, y + i + dy) as i64; - a += c * c; - b += c; - } - } - let a = round2_64(a, 2 * (bd - 8)); - let d = round2_64(b, bd - 8); - let p = (a * n - d * d).max(0); - let z = round2_64(p * s, SGRPROJ_MTABLE_BITS as u32); - let a2 = if z >= 255 { - 256 - } else if z == 0 { - 1 - } else { - ((z << SGRPROJ_SGR_BITS) + z / 2) / (z + 1) - }; - let b2 = ((1i64 << SGRPROJ_SGR_BITS) - a2) * b * one_over_n; + let (a, b) = box_sums(ctx, x + j, y + i, r); + let (a2, b2) = sgr_coefficients(a, b, &k, bd); a_arr[(i + 1) as usize][(j + 1) as usize] = a2; - b_arr[(i + 1) as usize][(j + 1) as usize] = round2_64(b2, SGRPROJ_RECIP_BITS as u32); + b_arr[(i + 1) as usize][(j + 1) as usize] = b2; } } for i in 0..h { @@ -245,13 +286,7 @@ fn box_filter(ctx: &LrCtx, x: i32, y: i32, w: i32, h: i32, set: usize, pass: usi let mut b = 0i64; for dy in -1..=1i32 { for dx in -1..=1i32 { - let weight = if pass == 0 { - if ((i + dy) & 1) != 0 { if dx == 0 { 6 } else { 5 } } else { 0 } - } else if dx == 0 || dy == 0 { - 4 - } else { - 3 - }; + let weight = box_weight(pass, i, dx, dy); a += weight * a_arr[(i + dy + 1) as usize][(j + dx + 1) as usize]; b += weight * b_arr[(i + dy + 1) as usize][(j + dx + 1) as usize]; }