fix(pyrowave): guard 4:4:4 modes that overflow the rate controller's block index

The vendored rate controller packs its wavelet block index into 16 bits
(RDOperation.block_offset_saving), so a mode whose 32x32-block count exceeds
u16::MAX wraps inside the controller and corrupts the bitstream — ~8K-class
4:4:4 territory. Compute the exact count (`block_count_32x32`, the counting walk
of upstream init_block_meta, pinned against the validated Apple WaveletLayout)
and expose `pyrowave_mode_fits_rdo`; the negotiator downgrades such a session to
4:2:0 before the Welcome (the honest-downgrade channel), and both encoders
refuse outright if one slips through rather than emit a wrapped stream.

Vendor patches: 0002-rdo-saving-clamp (analyze_rate_control.comp clamps the
saving accumulation to the target, same overrun class as 0001; slangmosh.hpp
regenerated), 0003-devel-encode-16bit-read (devel tool y4m 16-bit plane reads;
tool-only, kept so the vendored source stays honest).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-07-19 01:16:51 +02:00
parent ac0e73321c
commit 75474dcc90
11 changed files with 5877 additions and 1896 deletions
@@ -227,6 +227,14 @@ impl PyroWaveEncoder {
if !chroma.is_444() && (width % 2 != 0 || height % 2 != 0) {
bail!("pyrowave 4:2:0 needs even dimensions (got {width}x{height})");
}
if chroma.is_444() && !crate::pyrowave_mode_fits_rdo(width, height, true) {
// The negotiator downgrades these modes to 4:2:0 pre-Welcome; refuse if one
// slips through (e.g. the lab override) rather than wrap the RDO block index.
bail!(
"pyrowave 4:4:4 at {width}x{height} exceeds the rate controller's 16-bit \
block index (see pyrowave-sys patches/0002 note) — use 4:2:0 at this size"
);
}
// SAFETY: `open_inner` only issues Vulkan/pyrowave calls whose preconditions it
// establishes itself (valid instance/device, correctly-chained create-infos that
// `DeviceHold` keeps alive); all handles are freshly created and owned by the result.
+59
View File
@@ -53,6 +53,34 @@ pub(crate) fn stamp_color_bits(bitstream: &mut [u8], seq_offset: usize, bt2020_p
}
}
/// The wavelet block space's total 32x32-block count for a mode — the exact counting walk of
/// upstream `WaveletBuffers::init_block_meta` (also ported to the Apple `WaveletLayout`, whose
/// golden tests pin it against real host AUs). Needed because the vendored RDO pass packs the
/// block index into 16 bits (`RDOperation.block_offset_saving` — see
/// `patches/0002-rdo-saving-clamp.patch`): a mode whose count exceeds `u16::MAX` would wrap
/// inside the rate controller, so the host guards such modes out (≈8K 4:4:4 territory).
pub(crate) fn block_count_32x32(width: u32, height: u32, chroma444: bool) -> u32 {
const LEVELS: u32 = 5;
let align = |v: u32| ((v + 31) & !31).max(128);
let (aw, ah) = (align(width), align(height));
let mut count = 0u32;
for level in (0..LEVELS).rev() {
let lw = (aw / 2) >> level;
let lh = (ah / 2) >> level;
let blocks_x8 = lw.div_ceil(8);
let blocks_y8 = lh.div_ceil(8);
let per_band = blocks_x8.div_ceil(4) * blocks_y8.div_ceil(4);
let bands = if level == LEVELS - 1 { 4 } else { 3 };
for component in 0..3u32 {
if level == 0 && component != 0 && !chroma444 {
continue;
}
count += per_band * bands;
}
}
count
}
/// Frame pyrowave's `packets` (each an `(offset, size)` into `bitstream`) into the wire AU.
/// `wire_chunk = None` copies the single dense packet; `Some(chunk)` produces the windowed
/// datagram-aligned AU (a whole number of `chunk`-sized windows).
@@ -189,6 +217,37 @@ mod tests {
assert_eq!(packet_boundary(None, 777), 777);
}
#[test]
fn block_count_matches_the_apple_layout_invariant() {
// 256x144 (the golden-fixture geometry, aligned 256x160): recompute via the same walk
// the validated Apple WaveletLayout uses and pin a few mode-level facts.
let manual = |w: u32, h: u32, c444: bool| {
let align = |v: u32| ((v + 31) & !31).max(128);
let (aw, ah) = (align(w), align(h));
let mut n = 0u32;
for level in (0..5u32).rev() {
let per = ((aw / 2 >> level).div_ceil(8).div_ceil(4))
* ((ah / 2 >> level).div_ceil(8).div_ceil(4));
let bands = if level == 4 { 4 } else { 3 };
for c in 0..3 {
if level == 0 && c != 0 && !c444 {
continue;
}
n += per * bands;
}
}
n
};
for (w, h) in [(256, 144), (1920, 1080), (3840, 2160), (7680, 4320)] {
assert_eq!(block_count_32x32(w, h, false), manual(w, h, false));
assert_eq!(block_count_32x32(w, h, true), manual(w, h, true));
}
// 4:4:4 fits comfortably at 4K; the 16-bit RDO block index wraps around 8K 4:4:4.
assert!(block_count_32x32(3840, 2160, true) <= u16::MAX as u32);
assert!(block_count_32x32(7680, 4320, true) > u16::MAX as u32);
assert!(block_count_32x32(7680, 4320, false) <= u16::MAX as u32);
}
#[test]
fn stamp_color_bits_sets_range_and_hdr_bits() {
let mut bs = vec![0u8; 16];
@@ -182,6 +182,14 @@ impl PyroWaveEncoder {
if !chroma444 && (width % 2 != 0 || height % 2 != 0) {
bail!("pyrowave 4:2:0 needs even dimensions (got {width}x{height})");
}
if chroma444 && !crate::pyrowave_mode_fits_rdo(width, height, true) {
// The negotiator downgrades these modes to 4:2:0 pre-Welcome; refuse if one
// slips through (e.g. the lab override) rather than wrap the RDO block index.
bail!(
"pyrowave 4:4:4 at {width}x{height} exceeds the rate controller's 16-bit \
block index (see pyrowave-sys patches/0002 note) — use 4:2:0 at this size"
);
}
let fps = fps.max(1);
// Select pyrowave's device by the SELECTED render adapter's vendor/device-id — NOT by LUID:
// in Session 0 (the host service context) the Vulkan ICD reports `deviceLUIDValid = false`,
+12
View File
@@ -1326,6 +1326,18 @@ mod pyrowave;
#[path = "enc/pyrowave_wire.rs"]
mod pyrowave_wire;
/// Whether a PyroWave mode fits the vendored rate controller's packed 16-bit block index
/// (`patches/0002-rdo-saving-clamp.patch` note): false ≈ 8K-class 4:4:4. The negotiator
/// downgrades such a session to 4:2:0 before the Welcome; the encoders also refuse outright.
#[cfg(all(any(target_os = "linux", target_os = "windows"), feature = "pyrowave"))]
pub fn pyrowave_mode_fits_rdo(width: u32, height: u32, chroma444: bool) -> bool {
pyrowave_wire::block_count_32x32(width, height, chroma444) <= u16::MAX as u32
}
#[cfg(not(all(any(target_os = "linux", target_os = "windows"), feature = "pyrowave")))]
pub fn pyrowave_mode_fits_rdo(_width: u32, _height: u32, _chroma444: bool) -> bool {
false
}
#[cfg(test)]
mod tests {
use super::*;