diff --git a/Cargo.lock b/Cargo.lock index 333cb1b..1bbb9e5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -180,6 +180,37 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + [[package]] name = "fdeflate" version = "0.3.7" @@ -236,6 +267,7 @@ dependencies = [ "png", "qcms", "rav1d", + "rayon", "safe_unaligned_simd", "scuffle-h265", ] @@ -482,6 +514,26 @@ dependencies = [ "bitflags", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" diff --git a/Cargo.toml b/Cargo.toml index 0239265..b46dac8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,8 +8,10 @@ license-file = "LICENSE" repository = "https://github.com/ente-io/heic-decoder" [features] -default = ["std"] +default = ["std", "parallel-grid"] +decoder-tracing = ["std"] image-integration = ["dep:image"] +parallel-grid = ["std", "dep:rayon"] std = ["archmage/std"] [[bin]] @@ -28,6 +30,7 @@ image = { version = "0.25", default-features = false, features = ["png"], option libc = "0.2" moxcms = "0.8.1" png = "0.18" +rayon = { version = "1.10", optional = true } rav1d = { git = "https://github.com/ente-io/rav1d.git", branch = "arm32_stable", default-features = false, features = [ "bitdepth_8", "bitdepth_16", diff --git a/SPEED_IMPROVEMENTS.md b/SPEED_IMPROVEMENTS.md new file mode 100644 index 0000000..6eae29c --- /dev/null +++ b/SPEED_IMPROVEMENTS.md @@ -0,0 +1,353 @@ +# HEIC decoder speed improvements + +Date: 2026-07-16 + +Branch: `speed_improvements` + +## Goal + +Improve HEIC decoding performance for Ente's Rust ML indexing pipeline without +changing decoded pixels, orientation, color handling, public behavior, or code +readability. ONNX Runtime inference is deliberately outside this work. + +The implementation work is intentionally split into four independently +reviewable and revertible commits: + +1. compile diagnostic counters and CTU tracking out of production builds; +2. decode independent HEIF grid tiles with bounded parallelism while preserving + deterministic validation and paste order; +3. provide a direct RGB8 decode path so RGB-only consumers do not materialize + RGBA and then discard alpha; +4. add exact ARM NEON kernels for profiled HEVC hot paths. + +## Evidence and prioritization + +The optimized iPhone 15 Pro pipeline benchmark covers 14 representative images +with face and CLIP indexing enabled. Its summed release-mode medians are: + +| Stage | Time | Share of Rust-controlled non-inference time | +|---|---:|---:| +| Decode | 3,256.4 ms | 98.2% | +| Preprocessing | 49.1 ms | 1.5% | +| Alignment and postprocessing | 9.0 ms | 0.3% | + +The five HEIC fixtures account for 3,021.6 ms of decode time. Improving ML +preprocessing and postprocessing cannot materially change indexing latency; +the decoder is the meaningful Rust-side opportunity. + +Alternative decoder experiments reinforce this conclusion: + +| Decoder | HEIC decode sum | HEIC speedup | Full-corpus reduction | +|---|---:|---:|---:| +| Current decoder | 3,021.6 ms | 1.00x | - | +| ImageIO with current fallback | 1,846.9 ms | 1.64x | 30.3% | +| `heic` 0.1.6 with parallel grids | 1,875.5 ms | 1.61x | 29.7% | +| `libheif-rs` with libde265 | 1,625.3 ms | 1.86x | 35.6% | + +The parallel `heic` experiment was roughly even with the current decoder on +ordinary single-image HEICs but improved the two grid/problem fixtures by +3.16x and 2.67x. This makes grid scheduling the largest demonstrated +algorithmic opportunity in the current pure-Rust implementation. + +The alternatives passed Ente's existing downstream face and CLIP thresholds, +but that is weaker than byte identity. Changes within this repository must use +the current scalar/sequential implementation as an exact oracle wherever +possible. + +## Existing strengths to preserve + +- The decoder has pixel-for-pixel differential tests against `heif-dec` over + libheif samples, Ente camera fixtures, and a generated HEVC stress corpus. +- The image integration lazily decodes into the buffer supplied by + `image::ImageDecoder::read_image`, avoiding an extra owned RGBA handoff. +- Grid output is pasted directly into the final canvas and can apply supported + orientation transforms while pasting. +- The decoder exposes guardrails for pixel counts and input buffering. +- x86-64 transform, residual, dequantization, and 4:2:0 color conversion paths + already have scalar fallbacks suitable as exact SIMD test oracles. + +## 1. Compile out production diagnostics + +### Finding + +The default `std` feature currently enables debugging work in normal release +decodes even though `DEBUG_TRACE` is false and `SE_TRACE_LIMIT` is zero: + +- every syntax-element trace site increments a global atomic counter; +- every decoded CTU reads the CABAC position, locks a global mutex, and appends + to the tracker vector; +- large-coefficient tracking locks the same mutex; +- residual and NxN diagnostic counters increment in hot paths. + +The collected tracker is printed only when `DEBUG_TRACE` is true, so this work +has no production consumer. + +### Recommendation + +Introduce an opt-in `decoder-tracing` Cargo feature and compile all tracing +state and calls to no-ops when neither that feature nor tests are enabled. +Correctness checks that affect decode behavior should remain separate from +diagnostic logging. + +### Measurement + +Run a randomized release A/B on the same local corpus and command before and +after the change. Record the exact commands, selected files, aggregate time, +and per-file medians here when implemented. + +## 2. Bounded deterministic grid-tile parallelism + +### Finding + +HEIF grid tiles are independent coded image items, but both the planar grid +decode and direct RGBA grid paths decode them sequentially. Real problem +fixtures contain dozens of coded items, leaving most device cores idle while +decode dominates latency. + +### Recommendation + +- Add a small internal bounded worker scheduler rather than unbounded + thread-per-tile execution. +- Decode tiles in parallel, attach their original row-major indices, and + consume results in index order for deterministic error selection, + validation, and paste behavior. +- Decode the first tile synchronously because it establishes geometry, layout, + bit depth, range, matrix, and output allocation. +- Bound concurrency by available parallelism, remaining tile count, and a + conservative decoded-tile memory budget. +- Keep one sequential path available to tests as the exact oracle. +- Do not parallelize output writes unless non-overlap is mechanically proven; + ordered paste is cheap and avoids unsafe aliasing. + +The preferred memory shape is a bounded sliding window of decoded tiles, not a +`Vec` containing every decoded tile in a large panorama. + +## 3. Direct RGB8 output + +### Finding + +Ente's ML pipeline ultimately calls `DynamicImage::into_rgb8`. The HEIC image +hook advertises and fills RGBA8, after which `image` allocates RGB8, copies +three channels, and discards alpha. For a 55.9-megapixel panorama, the RGBA and +RGB allocations are approximately 224 MB and 168 MB respectively, excluding +YUV planes and decoder scratch. + +The measured post-decode RGBA-to-RGB copy is only tens of milliseconds across +the benchmark corpus, so the primary benefit is lower peak memory and safer +parallelism rather than a large standalone latency reduction. + +### Recommendation + +- Add first-class `DecodedRgbImage`/`DecodedRgbPixels` APIs and direct RGB8 + caller-buffer functions. +- Share color conversion and transform logic with RGBA rather than maintaining + a second decoder pipeline. +- Use direct RGB only when alpha is absent or the caller explicitly requests + alpha discard. Preserve RGBA behavior for alpha-bearing images. +- Keep the `image` hook's existing RGBA contract for compatibility unless the + consuming integration can explicitly request the RGB path. +- Verify direct RGB byte-for-byte against dropping alpha from the existing + final transformed RGBA8 result. + +## 4. Exact ARM NEON kernels + +### Finding + +Explicit decoder SIMD is currently x86-64 AVX2. AArch64 builds dispatch to the +scalar implementations for inverse transforms, dequantization, residual add, +and the fixed-point 4:2:0 YCbCr-to-RGB kernel. These operations are natural +NEON candidates and are relevant on every supported iPhone. + +### Recommendation + +Profile an AArch64 release build first and implement kernels in measured order. +Likely candidates are: + +1. residual add and clamp; +2. 4:2:0 YCbCr-to-RGB interleaving; +3. dequantization; +4. 8x8/16x16 inverse transforms, followed by larger transforms only if the + profile justifies them. + +Use `core::arch::aarch64` intrinsics behind compile-time architecture gates and +retain runtime dispatch conventions consistent with the existing x86 code. +Every kernel needs randomized scalar-vs-NEON unit tests covering saturation, +rounding, bit depths, odd widths, unaligned buffers, and scalar tails. Full +corpus output must remain pixel-identical. + +## Accuracy gates + +Each commit must pass, in increasing scope: + +1. `cargo fmt --check`; +2. `cargo clippy --all-targets --all-features -- -D warnings`; +3. `cargo test --all-features`; +4. quick pixel verification during development; +5. the repository's complete `scripts/heic_tests.sh all` suite after all four + changes. + +The final Ente integration must additionally preserve face detections, +landmarks, blur scores, face embeddings, and CLIP thresholds. Pet models were +not enabled in the original benchmark, so pet-enabled parity should be covered +separately before treating those results as validated. + +The HEIC corpus should include all EXIF orientations, `irot`/`imir`, ICC and +Display P3 content, 8/10/12-bit inputs, alpha, odd dimensions, ordinary images, +large grids, and malformed-input fallback behavior. + +## Benchmark discipline + +- Build Rust and the application in optimized release mode; reject debug or + accidentally unoptimized artifacts. +- Warm model sessions before measured indexing passes. +- Remove synchronous stage logging for final wall-clock comparisons. +- Randomize decoder/fixture order where practical and report medians plus + ranges or confidence intervals. +- Measure peak RSS alongside latency, particularly for grid parallelism. +- Cap thread counts explicitly so results are reproducible and nested + parallelism cannot oversubscribe the device. + +## Deferred opportunities + +- Cross-image pipelining may improve library throughput, but should follow the + decoder memory reduction because several simultaneous full-resolution HEICs + can consume gigabytes. +- ThinLTO and fewer codegen units are safe release-build experiments after the + larger algorithmic work. +- WPP row parallelism is substantially more complex than independent HEIF grid + tiles and should be considered only after profiles show remaining HEVC-core + limits. +- The CR2 Dart-to-JPEG fallback is both slow and inaccurate, but is separate + from this HEIC decoder work. + +## Implementation results + +### Production diagnostics + +Normal builds no longer contain syntax-element, residual, or NxN atomic +counters, CABAC-position collection, per-CTU mutex/vector tracking, or +large-coefficient tracking. The diagnostics remain available through the +opt-in `decoder-tracing` feature. + +The repository's full decode benchmark was run before and after with 12 files +and five release runs per file. Raw summed Rust averages moved from 19.408 s to +18.816 s (3.1% lower), but the external validator moved from 19.942 s to +19.220 s (3.6% lower) during the same sequential runs. The normalized +Rust/validator ratio therefore changed from 0.973x to 0.979x, showing that the +raw improvement was ambient machine variation rather than a defensible patch +effect. + +A tighter alternating comparison used independently built pre-change and +post-change release binaries: + +| Fixture | Pre-change median | Post-change median | Difference | +|---|---:|---:|---:| +| 55.9 MP grid panorama | 9.31 s | 9.29 s | -0.2% | +| Ordinary 8.6 MP HEIC | 1.65 s | 1.67 s | +1.2% | + +The corresponding averages were 9.392 s versus 9.334 s for the panorama and +1.658 s versus 1.670 s for the ordinary image. These differences are within +run-to-run noise. Both binaries produced identical output SHA-256 hashes. + +Conclusion: compiling diagnostics out is worthwhile production cleanup and +removes global synchronization and unused allocations, but it did not produce +a measurable end-to-end speedup on this Mac benchmark. + +### Bounded deterministic grid parallelism + +The default build now decodes independent HEIF grid tiles through Rayon's +shared worker pool. The first tile remains synchronous and establishes the +reference metadata. Remaining tiles are processed in windows capped at eight +workers and a conservative 64 MiB decoded-tile working estimate. Results are +then validated, converted, and pasted sequentially in row-major order. Builds +with `parallel-grid` disabled retain the sequential oracle, and +`decoder-tracing` builds stay sequential so their global trace remains useful. + +Alternating five-run release comparisons used the repository's image-adapter +checksum benchmark, which decodes directly into the same caller-buffer path +used by the `image` integration without adding PNG encoding time: + +| Fixture | Sequential median | Parallel median | Speedup | Reduction | +|---|---:|---:|---:|---:| +| 55.9 MP grid panorama | 1.20 s | 0.57 s | 2.11x | 52.5% | +| Text grid fixture | 0.29 s | 0.13 s | 2.23x | 55.2% | + +Every sequential and parallel run returned the same decoded checksum. A +separate full PNG hash comparison was also byte-identical for both fixtures. +On one warmed panorama run, peak RSS increased from 250,068,992 bytes to +282,771,456 bytes (13.1%) while elapsed adapter time fell from 1.18 s to +0.57 s. The bounded window makes that memory/speed tradeoff explicit. + +The final Rayon implementation passed the full differential verifier: 272 +files classified, 219 exact validator comparisons, 219 exact image-hook +comparisons, and zero failures. + +### Direct RGB8 output + +The decoder now exposes owned `decode_*_to_rgb8` APIs for coded HEIC/HEIF. +They share the existing YCbCr conversion, crop, orientation, grid validation, +and deterministic tile scheduling logic, but write three channels directly. +The existing RGBA APIs and `image` hook remain unchanged. Calling RGB8 is an +explicit request to discard auxiliary alpha. For high-bit-depth input, the +conversion uses the same `(sample + 128) / 257` reduction as `image`'s final +RGBA16-to-RGB8 conversion. + +The direct output was compared byte-for-byte with the existing image hook +followed by `DynamicImage::into_rgb8`. All 57 coded HEIC/HEIF files that this +decoder accepted from the 269-file local HEIF corpus matched exactly. This set +included the six Ente fixtures and 45 generated stress images covering 8-, +10-, and 12-bit decode paths. Unit tests also exercise transformed 8-bit and +10-bit synthetic images. + +Alternating five-run release measurements included both decode and the final +RGB conversion: + +| Fixture | RGBA then RGB median | Direct RGB median | Latency reduction | RGBA then RGB median peak RSS | Direct RGB median peak RSS | RSS reduction | +|---|---:|---:|---:|---:|---:|---:| +| 55.9 MP grid panorama | 0.55 s | 0.45 s | 18.2% | 449,150,976 B | 218,677,248 B | 51.3% | +| Ordinary 8.6 MP HEIC | 0.12 s | 0.11 s | 8.3% | 101,695,488 B | 50,872,320 B | 50.0% | + +The timing resolution is coarse for the ordinary image, but the panorama +result is stable across all five pairs. The much larger and consistent RSS +reduction is the primary result: the decoder no longer keeps the final RGBA +allocation alive while `image` allocates the replacement RGB buffer. + +### Exact ARM NEON kernels + +A fresh 1 ms Time Profiler capture used an optimized ARM64 build with debug +symbols and the 55.9 MP panorama's direct-RGB path. The largest self-sample +counts were residual parsing (321), CABAC bin decoding (198), YCbCr/RGB work +(160 combined), scaling-list dequantization (72), 32x32 IDCT (58), 16x16 IDCT +(27), and residual add/clamp (16). Entropy parsing is the dominant limit but +is serial and branch-heavy. Scaling-list dequantization and residual add were +therefore the best bounded SIMD targets; they are measurable hot paths and can +be expressed compactly with an exact scalar oracle. + +The AArch64 build now dispatches through NEON for: + +- scaling-list dequantization using i32 coefficient/matrix products widened to + i64 before the QP multiply, rounding shift, and saturating narrow; +- flat dequantization using i32 multiply/round/shift and saturating narrow; +- residual add using signed saturating addition followed by the exact HEVC + bit-depth clamp. + +The existing scalar implementation remains the fallback and test oracle. +Randomized NEON-vs-scalar tests cover 4/8/16/32 blocks, 8/10/12/14/16-bit clamp +ranges, extreme residuals, non-contiguous strides, scaling matrices, saturating +coefficients, and non-vector-length tails. The optimized library also checks +successfully for the `aarch64-apple-ios` target. + +In a follow-up profile, scaling-list dequantization fell from 72 to 30 self +samples and residual add from 16 to 9. Sample counts across separate captures +are directional rather than a precise benchmark, so the release A/B used 15 +alternating pairs instead: + +| Fixture | Scalar median | NEON median | Median reduction | Scalar average | NEON average | +|---|---:|---:|---:|---:|---:| +| 55.9 MP grid panorama | 0.461070 s | 0.455414 s | 1.2% | 0.460800 s | 0.456821 s | +| Ordinary 8.6 MP HEIC | 0.117213 s | 0.115700 s | 1.3% | 0.117346 s | 0.115630 s | + +The gain is intentionally modest because CABAC and residual parsing dominate. +The profile still shows 16x16/32x32 transforms as possible follow-up targets, +but porting their large AVX2 butterfly implementation would be a substantially +larger correctness and readability risk for the remaining single-digit share. diff --git a/src/bin/heif-image-adapter-bench.rs b/src/bin/heif-image-adapter-bench.rs index af18496..9378e9d 100644 --- a/src/bin/heif-image-adapter-bench.rs +++ b/src/bin/heif-image-adapter-bench.rs @@ -1,5 +1,5 @@ use heic_decoder::image_integration::register_image_decoder_hooks; -use heic_decoder::{DecodedRgbaPixels, decode_path_to_rgba}; +use heic_decoder::{DecodedRgbaPixels, decode_path_to_rgb8, decode_path_to_rgba}; use image::{DynamicImage, ImageReader}; use std::env; use std::error::Error; @@ -9,7 +9,10 @@ use std::path::Path; #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum DecodeMode { Direct, + Rgb, Adapter, + AdapterRgb, + VerifyRgb, } #[derive(Debug)] @@ -24,7 +27,7 @@ impl Display for CliError { impl Error for CliError {} fn usage() -> &'static str { - "Usage: heif-image-adapter-bench " + "Usage: heif-image-adapter-bench " } fn parse_args() -> Result<(DecodeMode, String), CliError> { @@ -35,7 +38,10 @@ fn parse_args() -> Result<(DecodeMode, String), CliError> { let mode = match args[1].as_str() { "direct" => DecodeMode::Direct, + "rgb" => DecodeMode::Rgb, "adapter" => DecodeMode::Adapter, + "adapter-rgb" => DecodeMode::AdapterRgb, + "verify-rgb" => DecodeMode::VerifyRgb, other => return Err(CliError(format!("Unsupported mode '{other}'. {}", usage()))), }; @@ -93,13 +99,49 @@ fn bench_adapter(input_path: &Path) -> Result> { Ok(((width as u64) << 32) ^ (height as u64) ^ checksum) } +fn bench_rgb(input_path: &Path) -> Result> { + let decoded = decode_path_to_rgb8(input_path)?; + Ok(((decoded.width as u64) << 32) ^ (decoded.height as u64) ^ small_checksum(&decoded.pixels)) +} + +fn decode_adapter_rgb(input_path: &Path) -> Result> { + let _ = register_image_decoder_hooks(); + Ok(ImageReader::open(input_path)?.decode()?.into_rgb8()) +} + +fn bench_adapter_rgb(input_path: &Path) -> Result> { + let decoded = decode_adapter_rgb(input_path)?; + Ok(((decoded.width() as u64) << 32) + ^ (decoded.height() as u64) + ^ small_checksum(decoded.as_raw())) +} + +fn verify_rgb(input_path: &Path) -> Result> { + let direct = decode_path_to_rgb8(input_path)?; + let adapter = decode_adapter_rgb(input_path)?; + if direct.width != adapter.width() + || direct.height != adapter.height() + || direct.pixels != *adapter.as_raw() + { + return Err(Box::new(CliError(format!( + "Direct RGB8 output does not match adapter RGB8 output for {}", + input_path.display() + )))); + } + + Ok(((direct.width as u64) << 32) ^ (direct.height as u64) ^ small_checksum(&direct.pixels)) +} + fn main() -> Result<(), Box> { let (mode, input_path) = parse_args()?; let input_path = Path::new(&input_path); let checksum = match mode { DecodeMode::Direct => bench_direct(input_path)?, + DecodeMode::Rgb => bench_rgb(input_path)?, DecodeMode::Adapter => bench_adapter(input_path)?, + DecodeMode::AdapterRgb => bench_adapter_rgb(input_path)?, + DecodeMode::VerifyRgb => verify_rgb(input_path)?, }; println!("{checksum}"); diff --git a/src/heic-decoder/hevc/ctu.rs b/src/heic-decoder/hevc/ctu.rs index 46e2173..3e2ea6f 100644 --- a/src/heic-decoder/hevc/ctu.rs +++ b/src/heic-decoder/hevc/ctu.rs @@ -8,10 +8,11 @@ use alloc::vec; use alloc::vec::Vec; +#[cfg(feature = "decoder-tracing")] use core::sync::atomic::{AtomicU32, Ordering}; -/// Set to true to enable verbose debug tracing to stderr -const DEBUG_TRACE: bool = false; +/// Verbose decoder tracing is available only in explicitly opted-in builds. +const DEBUG_TRACE: bool = cfg!(feature = "decoder-tracing"); /// Debug print macro gated behind DEBUG_TRACE const macro_rules! debug_trace { @@ -34,6 +35,8 @@ use super::residual::{self, ScanOrder}; use super::sao::SaoMap; use super::slice::{IntraPredMode, PartMode, PredMode, SliceHeader}; use super::transform; +#[cfg(target_arch = "aarch64")] +use super::transform_simd::add_residual_block_neon; use super::transform_simd::add_residual_block_scalar; #[cfg(target_arch = "x86_64")] use super::transform_simd::add_residual_block_v3; @@ -43,11 +46,14 @@ use archmage::incant; type Result = core::result::Result; /// Global SE counter for syntax element tracing +#[cfg(feature = "decoder-tracing")] pub static SE_COUNTER: AtomicU32 = AtomicU32::new(0); +#[cfg(feature = "decoder-tracing")] pub const SE_TRACE_LIMIT: u32 = 0; /// Log a syntax element decode for differential testing. -/// Set SE_TRACE_LIMIT > 0 to enable tracing. +/// Increase `SE_TRACE_LIMIT` to include syntax-element details in tracing. +#[cfg(feature = "decoder-tracing")] #[allow(clippy::absurd_extreme_comparisons)] fn se_trace(name: &str, val: i64, cabac: &CabacDecoder) { let num = SE_COUNTER.fetch_add(1, Ordering::Relaxed); @@ -65,6 +71,10 @@ fn se_trace(name: &str, val: i64, cabac: &CabacDecoder) { let _ = (name, val, cabac); } +#[cfg(not(feature = "decoder-tracing"))] +#[inline(always)] +fn se_trace(_name: &str, _val: i64, _cabac: &CabacDecoder) {} + /// Chroma QP mapping table (H.265 Table 8-10) /// Maps qPi (0-57) to QpC for 8-bit video fn chroma_qp_mapping(qp_i: i32) -> i32 { @@ -287,7 +297,7 @@ impl<'a> SliceContext<'a> { /// Decode all CTUs in the slice pub fn decode_slice(&mut self, frame: &mut DecodedFrame) -> Result<()> { - // Initialize CABAC tracker for debugging + #[cfg(feature = "decoder-tracing")] debug::init_tracker(); let ctb_size = self.sps.ctb_size(); @@ -322,22 +332,23 @@ impl<'a> SliceContext<'a> { let x_ctb = self.ctb_x * ctb_size; let y_ctb = self.ctb_y * ctb_size; - // Track CTU position for debugging - let (byte_pos, _, _) = self.cabac.get_position(); - debug::track_ctu_start(ctu_count, byte_pos); + #[cfg(feature = "decoder-tracing")] + { + let (byte_pos, _, _) = self.cabac.get_position(); + debug::track_ctu_start(ctu_count, byte_pos); - // DEBUG: Print CTU state periodically - if ctu_count.is_multiple_of(50) || ctu_count <= 3 { - let (range, offset) = self.cabac.get_state(); - debug_trace!( - "DEBUG: CTU {} byte={} cabac=({},{}) x={} y={}", - ctu_count, - byte_pos, - range, - offset, - self.ctb_x, - self.ctb_y - ); + if ctu_count.is_multiple_of(50) || ctu_count <= 3 { + let (range, offset) = self.cabac.get_state(); + debug_trace!( + "DEBUG: CTU {} byte={} cabac=({},{}) x={} y={}", + ctu_count, + byte_pos, + range, + offset, + self.ctb_x, + self.ctb_y + ); + } } // Enable debug for CTU 1 (where first large coefficient occurs) self.debug_ctu = ctu_count == 1; @@ -396,6 +407,7 @@ impl<'a> SliceContext<'a> { } } + #[cfg(feature = "decoder-tracing")] if DEBUG_TRACE { debug::print_tracker_summary(); } @@ -802,6 +814,7 @@ impl<'a> SliceContext<'a> { // At minimum size, can be 2Nx2N or NxN let pm = self.decode_part_mode(pred_mode, log2_cb_size)?; // Debug: log part_mode for first CTU (and count NxN) + #[cfg(feature = "decoder-tracing")] if pm == PartMode::PartNxN { static NXN_COUNT: core::sync::atomic::AtomicU32 = core::sync::atomic::AtomicU32::new(0); @@ -1609,7 +1622,7 @@ impl<'a> SliceContext<'a> { size, max_val ), - [v3] + [v3, neon] ); } else { for py in 0..size { diff --git a/src/heic-decoder/hevc/debug.rs b/src/heic-decoder/hevc/debug.rs index b8a8f8c..dd66fa4 100644 --- a/src/heic-decoder/hevc/debug.rs +++ b/src/heic-decoder/hevc/debug.rs @@ -255,17 +255,17 @@ impl CabacTracker { } /// Global tracker instance -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "decoder-tracing"))] static TRACKER: std::sync::Mutex> = std::sync::Mutex::new(None); /// Initialize global tracker -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "decoder-tracing"))] pub fn init_tracker() { *TRACKER.lock().unwrap() = Some(CabacTracker::new()); } /// Record CTU start in global tracker -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "decoder-tracing"))] pub fn track_ctu_start(ctu_idx: u32, byte_pos: usize) { if let Some(tracker) = TRACKER.lock().unwrap().as_mut() { tracker.record_ctu_start(ctu_idx, byte_pos); @@ -273,7 +273,7 @@ pub fn track_ctu_start(ctu_idx: u32, byte_pos: usize) { } /// Record large coefficient in global tracker -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "decoder-tracing"))] pub fn track_large_coeff(byte_pos: usize) { if let Some(tracker) = TRACKER.lock().unwrap().as_mut() { tracker.record_large_coeff(byte_pos); @@ -281,24 +281,28 @@ pub fn track_large_coeff(byte_pos: usize) { } /// Print tracker summary -#[cfg(feature = "std")] +#[cfg(all(feature = "std", feature = "decoder-tracing"))] pub fn print_tracker_summary() { if let Some(tracker) = TRACKER.lock().unwrap().as_ref() { tracker.print_summary(); } } -/// No-op stubs for no_std -#[cfg(not(feature = "std"))] +/// No-op stubs when decoder tracing is disabled. +#[cfg(not(all(feature = "std", feature = "decoder-tracing")))] +#[inline(always)] pub fn init_tracker() {} -/// No-op stub for no_std -#[cfg(not(feature = "std"))] +/// No-op stub when decoder tracing is disabled. +#[cfg(not(all(feature = "std", feature = "decoder-tracing")))] +#[inline(always)] pub fn track_ctu_start(_ctu_idx: u32, _byte_pos: usize) {} -/// No-op stub for no_std -#[cfg(not(feature = "std"))] +/// No-op stub when decoder tracing is disabled. +#[cfg(not(all(feature = "std", feature = "decoder-tracing")))] +#[inline(always)] pub fn track_large_coeff(_byte_pos: usize) {} -/// No-op stub for no_std -#[cfg(not(feature = "std"))] +/// No-op stub when decoder tracing is disabled. +#[cfg(not(all(feature = "std", feature = "decoder-tracing")))] +#[inline(always)] pub fn print_tracker_summary() {} #[cfg(test)] diff --git a/src/heic-decoder/hevc/residual.rs b/src/heic-decoder/hevc/residual.rs index 0a5ab1a..402c74b 100644 --- a/src/heic-decoder/hevc/residual.rs +++ b/src/heic-decoder/hevc/residual.rs @@ -197,6 +197,7 @@ impl CoeffBuffer { /// Decode residual coefficients for a transform unit /// Debug counter to identify specific TU calls +#[cfg(feature = "decoder-tracing")] pub static DEBUG_RESIDUAL_COUNTER: core::sync::atomic::AtomicU32 = core::sync::atomic::AtomicU32::new(0); @@ -213,10 +214,11 @@ pub fn decode_residual( _x0: u32, _y0: u32, ) -> Result<(CoeffBuffer, bool)> { + #[cfg(feature = "decoder-tracing")] DEBUG_RESIDUAL_COUNTER.fetch_add(1, core::sync::atomic::Ordering::Relaxed); // SE trace for differential testing against libde265 - #[cfg(feature = "std")] + #[cfg(feature = "decoder-tracing")] { let se_num = crate::heic_decoder::hevc::ctu::SE_COUNTER .fetch_add(1, core::sync::atomic::Ordering::Relaxed); @@ -726,6 +728,7 @@ pub fn decode_residual( buffer.set(x, y, coeff_values[n]); // Track large coefficients (indicates CABAC desync) + #[cfg(feature = "decoder-tracing")] if coeff_values[n].abs() > 500 { let (byte_pos, _, _) = cabac.get_position(); debug::track_large_coeff(byte_pos); diff --git a/src/heic-decoder/hevc/transform.rs b/src/heic-decoder/hevc/transform.rs index f8eea4c..b719811 100644 --- a/src/heic-decoder/hevc/transform.rs +++ b/src/heic-decoder/hevc/transform.rs @@ -4,11 +4,15 @@ //! - 4x4 Inverse DST (for intra 4x4 luma) //! - 4x4, 8x8, 16x16, 32x32 Inverse DCT //! -//! The 8x8 and 16x16 IDCTs dispatch to AVX2 SIMD via `incant!` when available. +//! The 8x8, 16x16, and 32x32 IDCTs dispatch to AVX2 when available. Flat and +//! scaling-list dequantization also use exact NEON kernels on AArch64. // Transform and inverse quantization for HEVC +#[cfg(target_arch = "aarch64")] +use super::transform_simd::{dequantize_neon, dequantize_scaled_neon}; use super::transform_simd::{ - dequantize_scalar, idct8_scalar, idct16_scalar, idct32_scalar, idst4_scalar, + dequantize_scalar, dequantize_scaled_scalar, idct8_scalar, idct16_scalar, idct32_scalar, + idst4_scalar, }; #[cfg(target_arch = "x86_64")] use super::transform_simd::{dequantize_v3, idct8_v3, idct16_v3, idct32_v3, idst4_v3}; @@ -669,7 +673,7 @@ pub fn dequantize(coeffs: &mut [i16], params: DequantParams) { let add = if shift > 0 { 1 << (shift - 1) } else { 0 }; if shift >= 0 && combined_scale <= 32767 { - incant!(dequantize(coeffs, combined_scale, shift, add), [v3]); + incant!(dequantize(coeffs, combined_scale, shift, add), [v3, neon]); } else if shift >= 0 { // Large scale (high QP at high bit depth): coef * scale overflows // i32, evaluate in i64 (mirrors libde265's int64 fallback). @@ -702,14 +706,19 @@ pub fn dequantize_scaled(coeffs: &mut [i16], params: DequantParams, scaling_matr let add = if bd_shift > 0 { 1 << (bd_shift - 1) } else { 0 }; if bd_shift >= 0 { - for (i, coef) in coeffs.iter_mut().enumerate() { - let m = scaling_matrix.get(i).copied().unwrap_or(16) as i64; - // i64: m (up to 255) * levelScale (72) << qp_per overflows i32 - // already at moderate QPs (libde265 uses int64 here too). - let value = - (*coef as i64 * m * level_scale as i64 * (1i64 << qp_per) + add as i64) >> bd_shift; - *coef = value.clamp(-32768, 32767) as i16; - } + // i64: m (up to 255) * levelScale (72) << qp_per overflows i32 + // already at moderate QPs (libde265 uses int64 here too). + let combined_scale = level_scale * (1 << qp_per); + incant!( + dequantize_scaled( + coeffs, + scaling_matrix, + combined_scale, + bd_shift, + i64::from(add) + ), + [neon] + ); } else { let neg_shift = -bd_shift; for (i, coef) in coeffs.iter_mut().enumerate() { diff --git a/src/heic-decoder/hevc/transform_simd.rs b/src/heic-decoder/hevc/transform_simd.rs index 9739a32..3b6a296 100644 --- a/src/heic-decoder/hevc/transform_simd.rs +++ b/src/heic-decoder/hevc/transform_simd.rs @@ -6,6 +6,18 @@ use archmage::prelude::*; +#[cfg(target_arch = "aarch64")] +use core::arch::aarch64::{ + vaddq_s32, vaddq_s64, vcombine_s16, vcombine_s32, vdup_n_s32, vdupq_n_s16, vdupq_n_s32, + vdupq_n_s64, vget_high_s16, vget_high_s32, vget_low_s16, vget_low_s32, vmaxq_s16, vminq_s16, + vmovl_s16, vmull_s32, vmulq_n_s32, vqaddq_s16, vqmovn_s32, vqmovn_s64, vreinterpretq_s16_u16, + vreinterpretq_u16_s16, vshlq_s32, vshlq_s64, +}; +#[cfg(target_arch = "aarch64")] +use safe_unaligned_simd::aarch64::{ + vld1_s16, vld1q_s16, vld1q_s32, vld1q_u16, vst1_s16, vst1q_s16, vst1q_u16, +}; + #[cfg(target_arch = "x86_64")] use safe_unaligned_simd::x86_64::{ _mm_loadu_si128, _mm_storeu_si128, _mm256_loadu_si256, _mm256_storeu_si256, @@ -1349,6 +1361,55 @@ pub(crate) fn add_residual_block_v3( } } +/// AArch64 NEON residual add. When prediction samples fit signed i16, signed +/// saturating addition followed by the codec-range clamp is equivalent to the +/// scalar i32 calculation even when the mathematical sum exceeds i16. Wider +/// non-profile inputs retain the scalar path. +#[cfg(target_arch = "aarch64")] +#[allow(clippy::too_many_arguments)] +#[arcane] +pub(crate) fn add_residual_block_neon( + _token: NeonToken, + plane: &mut [u16], + stride: usize, + x0: usize, + y0: usize, + residual: &[i16], + size: usize, + max_val: i32, +) { + if max_val > i32::from(i16::MAX) { + add_residual_block_scalar(ScalarToken, plane, stride, x0, y0, residual, size, max_val); + return; + } + + let zero = vdupq_n_s16(0); + let max_v = vdupq_n_s16(max_val as i16); + + for py in 0..size { + let row_start = (y0 + py) * stride + x0; + let row = &mut plane[row_start..row_start + size]; + let res_row = &residual[py * size..(py + 1) * size]; + let chunks = size / 8; + for chunk in 0..chunks { + let offset = chunk * 8; + let prediction = vld1q_u16((&row[offset..offset + 8]).try_into().unwrap()); + let residual = vld1q_s16((&res_row[offset..offset + 8]).try_into().unwrap()); + let sum = vqaddq_s16(vreinterpretq_s16_u16(prediction), residual); + let clamped = vminq_s16(vmaxq_s16(sum, zero), max_v); + vst1q_u16( + (&mut row[offset..offset + 8]).try_into().unwrap(), + vreinterpretq_u16_s16(clamped), + ); + } + for i in (chunks * 8)..size { + let prediction = i32::from(row[i]); + let residual = i32::from(res_row[i]); + row[i] = (prediction + residual).clamp(0, max_val) as u16; + } + } +} + /// Scalar fallback for residual block add #[allow(clippy::too_many_arguments)] pub(crate) fn add_residual_block_scalar( @@ -1430,6 +1491,38 @@ pub(crate) fn dequantize_v3( } } +/// AArch64 NEON flat dequantization, four coefficients at a time. Widening to +/// i32 before multiplication preserves the scalar overflow and rounding +/// behavior, and the saturating narrow implements the final codec clamp. +#[cfg(target_arch = "aarch64")] +#[arcane] +pub(crate) fn dequantize_neon( + _token: NeonToken, + coeffs: &mut [i16], + combined_scale: i32, + shift: i32, + add: i32, +) { + let shift_v = vdupq_n_s32(-shift); + let add_v = vdupq_n_s32(add); + let chunks = coeffs.len() / 4; + for chunk in 0..chunks { + let offset = chunk * 4; + let source = vld1_s16((&coeffs[offset..offset + 4]).try_into().unwrap()); + let product = vmulq_n_s32(vmovl_s16(source), combined_scale); + let shifted = vshlq_s32(vaddq_s32(product, add_v), shift_v); + vst1_s16( + (&mut coeffs[offset..offset + 4]).try_into().unwrap(), + vqmovn_s32(shifted), + ); + } + + for coefficient in coeffs.iter_mut().skip(chunks * 4) { + let value = (i32::from(*coefficient) * combined_scale + add) >> shift; + *coefficient = value.clamp(-32768, 32767) as i16; + } +} + /// Scalar fallback for dequantize pub(crate) fn dequantize_scalar( _token: ScalarToken, @@ -1443,3 +1536,181 @@ pub(crate) fn dequantize_scalar( *coef = value.clamp(-32768, 32767) as i16; } } + +/// AArch64 NEON scaling-list dequantization. The coefficient×matrix product +/// fits i32; widening multiplication by the QP scale then preserves the +/// scalar path's required i64 range before exact rounding and saturation. +#[cfg(target_arch = "aarch64")] +#[arcane] +pub(crate) fn dequantize_scaled_neon( + _token: NeonToken, + coeffs: &mut [i16], + scaling_matrix: &[u8], + combined_scale: i32, + shift: i32, + add: i64, +) { + let chunks = coeffs.len().min(scaling_matrix.len()) / 4; + let scale_v = vdup_n_s32(combined_scale); + let add_v = vdupq_n_s64(add); + let shift_v = vdupq_n_s64(i64::from(-shift)); + + for chunk in 0..chunks { + let offset = chunk * 4; + let coefficients = vmovl_s16(vld1_s16((&coeffs[offset..offset + 4]).try_into().unwrap())); + let matrix_values = [ + i32::from(scaling_matrix[offset]), + i32::from(scaling_matrix[offset + 1]), + i32::from(scaling_matrix[offset + 2]), + i32::from(scaling_matrix[offset + 3]), + ]; + let coefficient_matrix = + core::arch::aarch64::vmulq_s32(coefficients, vld1q_s32(&matrix_values)); + let product_low = vmull_s32(vget_low_s32(coefficient_matrix), scale_v); + let product_high = vmull_s32(vget_high_s32(coefficient_matrix), scale_v); + let rounded_low = vshlq_s64(vaddq_s64(product_low, add_v), shift_v); + let rounded_high = vshlq_s64(vaddq_s64(product_high, add_v), shift_v); + let narrowed = vcombine_s32(vqmovn_s64(rounded_low), vqmovn_s64(rounded_high)); + vst1_s16( + (&mut coeffs[offset..offset + 4]).try_into().unwrap(), + vqmovn_s32(narrowed), + ); + } + + for (index, coefficient) in coeffs.iter_mut().enumerate().skip(chunks * 4) { + let matrix = i64::from(scaling_matrix.get(index).copied().unwrap_or(16)); + let value = (i64::from(*coefficient) * matrix * i64::from(combined_scale) + add) >> shift; + *coefficient = value.clamp(-32768, 32767) as i16; + } +} + +pub(crate) fn dequantize_scaled_scalar( + _token: ScalarToken, + coeffs: &mut [i16], + scaling_matrix: &[u8], + combined_scale: i32, + shift: i32, + add: i64, +) { + for (index, coefficient) in coeffs.iter_mut().enumerate() { + let matrix = i64::from(scaling_matrix.get(index).copied().unwrap_or(16)); + let value = (i64::from(*coefficient) * matrix * i64::from(combined_scale) + add) >> shift; + *coefficient = value.clamp(-32768, 32767) as i16; + } +} + +#[cfg(all(test, target_arch = "aarch64"))] +mod neon_tests { + use super::*; + + fn next_random(state: &mut u64) -> u32 { + *state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + (*state >> 32) as u32 + } + + #[test] + fn neon_residual_add_matches_scalar_for_all_block_sizes_and_bit_depths() { + let neon = NeonToken::summon().expect("AArch64 must provide NEON"); + let mut state = 0x4f1b_d38a_9207_c6e5; + for size in [4_usize, 8, 16, 32] { + for bit_depth in [8_u8, 10, 12, 14, 16] { + let stride = size + 5; + let rows = size + 2; + let max_value = (1_i32 << bit_depth) - 1; + let mut scalar_plane = (0..stride * rows) + .map(|_| (next_random(&mut state) % (max_value as u32 + 1)) as u16) + .collect::>(); + let mut neon_plane = scalar_plane.clone(); + let mut residual = (0..size * size) + .map(|_| next_random(&mut state) as i16) + .collect::>(); + residual[0] = i16::MIN; + let last_residual = residual.len() - 1; + residual[last_residual] = i16::MAX; + + add_residual_block_scalar( + ScalarToken, + &mut scalar_plane, + stride, + 3, + 1, + &residual, + size, + max_value, + ); + add_residual_block_neon( + neon, + &mut neon_plane, + stride, + 3, + 1, + &residual, + size, + max_value, + ); + assert_eq!(neon_plane, scalar_plane, "size={size}, depth={bit_depth}"); + } + } + } + + #[test] + fn neon_flat_dequantization_matches_scalar_with_unaligned_tails() { + let neon = NeonToken::summon().expect("AArch64 must provide NEON"); + let mut state = 0x8675_309d_4e21_bacf; + for length in [1_usize, 4, 7, 16, 63, 256, 1023] { + for (combined_scale, shift) in [(40, 0), (720, 3), (18_432, 8), (32_767, 10)] { + let mut scalar = (0..length) + .map(|_| next_random(&mut state) as i16) + .collect::>(); + let mut simd = scalar.clone(); + let add = if shift == 0 { 0 } else { 1 << (shift - 1) }; + dequantize_scalar(ScalarToken, &mut scalar, combined_scale, shift, add); + dequantize_neon(neon, &mut simd, combined_scale, shift, add); + assert_eq!( + simd, scalar, + "length={length}, scale={combined_scale}, shift={shift}" + ); + } + } + } + + #[test] + fn neon_scaling_list_dequantization_matches_i64_scalar() { + let neon = NeonToken::summon().expect("AArch64 must provide NEON"); + let mut state = 0xa913_57fc_d204_68be; + for (length, matrix_length) in [(4_usize, 4_usize), (7, 5), (64, 64), (255, 255)] { + for (combined_scale, shift) in [(40, 3), (720, 7), (18_432, 14)] { + let mut scalar = (0..length) + .map(|_| next_random(&mut state) as i16) + .collect::>(); + let mut simd = scalar.clone(); + let scaling_matrix = (0..matrix_length) + .map(|_| next_random(&mut state) as u8) + .collect::>(); + let add = 1_i64 << (shift - 1); + dequantize_scaled_scalar( + ScalarToken, + &mut scalar, + &scaling_matrix, + combined_scale, + shift, + add, + ); + dequantize_scaled_neon( + neon, + &mut simd, + &scaling_matrix, + combined_scale, + shift, + add, + ); + assert_eq!( + simd, scalar, + "length={length}, scale={combined_scale}, shift={shift}" + ); + } + } + } +} diff --git a/src/lib.rs b/src/lib.rs index f6addb8..aabd03e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -30,6 +30,8 @@ use rav1d::{ Dav1dResult, Rav1dError, dav1d_close, dav1d_data_create, dav1d_data_unref, dav1d_default_settings, dav1d_get_picture, dav1d_open, dav1d_picture_unref, dav1d_send_data, }; +#[cfg(feature = "parallel-grid")] +use rayon::prelude::*; use scuffle_h265::NALUnitType; use source::{ FileSource, RandomAccessSource, SourceReadError, TempFileSpoolOptions, TempFileSpoolSource, @@ -41,6 +43,8 @@ use std::fmt::{Display, Formatter}; use std::fs::File; use std::io::{BufRead, BufWriter, Read}; use std::mem::MaybeUninit; +#[cfg(feature = "parallel-grid")] +use std::mem::size_of; use std::path::{Path, PathBuf}; use std::ptr::{self, NonNull}; @@ -610,6 +614,20 @@ pub struct DecodedRgbaImage { pub icc_profile: Option>, } +/// Decoded RGB8 image buffer for consumers that intentionally discard alpha. +/// +/// Pixels are byte-for-byte equivalent to converting this decoder's final +/// RGBA output through `image::DynamicImage::into_rgb8`, but the RGB path does +/// not allocate or write an unused alpha channel. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct DecodedRgbImage { + pub width: u32, + pub height: u32, + pub source_bit_depth: u8, + pub pixels: Vec, + pub icc_profile: Option>, +} + #[cfg(feature = "image-integration")] #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) struct DecodedRgbaLayout { @@ -4160,45 +4178,285 @@ fn decode_primary_heic_grid_to_image( }); } - let mut output = None; - let mut tile_width = 0; - let mut tile_height = 0; + let first_tile = decode_heic_grid_tile_to_image(&grid_data.tiles[0])?; + let tile_width = first_tile.width; + let tile_height = first_tile.height; + let mut output = create_decoded_heic_grid_output(descriptor, &first_tile)?; - for row in 0..rows { - for column in 0..columns { - let tile_index = row - .checked_mul(columns) - .and_then(|idx| idx.checked_add(column)) - .ok_or_else(|| DecodeHeicError::InvalidDecodedFrame { - detail: format!("grid tile index overflow at row={row}, column={column}"), - })?; - let tile = decode_heic_grid_tile_to_image(&grid_data.tiles[tile_index])?; + for_each_decoded_heic_grid_tile(grid_data, first_tile, |tile_index, tile| { + let row = tile_index / columns; + let column = tile_index % columns; + paste_decoded_heic_grid_tile( + &tile, + &mut output, + tile_width, + tile_height, + row, + column, + tile_index, + ) + })?; + + Ok(output) +} + +#[cfg(feature = "parallel-grid")] +const MAX_GRID_TILE_DECODE_THREADS: usize = 8; +#[cfg(feature = "parallel-grid")] +const GRID_TILE_DECODE_MEMORY_BUDGET: u64 = 64 * 1024 * 1024; + +#[cfg(feature = "parallel-grid")] +#[derive(Default)] +struct HeicGridTileStreamMemoryEstimate { + normalized_stream_bytes: u64, + rbsp_bytes: u64, + emulation_prevention_position_bytes: u64, + nal_count: u64, +} + +#[cfg(feature = "parallel-grid")] +impl HeicGridTileStreamMemoryEstimate { + fn add_nal(&mut self, nal_unit: &[u8]) -> Result<(), DecodeHeicError> { + let nal_size = + u32::try_from(nal_unit.len()).map_err(|_| DecodeHeicError::NalUnitTooLarge { + nal_size: nal_unit.len(), + })?; + self.normalized_stream_bytes = self + .normalized_stream_bytes + .saturating_add(u64::from(nal_size)) + .saturating_add(4); + self.nal_count = self.nal_count.saturating_add(1); + + let Some(rbsp) = nal_unit.get(2..) else { + return Ok(()); + }; + self.rbsp_bytes = self.rbsp_bytes.saturating_add(rbsp.len() as u64); + + let removed_bytes = count_hevc_emulation_prevention_bytes(rbsp); + if removed_bytes > 0 { + // `skipped_byte_positions` grows from an empty Vec. Twice + // the populated size plus four elements covers its amortized + // growth, including the initial small allocation. + let position_capacity = removed_bytes.saturating_mul(2).saturating_add(4); + self.emulation_prevention_position_bytes = self + .emulation_prevention_position_bytes + .saturating_add(position_capacity.saturating_mul(size_of::() as u64)); + } + Ok(()) + } + + fn estimated_decoder_bytes(&self) -> u64 { + // Stream assembly appends to a growing Vec, so allow up to twice its + // populated length. Each parsed NAL then owns an RBSP Vec whose + // requested capacity equals the transmitted payload length. + let stream_and_rbsp_bytes = self + .normalized_stream_bytes + .saturating_mul(2) + .saturating_add(self.rbsp_bytes); + let backend_nal_metadata_bytes = conservative_vec_storage_bytes( + self.nal_count, + size_of::>(), + ); + let length_prefixed_nal_metadata_bytes = conservative_vec_storage_bytes( + self.nal_count, + size_of::>(), + ); + + stream_and_rbsp_bytes + .saturating_add(self.emulation_prevention_position_bytes) + .saturating_add(backend_nal_metadata_bytes) + .saturating_add(length_prefixed_nal_metadata_bytes) + } +} + +#[cfg(feature = "parallel-grid")] +fn conservative_vec_storage_bytes(element_count: u64, element_size: usize) -> u64 { + if element_count == 0 { + return 0; + } + element_count + .saturating_mul(2) + .saturating_add(4) + .saturating_mul(element_size as u64) +} + +#[cfg(feature = "parallel-grid")] +fn count_hevc_emulation_prevention_bytes(bytes: &[u8]) -> u64 { + let mut count = 0_u64; + let mut cursor = 0_usize; + while cursor < bytes.len() { + if cursor + 2 < bytes.len() + && bytes[cursor] == 0 + && bytes[cursor + 1] == 0 + && bytes[cursor + 2] == 3 + { + count = count.saturating_add(1); + cursor += 3; + } else { + cursor += 1; + } + } + count +} + +#[cfg(feature = "parallel-grid")] +fn estimate_heic_grid_sps_decode_bytes( + sps: &heic_decoder::hevc::params::Sps, +) -> Result { + let _ = hevc_metadata_from_sps(sps)?; + let allocation_layout = heic_layout_from_sps_chroma_array_type(sps.chroma_format_idc)?; + let working_bytes_per_pixel = match allocation_layout { + HeicPixelLayout::Yuv400 => 4_u64, + HeicPixelLayout::Yuv420 => 6, + HeicPixelLayout::Yuv422 => 8, + HeicPixelLayout::Yuv444 => 12, + }; + Ok(u64::from(sps.pic_width_in_luma_samples) + .saturating_mul(u64::from(sps.pic_height_in_luma_samples)) + .saturating_mul(working_bytes_per_pixel)) +} + +/// Estimate the normalized decoder stream plus retained YUV output and decoder +/// working storage for one tile. Raw SPS geometry is deliberately used here: +/// a tile-level clean aperture may make the final tile small, but the decoder +/// must still materialize the uncropped frame first. +#[cfg(feature = "parallel-grid")] +fn estimate_heic_grid_tile_decode_bytes( + tile: &isobmff::HeicGridTileItemData, +) -> Result { + // The backend currently retains the last SPS it parses. Use the largest + // allocation advertised by any SPS instead: this remains safe for files + // carrying unused or in-stream replacement parameter sets and avoids + // coupling the memory bound to the backend's SPS-selection details. + let mut decoded_bytes = None::; + let mut stream_memory = HeicGridTileStreamMemoryEstimate::default(); + walk_hevc_nals_from_hvcc_or_payload(&tile.hvcc, &tile.payload, |offset, nal_unit| { + stream_memory.add_nal(nal_unit)?; + let unit = LengthPrefixedHevcNalUnit { + offset, + bytes: nal_unit, + }; + if unit.nal_unit_type() == Some(NALUnitType::SpsNut) { + let sps = hevc_sps_from_nal(nal_unit, offset)?; + let candidate_bytes = estimate_heic_grid_sps_decode_bytes(&sps)?; + decoded_bytes = Some(decoded_bytes.unwrap_or_default().max(candidate_bytes)); + } + Ok(false) + })?; + + let decoded_bytes = decoded_bytes.ok_or(DecodeHeicError::MissingSpsNalUnit)?; + Ok(decoded_bytes + .saturating_add(stream_memory.estimated_decoder_bytes()) + .max(1)) +} + +/// Choose a conservative number of simultaneously decoded tiles from each +/// tile's own preflight estimate. An estimate failure stops the parallel +/// window before that tile so normal row-major decoding still selects the +/// user-visible error. +#[cfg(feature = "parallel-grid")] +fn heic_grid_tile_decode_window_from_estimates( + estimates: impl IntoIterator>, + available_threads: usize, +) -> usize { + let available_threads = available_threads.max(1); + let mut window = 0_usize; + let mut estimated_bytes = 0_u64; + for tile_bytes in estimates.into_iter().take(available_threads) { + let Some(tile_bytes) = tile_bytes else { + break; + }; + let next_estimated_bytes = estimated_bytes.saturating_add(tile_bytes.max(1)); + if window > 0 && next_estimated_bytes > GRID_TILE_DECODE_MEMORY_BUDGET { + break; + } + window += 1; + estimated_bytes = next_estimated_bytes; + if estimated_bytes >= GRID_TILE_DECODE_MEMORY_BUDGET { + break; + } + } + window.max(1) +} + +/// Choose a conservative number of simultaneously decoded tiles. Large, +/// malformed, or unusual tiles stay sequential instead of inheriting the +/// first grid tile's geometry and multiplying peak memory. +#[cfg(feature = "parallel-grid")] +fn heic_grid_tile_decode_window(remaining_tiles: &[isobmff::HeicGridTileItemData]) -> usize { + if remaining_tiles.len() <= 1 || cfg!(feature = "decoder-tracing") { + return 1; + } + + let available_threads = rayon::current_num_threads() + .min(MAX_GRID_TILE_DECODE_THREADS) + .min(remaining_tiles.len()); + heic_grid_tile_decode_window_from_estimates( + remaining_tiles + .iter() + .map(|tile| estimate_heic_grid_tile_decode_bytes(tile).ok()), + available_threads, + ) +} + +/// Decode grid tiles in bounded batches, but deliver them to the caller in +/// row-major order. Keeping validation and paste work on the caller thread +/// preserves deterministic errors and output while independent HEVC payloads +/// use the available cores. +fn for_each_decoded_heic_grid_tile( + grid_data: &isobmff::HeicGridPrimaryItemData, + first_tile: DecodedHeicImage, + mut consume: impl FnMut(usize, DecodedHeicImage) -> Result<(), E>, +) -> Result<(), E> +where + E: From, +{ + let remaining_tiles = &grid_data.tiles[1..]; + + consume(0, first_tile)?; - if tile_index == 0 { - tile_width = tile.width; - tile_height = tile.height; - output = Some(create_decoded_heic_grid_output(descriptor, &tile)?); + if remaining_tiles.is_empty() { + return Ok(()); + } + + #[cfg(feature = "parallel-grid")] + { + let mut next_tile_index = 1_usize; + while next_tile_index < grid_data.tiles.len() { + let undecoded_tiles = &grid_data.tiles[next_tile_index..]; + let window = heic_grid_tile_decode_window(undecoded_tiles); + if window > 1 { + let decoded = undecoded_tiles[..window] + .par_iter() + .map(decode_heic_grid_tile_to_image) + .collect::>(); + + for (offset, tile) in decoded.into_iter().enumerate() { + consume(next_tile_index + offset, tile.map_err(E::from)?)?; + } + next_tile_index += window; + } else { + consume( + next_tile_index, + decode_heic_grid_tile_to_image(&grid_data.tiles[next_tile_index]) + .map_err(E::from)?, + )?; + next_tile_index += 1; } + } + Ok(()) + } - paste_decoded_heic_grid_tile( - &tile, - output - .as_mut() - .ok_or_else(|| DecodeHeicError::InvalidDecodedFrame { - detail: "grid output was not initialized".to_string(), - })?, - tile_width, - tile_height, - row, - column, - tile_index, + #[cfg(not(feature = "parallel-grid"))] + { + for (offset, tile) in remaining_tiles.iter().enumerate() { + consume( + offset + 1, + decode_heic_grid_tile_to_image(tile).map_err(E::from)?, )?; } + Ok(()) } - - output.ok_or_else(|| DecodeHeicError::InvalidDecodedFrame { - detail: "grid tile list cannot be empty".to_string(), - }) } fn decode_heic_grid_tile_to_image( @@ -4321,6 +4579,143 @@ fn decode_primary_heic_grid_to_rgba_image( } } +fn decode_primary_heic_grid_to_rgb8_image( + grid_data: &isobmff::HeicGridPrimaryItemData, + transforms: &[isobmff::PrimaryItemTransformProperty], + icc_profile: Option>, +) -> Result { + validate_heic_grid_descriptor_and_tile_count(grid_data)?; + let (first_tile, source_bit_depth) = decode_and_validate_heic_grid_first_tile(grid_data)?; + let reference = HeicGridTileReference::from_first_tile(&first_tile, &grid_data.colr); + let transform_plan = RgbaTransformPlan::from_primary_transforms( + grid_data.descriptor.output_width, + grid_data.descriptor.output_height, + transforms, + )?; + let destination_width = usize::try_from(transform_plan.destination_width).map_err(|_| { + DecodeHeicError::InvalidDecodedFrame { + detail: "transformed RGB grid width cannot be represented".to_string(), + } + })?; + let source_width = usize::try_from(grid_data.descriptor.output_width).map_err(|_| { + DecodeHeicError::InvalidDecodedFrame { + detail: "RGB grid source width cannot be represented".to_string(), + } + })?; + let source_height = usize::try_from(grid_data.descriptor.output_height).map_err(|_| { + DecodeHeicError::InvalidDecodedFrame { + detail: "RGB grid source height cannot be represented".to_string(), + } + })?; + let mut output = vec![ + 0_u8; + checked_interleaved_sample_count( + transform_plan.destination_width, + transform_plan.destination_height, + 3, + )? + ]; + + if !heic_grid_tiles_cover_descriptor(&grid_data.descriptor, &reference) { + let gap_pixel = heic_grid_gap_rgb8_pixel(&reference)?; + for pixel in output.chunks_exact_mut(3) { + pixel.copy_from_slice(&gap_pixel); + } + } + + let columns = usize::from(grid_data.descriptor.columns); + let mut tile_pixels = Vec::new(); + for_each_decoded_heic_grid_tile( + grid_data, + first_tile, + |tile_index, mut tile| -> Result<(), DecodeError> { + validate_decoded_heic_grid_tile_reference(&tile, &reference, tile_index)?; + tile.ycbcr_range = reference.conversion_ycbcr_range; + tile.ycbcr_matrix = reference.conversion_ycbcr_matrix; + convert_heic_to_rgb8_into(&tile, &mut tile_pixels)?; + + let tile_width = + usize::try_from(tile.width).map_err(|_| DecodeHeicError::InvalidDecodedFrame { + detail: format!("RGB grid tile width {} cannot be represented", tile.width), + })?; + let tile_height = + usize::try_from(tile.height).map_err(|_| DecodeHeicError::InvalidDecodedFrame { + detail: format!("RGB grid tile height {} cannot be represented", tile.height), + })?; + let expected_tile_samples = tile_width + .checked_mul(tile_height) + .and_then(|pixels| pixels.checked_mul(3)) + .ok_or_else(|| DecodeHeicError::InvalidDecodedFrame { + detail: "RGB grid tile sample count overflow".to_string(), + })?; + if tile_pixels.len() != expected_tile_samples { + return Err(DecodeHeicError::InvalidDecodedFrame { + detail: format!( + "RGB grid tile has {} samples, expected {expected_tile_samples}", + tile_pixels.len() + ), + } + .into()); + } + + let row = tile_index / columns; + let column = tile_index % columns; + let (x_origin, y_origin) = + heic_grid_tile_origin(reference.tile_width, reference.tile_height, row, column)?; + validate_heic_grid_tile_origin_alignment(reference.layout, x_origin, y_origin)?; + let x_origin = + usize::try_from(x_origin).map_err(|_| DecodeHeicError::InvalidDecodedFrame { + detail: "RGB grid tile x-origin cannot be represented".to_string(), + })?; + let y_origin = + usize::try_from(y_origin).map_err(|_| DecodeHeicError::InvalidDecodedFrame { + detail: "RGB grid tile y-origin cannot be represented".to_string(), + })?; + + for tile_y in 0..tile_height { + let source_y = y_origin.checked_add(tile_y).ok_or_else(|| { + DecodeHeicError::InvalidDecodedFrame { + detail: "RGB grid tile y-coordinate overflow".to_string(), + } + })?; + if source_y >= source_height { + break; + } + for tile_x in 0..tile_width { + let source_x = x_origin.checked_add(tile_x).ok_or_else(|| { + DecodeHeicError::InvalidDecodedFrame { + detail: "RGB grid tile x-coordinate overflow".to_string(), + } + })?; + if source_x >= source_width { + break; + } + let Some((destination_x, destination_y)) = + transform_plan.map_source_pixel(source_x, source_y)? + else { + continue; + }; + let source_sample = (tile_y * tile_width + tile_x) * 3; + let destination_sample = + (destination_y * destination_width + destination_x) * 3; + output[destination_sample..destination_sample + 3] + .copy_from_slice(&tile_pixels[source_sample..source_sample + 3]); + } + } + + Ok(()) + }, + )?; + + Ok(DecodedRgbImage { + width: transform_plan.destination_width, + height: transform_plan.destination_height, + source_bit_depth, + pixels: output, + icc_profile, + }) +} + fn validate_heic_grid_descriptor_and_tile_count( grid_data: &isobmff::HeicGridPrimaryItemData, ) -> Result<(), DecodeHeicError> { @@ -4584,6 +4979,17 @@ fn heic_grid_gap_rgba_pixel( reference: &HeicGridTileReference, convert_tile: fn(&DecodedHeicImage, &mut Vec) -> Result<(), DecodeHeicError>, ) -> Result<[T; 4], DecodeHeicError> { + heic_grid_gap_pixel(reference, convert_tile) +} + +fn heic_grid_gap_rgb8_pixel(reference: &HeicGridTileReference) -> Result<[u8; 3], DecodeHeicError> { + heic_grid_gap_pixel(reference, convert_heic_to_rgb8_into) +} + +fn heic_grid_gap_pixel( + reference: &HeicGridTileReference, + convert_tile: fn(&DecodedHeicImage, &mut Vec) -> Result<(), DecodeHeicError>, +) -> Result<[T; CHANNELS], DecodeHeicError> { let zero_plane = HeicPlane { width: 1, height: 1, @@ -4608,10 +5014,10 @@ fn heic_grid_gap_rgba_pixel( }; let mut pixel = Vec::new(); convert_tile(&zero_image, &mut pixel)?; - <[T; 4]>::try_from(pixel.as_slice()).map_err(|_| DecodeHeicError::InvalidDecodedFrame { + <[T; CHANNELS]>::try_from(pixel.as_slice()).map_err(|_| DecodeHeicError::InvalidDecodedFrame { detail: format!( - "grid gap pixel conversion produced {} samples, expected 4", - pixel.len() + "grid gap pixel conversion produced {} samples, expected {CHANNELS}", + pixel.len(), ), }) } @@ -4667,18 +5073,8 @@ fn for_each_heic_grid_tile_rgba( mut paste: impl FnMut(&DecodedHeicImage, &[T], u32, u32) -> Result<(), DecodeError>, ) -> Result<(), DecodeError> { let columns = usize::from(grid_data.descriptor.columns); - let mut first_tile = Some(first_tile); let mut tile_pixels = Vec::new(); - for tile_index in 0..grid_data.tiles.len() { - let mut tile = if tile_index == 0 { - first_tile - .take() - .ok_or_else(|| DecodeHeicError::InvalidDecodedFrame { - detail: "first grid tile was already consumed".to_string(), - })? - } else { - decode_heic_grid_tile_to_image(&grid_data.tiles[tile_index])? - }; + for_each_decoded_heic_grid_tile(grid_data, first_tile, |tile_index, mut tile| { validate_decoded_heic_grid_tile_reference(&tile, reference, tile_index)?; tile.ycbcr_range = reference.conversion_ycbcr_range; tile.ycbcr_matrix = reference.conversion_ycbcr_matrix; @@ -4689,9 +5085,8 @@ fn for_each_heic_grid_tile_rgba( heic_grid_tile_origin(reference.tile_width, reference.tile_height, row, column)?; validate_heic_grid_tile_origin_alignment(reference.layout, x_origin, y_origin)?; paste(&tile, &tile_pixels, x_origin, y_origin)?; - } - - Ok(()) + Ok(()) + }) } fn paste_heic_grid_tiles_to_rgba( @@ -4822,7 +5217,6 @@ struct RgbaOrientationTransform { /// transform) makes clean-aperture clipping explicit and preserves the exact /// transform order used by `apply_primary_item_transforms_rgba`. #[derive(Clone, Debug)] -#[cfg(feature = "image-integration")] struct RgbaTransformPlan { source_width: u32, source_height: u32, @@ -4832,7 +5226,6 @@ struct RgbaTransformPlan { } #[derive(Clone, Copy, Debug)] -#[cfg(feature = "image-integration")] enum ResolvedRgbaTransformStep { CleanAperture { left: u32, @@ -4852,7 +5245,6 @@ enum ResolvedRgbaTransformStep { }, } -#[cfg(feature = "image-integration")] impl RgbaTransformPlan { fn from_primary_transforms( width: u32, @@ -6467,13 +6859,25 @@ fn hevc_metadata_from_sps_nal( nal_bytes: &[u8], nal_offset: usize, ) -> Result { - let sps = heic_decoder::hevc::bitstream::parse_single_nal(nal_bytes) + let sps = hevc_sps_from_nal(nal_bytes, nal_offset)?; + hevc_metadata_from_sps(&sps) +} + +fn hevc_sps_from_nal( + nal_bytes: &[u8], + nal_offset: usize, +) -> Result { + heic_decoder::hevc::bitstream::parse_single_nal(nal_bytes) .and_then(|nal| heic_decoder::hevc::params::parse_sps(&nal.payload)) .map_err(|err| DecodeHeicError::SpsParseFailed { offset: nal_offset, detail: err.to_string(), - })?; + }) +} +fn hevc_metadata_from_sps( + sps: &heic_decoder::hevc::params::Sps, +) -> Result { let (sub_width_c, sub_height_c) = match sps.chroma_format_idc { 1 => (2u32, 2u32), 2 => (2, 1), @@ -6514,14 +6918,14 @@ fn hevc_metadata_from_sps_nal( }) } -/// SPS-derived metadata from the hvcC parameter-set arrays alone, or `None` +/// Parse the SPS from the hvcC parameter-set arrays alone, or return `None` /// when the arrays carry no SPS (`hev1` items may keep parameter sets only /// in-stream). Selection matches the assembled-stream scan: first NAL whose /// own header says SPS, in hvcC array order. #[cfg(feature = "image-integration")] -fn hevc_metadata_from_hvcc_nal_arrays( +fn hevc_sps_from_hvcc_nal_arrays( hvcc: &isobmff::HevcDecoderConfigurationBox, -) -> Result, DecodeHeicError> { +) -> Result, DecodeHeicError> { for nal_array in &hvcc.nal_arrays { for nal_unit in &nal_array.nal_units { let unit = LengthPrefixedHevcNalUnit { @@ -6531,40 +6935,79 @@ fn hevc_metadata_from_hvcc_nal_arrays( if unit.nal_unit_type() != Some(NALUnitType::SpsNut) { continue; } - return hevc_metadata_from_sps_nal(nal_unit, 0).map(Some); + return hevc_sps_from_nal(nal_unit, 0).map(Some); } } Ok(None) } -/// SPS-derived metadata without assembling a decoder stream: prefer the hvcC +/// Visit hvcC and in-stream NAL units in the same order used to assemble the +/// decoder stream. Returning `true` stops the walk after the current unit. +#[cfg(feature = "parallel-grid")] +fn walk_hevc_nals_from_hvcc_or_payload( + hvcc: &isobmff::HevcDecoderConfigurationBox, + payload: &[u8], + mut visit: impl FnMut(usize, &[u8]) -> Result, +) -> Result<(), DecodeHeicError> { + let nal_length_size = hvcc.nal_length_size; + if !(1..=4).contains(&nal_length_size) { + return Err(DecodeHeicError::InvalidNalLengthSize { nal_length_size }); + } + for nal_array in &hvcc.nal_arrays { + for nal_unit in &nal_array.nal_units { + if visit(0, nal_unit)? { + return Ok(()); + } + } + } + walk_length_prefixed_payload_nals(payload, usize::from(nal_length_size), visit) +} + +#[cfg(feature = "image-integration")] +fn hevc_metadata_from_hvcc_nal_arrays( + hvcc: &isobmff::HevcDecoderConfigurationBox, +) -> Result, DecodeHeicError> { + hevc_sps_from_hvcc_nal_arrays(hvcc)? + .map(|sps| hevc_metadata_from_sps(&sps)) + .transpose() +} + +/// Parse the SPS without assembling a decoder stream: prefer the hvcC /// parameter-set arrays, then scan the item payload's length-prefixed NAL /// units in place. This visits NAL units in the same order as assembling the -/// stream and scanning it, but copies no payload bytes — the layout probe -/// runs this once per image-hook decode. +/// stream and scanning it, but copies no payload bytes. #[cfg(feature = "image-integration")] -fn decode_hevc_metadata_from_hvcc_or_payload( +fn decode_hevc_sps_from_hvcc_or_payload( hvcc: &isobmff::HevcDecoderConfigurationBox, payload: &[u8], -) -> Result { +) -> Result { let nal_length_size = hvcc.nal_length_size; if !(1..=4).contains(&nal_length_size) { return Err(DecodeHeicError::InvalidNalLengthSize { nal_length_size }); } - if let Some(metadata) = hevc_metadata_from_hvcc_nal_arrays(hvcc)? { - return Ok(metadata); + if let Some(sps) = hevc_sps_from_hvcc_nal_arrays(hvcc)? { + return Ok(sps); } - let mut metadata = None; + let mut sps = None; walk_length_prefixed_payload_nals(payload, usize::from(nal_length_size), |offset, nal| { let unit = LengthPrefixedHevcNalUnit { offset, bytes: nal }; if unit.nal_unit_type() != Some(NALUnitType::SpsNut) { return Ok(false); } - metadata = Some(hevc_metadata_from_sps_nal(nal, offset)?); + sps = Some(hevc_sps_from_nal(nal, offset)?); Ok(true) })?; - metadata.ok_or(DecodeHeicError::MissingSpsNalUnit) + sps.ok_or(DecodeHeicError::MissingSpsNalUnit) +} + +#[cfg(feature = "image-integration")] +fn decode_hevc_metadata_from_hvcc_or_payload( + hvcc: &isobmff::HevcDecoderConfigurationBox, + payload: &[u8], +) -> Result { + let sps = decode_hevc_sps_from_hvcc_or_payload(hvcc, payload)?; + hevc_metadata_from_sps(&sps) } fn heic_layout_from_sps_chroma_array_type( @@ -7515,6 +7958,73 @@ fn decode_heif_source_to_rgba( ) } +fn decode_heif_bytes_to_rgb8( + input: &[u8], + guardrails: DecodeGuardrails, +) -> Result { + let transforms = isobmff::parse_primary_item_transform_properties(input) + .map_err(DecodeHeicError::ParsePrimaryTransforms)? + .transforms; + let icc_profile = primary_icc_profile_from_heic(input); + let mut source: Option<&mut dyn RandomAccessSource> = None; + decode_primary_heic_to_rgb8_from_resolved_input( + input, + &mut source, + guardrails, + &transforms, + icc_profile, + ) +} + +fn decode_heif_source_to_rgb8( + source: &mut S, + input: &[u8], + guardrails: DecodeGuardrails, +) -> Result { + let transforms = isobmff::parse_primary_item_transform_properties(input) + .map_err(DecodeHeicError::ParsePrimaryTransforms)? + .transforms; + let icc_profile = primary_icc_profile_from_heic(input); + let mut source: Option<&mut dyn RandomAccessSource> = Some(source); + decode_primary_heic_to_rgb8_from_resolved_input( + input, + &mut source, + guardrails, + &transforms, + icc_profile, + ) +} + +fn decode_primary_heic_to_rgb8_from_resolved_input( + input: &[u8], + source: &mut Option<&mut dyn RandomAccessSource>, + guardrails: DecodeGuardrails, + transforms: &[isobmff::PrimaryItemTransformProperty], + icc_profile: Option>, +) -> Result { + let primary_with_grid = if let Some(source) = source.as_mut() { + isobmff::extract_primary_heic_item_data_with_grid_from_source(source, input) + .map_err(DecodeHeicError::from)? + } else { + isobmff::extract_primary_heic_item_data_with_grid(input).map_err(DecodeHeicError::from)? + }; + + match primary_with_grid { + isobmff::HeicPrimaryItemDataWithGrid::Grid(grid_data) => { + guardrails.enforce_pixel_count( + grid_data.descriptor.output_width, + grid_data.descriptor.output_height, + )?; + decode_primary_heic_grid_to_rgb8_image(&grid_data, transforms, icc_profile) + } + isobmff::HeicPrimaryItemDataWithGrid::Coded(item_data) => { + let decoded = decode_primary_heic_coded_item_to_image(input, &item_data)?; + guardrails.enforce_pixel_count(decoded.width, decoded.height)?; + decoded_heic_to_rgb8_image(decoded, transforms, icc_profile) + } + } +} + fn decode_primary_heic_to_rgba_from_resolved_input( input: &[u8], source: &mut Option<&mut dyn RandomAccessSource>, @@ -9093,6 +9603,19 @@ fn decode_bytes_to_rgba_with_hint_and_guardrails( } } +fn decode_bytes_to_rgb8_with_hint_and_guardrails( + input: &[u8], + hint: Option, + guardrails: DecodeGuardrails, +) -> Result { + match enforce_and_resolve_input_family(input, hint, &guardrails)? { + HeifInputFamily::Heif => decode_heif_bytes_to_rgb8(input, guardrails), + HeifInputFamily::Avif => Err(DecodeError::Unsupported( + "direct RGB8 output currently supports coded HEIC/HEIF inputs".to_string(), + )), + } +} + fn decode_bytes_to_png_with_hint_and_guardrails( input: &[u8], hint: Option, @@ -9155,6 +9678,27 @@ fn decode_source_to_rgba_with_hint_and_guardrails( } } +fn decode_source_to_rgb8_with_hint_and_guardrails( + source: &mut S, + hint: Option, + guardrails: DecodeGuardrails, +) -> Result { + guardrails.enforce_input_bytes(source.len())?; + let selected = + read_selected_top_level_boxes_from_source(source, &[FTYP_BOX_TYPE, META_BOX_TYPE])?; + let source_family_hint = detect_input_family_from_source_selected_boxes(&selected)?; + let input = encode_source_selected_top_level_boxes(&selected); + match source_family_hint + .or(hint) + .ok_or_else(unknown_input_family_error)? + { + HeifInputFamily::Heif => decode_heif_source_to_rgb8(source, &input, guardrails), + HeifInputFamily::Avif => Err(DecodeError::Unsupported( + "direct RGB8 output currently supports coded HEIC/HEIF inputs".to_string(), + )), + } +} + fn decode_source_to_png_with_hint_and_guardrails( source: &mut S, hint: Option, @@ -9205,6 +9749,20 @@ pub fn decode_bytes_to_rgba(input: &[u8]) -> Result Result { + decode_bytes_to_rgb8_with_hint_and_guardrails(input, Some(HeifInputFamily::Heif), guardrails) +} + +/// Decode coded HEIC/HEIF bytes directly into an owned RGB8 buffer. +pub fn decode_bytes_to_rgb8(input: &[u8]) -> Result { + decode_bytes_to_rgb8_with_guardrails(input, DecodeGuardrails::default()) +} + /// Decode a `Read` source with configurable guardrails into an owned RGBA buffer. pub fn decode_read_to_rgba_with_guardrails( input_reader: R, @@ -9257,6 +9815,31 @@ pub fn decode_path_to_rgba(input_path: &Path) -> Result Result { + if !input_path.exists() { + return Err(DecodeError::Unsupported(format!( + "Input file does not exist: {}", + input_path.display() + ))); + } + let mut source = FileSource::open(input_path).map_err(decode_error_from_source_read_error)?; + decode_source_to_rgb8_with_hint_and_guardrails( + &mut source, + extension_family_hint(input_path), + guardrails, + ) +} + +/// Decode a coded HEIC/HEIF path directly into an owned RGB8 buffer. +pub fn decode_path_to_rgb8(input_path: &Path) -> Result { + decode_path_to_rgb8_with_guardrails(input_path, DecodeGuardrails::default()) +} + /// Backward-compatible alias for [`decode_path_to_rgba`]. pub fn decode_file_to_rgba(input_path: &Path) -> Result { decode_path_to_rgba(input_path) @@ -9863,6 +10446,83 @@ fn decoded_heic_to_rgba_image( }) } +fn decoded_heic_to_rgb8_image( + mut decoded: DecodedHeicImage, + transforms: &[isobmff::PrimaryItemTransformProperty], + icc_profile: Option>, +) -> Result { + let (cropped, remaining_transforms) = + crop_heic_by_leading_chroma_aligned_clean_apertures(decoded, transforms)?; + decoded = cropped; + let source_bit_depth = heic_bit_depth_for_png_conversion(&decoded)?; + let pixels = convert_heic_to_rgb8(&decoded)?; + let (width, height, pixels) = apply_primary_item_transforms_rgb8( + decoded.width, + decoded.height, + pixels, + remaining_transforms, + )?; + Ok(DecodedRgbImage { + width, + height, + source_bit_depth, + pixels, + icc_profile, + }) +} + +fn apply_primary_item_transforms_rgb8( + width: u32, + height: u32, + pixels: Vec, + transforms: &[isobmff::PrimaryItemTransformProperty], +) -> Result<(u32, u32, Vec), DecodeError> { + let expected = checked_interleaved_sample_count(width, height, 3)?; + if pixels.len() != expected { + return Err(DecodeError::Unsupported(format!( + "RGB8 transform input has {} samples, expected {expected}", + pixels.len() + ))); + } + + let plan = RgbaTransformPlan::from_primary_transforms(width, height, transforms)?; + if plan.is_identity() { + return Ok((width, height, pixels)); + } + + let source_width = usize::try_from(width).map_err(|_| { + DecodeError::Unsupported(format!("RGB8 source width cannot be represented ({width})")) + })?; + let destination_width = usize::try_from(plan.destination_width).map_err(|_| { + DecodeError::Unsupported(format!( + "RGB8 destination width cannot be represented ({})", + plan.destination_width + )) + })?; + let destination_height = usize::try_from(plan.destination_height).map_err(|_| { + DecodeError::Unsupported(format!( + "RGB8 destination height cannot be represented ({})", + plan.destination_height + )) + })?; + let mut transformed = + vec![ + 0_u8; + checked_interleaved_sample_count(plan.destination_width, plan.destination_height, 3,)? + ]; + for destination_y in 0..destination_height { + for destination_x in 0..destination_width { + let (source_x, source_y) = plan.map_destination_pixel(destination_x, destination_y)?; + let source_sample = (source_y * source_width + source_x) * 3; + let destination_sample = (destination_y * destination_width + destination_x) * 3; + transformed[destination_sample..destination_sample + 3] + .copy_from_slice(&pixels[source_sample..source_sample + 3]); + } + } + + Ok((plan.destination_width, plan.destination_height, transformed)) +} + fn decoded_uncompressed_to_rgba_image( decoded: DecodedUncompressedImage, transforms: &[isobmff::PrimaryItemTransformProperty], @@ -10152,10 +10812,18 @@ fn apply_primary_item_transforms_rgba( } fn checked_rgba_sample_count(width: u32, height: u32) -> Result { + checked_interleaved_sample_count(width, height, 4) +} + +fn checked_interleaved_sample_count( + width: u32, + height: u32, + channels: u64, +) -> Result { let pixel_count = u64::from(width).checked_mul(u64::from(height)).ok_or({ DecodeError::TransformGuard(TransformGuardError::PixelCountOverflow { width, height }) })?; - let sample_count = pixel_count.checked_mul(4).ok_or({ + let sample_count = pixel_count.checked_mul(channels).ok_or({ DecodeError::TransformGuard(TransformGuardError::SampleCountOverflow { width, height }) })?; usize::try_from(sample_count).map_err(|_| { @@ -11018,25 +11686,56 @@ fn convert_heic_to_rgba8_into( ) -> Result<(), DecodeHeicError> { let output_len = checked_heic_rgba_output_len(decoded)?; out.resize(output_len, 0); - convert_heic_to_rgba8_slice(decoded, out) + convert_heic_to_interleaved_rgb8_slice::<4>(decoded, out, scale_sample_to_u8, "RGBA8") +} + +fn convert_heic_to_rgb8(decoded: &DecodedHeicImage) -> Result, DecodeHeicError> { + let mut out = Vec::new(); + convert_heic_to_rgb8_into(decoded, &mut out)?; + Ok(out) +} + +fn convert_heic_to_rgb8_into( + decoded: &DecodedHeicImage, + out: &mut Vec, +) -> Result<(), DecodeHeicError> { + let output_len = checked_heic_interleaved_output_len(decoded, 3, "RGB")?; + out.resize(output_len, 0); + convert_heic_to_interleaved_rgb8_slice::<3>( + decoded, + out, + scale_heic_sample_to_image_rgb8, + "RGB8", + ) } fn checked_heic_rgba_output_len(decoded: &DecodedHeicImage) -> Result { + checked_heic_interleaved_output_len(decoded, 4, "RGBA") +} + +fn checked_heic_interleaved_output_len( + decoded: &DecodedHeicImage, + channels: usize, + label: &str, +) -> Result { let expected_y_samples = heic_sample_count(decoded.width, decoded.height, "Y")?; expected_y_samples - .checked_mul(4) + .checked_mul(channels) .ok_or_else(|| DecodeHeicError::InvalidDecodedFrame { detail: format!( - "RGBA output sample count overflow for {}x{}", + "{label} output sample count overflow for {}x{}", decoded.width, decoded.height ), }) } -fn convert_heic_to_rgba8_slice( +fn convert_heic_to_interleaved_rgb8_slice( decoded: &DecodedHeicImage, out: &mut [u8], + scale_sample: fn(u16, u8) -> u8, + output_label: &str, ) -> Result<(), DecodeHeicError> { + debug_assert!(matches!(CHANNELS, 3 | 4)); let ycbcr_transform = ycbcr_transform_from_matrix(decoded.ycbcr_matrix).map_err(|matrix_coefficients| { DecodeHeicError::UnsupportedMatrixCoefficients { @@ -11065,19 +11764,11 @@ fn convert_heic_to_rgba8_slice( usize::try_from(decoded.height).map_err(|_| DecodeHeicError::InvalidDecodedFrame { detail: format!("HEIC height does not fit in usize ({})", decoded.height), })?; - let output_len = - expected_y_samples - .checked_mul(4) - .ok_or_else(|| DecodeHeicError::InvalidDecodedFrame { - detail: format!( - "RGBA output sample count overflow for {}x{}", - decoded.width, decoded.height - ), - })?; + let output_len = checked_heic_interleaved_output_len(decoded, CHANNELS, output_label)?; if out.len() != output_len { return Err(DecodeHeicError::InvalidDecodedFrame { detail: format!( - "RGBA8 output has {} samples, expected {output_len}", + "{output_label} output has {} samples, expected {output_len}", out.len() ), }); @@ -11098,7 +11789,7 @@ fn convert_heic_to_rgba8_slice( for y in 0..height { let row_start = y * width; - let out_row_start = row_start * 4; + let out_row_start = row_start * CHANNELS; for x in 0..width { let y_index = row_start + x; @@ -11128,11 +11819,13 @@ fn convert_heic_to_rgba8_slice( } else { converter.convert(y_sample, cb_sample, cr_sample) }; - let out_index = out_row_start + (x * 4); - out[out_index] = scale_sample_to_u8(r, bit_depth); - out[out_index + 1] = scale_sample_to_u8(g, bit_depth); - out[out_index + 2] = scale_sample_to_u8(b, bit_depth); - out[out_index + 3] = u8::MAX; + let out_index = out_row_start + (x * CHANNELS); + out[out_index] = scale_sample(r, bit_depth); + out[out_index + 1] = scale_sample(g, bit_depth); + out[out_index + 2] = scale_sample(b, bit_depth); + if CHANNELS == 4 { + out[out_index + 3] = u8::MAX; + } } } @@ -11903,6 +12596,17 @@ fn scale_sample_to_u16(sample: u16, bit_depth: u8) -> u16 { as u16 } +/// Match `image`'s RGBA16-to-RGB8 conversion for high-bit-depth HEIC while +/// avoiding the intermediate RGBA16 allocation. Eight-bit inputs retain their +/// existing direct scaling path. +fn scale_heic_sample_to_image_rgb8(sample: u16, bit_depth: u8) -> u8 { + if bit_depth <= 8 { + return scale_sample_to_u8(sample, bit_depth); + } + + ((u32::from(scale_sample_to_u16(sample, bit_depth)) + 128) / 257) as u8 +} + #[derive(Default)] struct DecoderContextGuard(Option); @@ -12461,7 +13165,169 @@ mod tests { assert_eq!(direct, transformed); } - #[cfg(feature = "image-integration")] + #[cfg(feature = "parallel-grid")] + #[test] + fn grid_parallel_window_uses_each_candidate_tile_estimate() { + const MIB: u64 = 1024 * 1024; + + assert_eq!( + super::heic_grid_tile_decode_window_from_estimates( + [Some(4 * MIB), Some(80 * MIB), Some(4 * MIB)], + 8, + ), + 1, + "a later oversized tile must not inherit the first candidate's small estimate" + ); + assert_eq!( + super::heic_grid_tile_decode_window_from_estimates( + [Some(32 * MIB), Some(32 * MIB), Some(MIB)], + 8, + ), + 2, + "a window may fill the memory budget exactly" + ); + assert_eq!( + super::heic_grid_tile_decode_window_from_estimates( + [Some(4 * MIB), None, Some(4 * MIB)], + 8, + ), + 1, + "a tile whose metadata cannot be preflighted must stay outside the parallel window" + ); + assert_eq!( + super::heic_grid_tile_decode_window_from_estimates( + core::iter::repeat_n(Some(MIB), 16), + 8, + ), + 8, + "the thread cap must remain effective for small tiles" + ); + } + + #[cfg(feature = "parallel-grid")] + #[test] + fn grid_parallel_estimate_uses_coded_sps_geometry_and_both_stream_copies() { + use super::isobmff::{HeicGridTileItemData, HevcNalArray}; + + // x265 SPS for a 126x126 visible 4:2:0 image whose coded frame is + // 128x128. The decoder allocates the coded geometry before applying + // the two-pixel SPS conformance window. + let sps_nal = vec![ + 0x42, 0x01, 0x01, 0x03, 0x70, 0x00, 0x00, 0x03, 0x00, 0x90, 0x00, 0x00, 0x03, 0x00, + 0x00, 0x03, 0x00, 0x1e, 0xa0, 0x10, 0x20, 0x20, 0x75, 0x59, 0x65, 0x66, 0x92, 0x4c, + 0xae, 0x01, 0x00, 0x00, 0x03, 0x03, 0xe8, 0x00, 0x00, 0x03, 0x03, 0xe8, 0x08, + ]; + let tile = HeicGridTileItemData { + item_id: 1, + construction_method: 0, + hvcc: test_hvcc( + vec![HevcNalArray { + array_completeness: true, + nal_unit_type: 33, + nal_units: vec![sps_nal.clone()], + }], + 4, + ), + colr: Default::default(), + transforms: Vec::new(), + payload: Vec::new(), + }; + + let sps = super::hevc_sps_from_nal(&sps_nal, 0).expect("SPS should parse"); + let metadata = super::hevc_metadata_from_sps(&sps).expect("SPS metadata should parse"); + assert_eq!((metadata.width, metadata.height), (126, 126)); + + let decoded_bytes = super::estimate_heic_grid_sps_decode_bytes(&sps) + .expect("SPS allocation estimate should parse"); + let mut stream_memory = super::HeicGridTileStreamMemoryEstimate::default(); + stream_memory + .add_nal(&sps_nal) + .expect("SPS stream estimate should fit"); + assert_eq!( + super::estimate_heic_grid_tile_decode_bytes(&tile).expect("tile estimate should parse"), + decoded_bytes + stream_memory.estimated_decoder_bytes(), + ); + assert_eq!(decoded_bytes, 128 * 128 * 6); + } + + #[cfg(feature = "parallel-grid")] + #[test] + fn grid_parallel_estimate_uses_largest_sps_across_hvcc_and_payload() { + use super::isobmff::{HeicGridTileItemData, HevcNalArray}; + + let small_sps = vec![ + 0x42, 0x01, 0x01, 0x03, 0x70, 0x00, 0x00, 0x03, 0x00, 0x90, 0x00, 0x00, 0x03, 0x00, + 0x00, 0x03, 0x00, 0x1e, 0xa0, 0x20, 0x81, 0x05, 0x96, 0xea, 0xae, 0x9a, 0xe6, 0xe0, + 0x21, 0xa0, 0xc0, 0x80, 0x00, 0x00, 0x0c, 0x80, 0x00, 0x00, 0x03, 0x00, 0x84, + ]; + let large_sps = vec![ + 0x42, 0x01, 0x01, 0x03, 0x70, 0x00, 0x00, 0x03, 0x00, 0x90, 0x00, 0x00, 0x03, 0x00, + 0x00, 0x03, 0x00, 0x78, 0xa0, 0x07, 0xd2, 0x00, 0x53, 0x9f, 0x59, 0x6e, 0xa4, 0x92, + 0x9a, 0xe6, 0xe0, 0x21, 0xa0, 0xc0, 0x80, 0x00, 0x00, 0x0c, 0x80, 0x00, 0x00, 0x03, + 0x00, 0x84, + ]; + let mut payload = Vec::new(); + payload.extend_from_slice(&(large_sps.len() as u32).to_be_bytes()); + payload.extend_from_slice(&large_sps); + let tile = HeicGridTileItemData { + item_id: 1, + construction_method: 0, + hvcc: test_hvcc( + vec![HevcNalArray { + array_completeness: true, + nal_unit_type: 33, + nal_units: vec![small_sps.clone()], + }], + 4, + ), + colr: Default::default(), + transforms: Vec::new(), + payload, + }; + + let small = super::hevc_sps_from_nal(&small_sps, 0).expect("small SPS should parse"); + let large = super::hevc_sps_from_nal(&large_sps, 0).expect("large SPS should parse"); + let small_bytes = super::estimate_heic_grid_sps_decode_bytes(&small) + .expect("small SPS estimate should parse"); + let large_bytes = super::estimate_heic_grid_sps_decode_bytes(&large) + .expect("large SPS estimate should parse"); + assert!(large_bytes > small_bytes); + + let mut stream_memory = super::HeicGridTileStreamMemoryEstimate::default(); + stream_memory + .add_nal(&small_sps) + .expect("small SPS stream estimate should fit"); + stream_memory + .add_nal(&large_sps) + .expect("large SPS stream estimate should fit"); + assert_eq!( + super::estimate_heic_grid_tile_decode_bytes(&tile).expect("tile estimate should parse"), + large_bytes + stream_memory.estimated_decoder_bytes(), + ); + } + + #[cfg(feature = "parallel-grid")] + #[test] + fn grid_parallel_stream_estimate_counts_epb_positions_and_nal_metadata() { + let nal = [ + 0x4e, 0x01, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x01, 0x00, 0x00, 0x03, 0x02, + ]; + let mut stream_memory = super::HeicGridTileStreamMemoryEstimate::default(); + stream_memory + .add_nal(&nal) + .expect("NAL stream estimate should fit"); + + assert_eq!(stream_memory.normalized_stream_bytes, nal.len() as u64 + 4); + assert_eq!(stream_memory.rbsp_bytes, nal.len() as u64 - 2); + assert_eq!(stream_memory.emulation_prevention_position_bytes, 40); + assert!( + stream_memory.estimated_decoder_bytes() + > stream_memory.normalized_stream_bytes * 2 + stream_memory.rbsp_bytes, + "parsed-NAL metadata and EPB positions must add memory beyond the stream and RBSP copies" + ); + } + + #[cfg(any(feature = "image-integration", feature = "parallel-grid"))] fn test_hvcc( nal_arrays: Vec, nal_length_size: u8, @@ -12771,6 +13637,76 @@ mod tests { .expect("reference RGBA transform should succeed") } + fn rgba8_to_rgb8(pixels: &[u8]) -> Vec { + pixels + .chunks_exact(4) + .flat_map(|pixel| pixel[..3].iter().copied()) + .collect() + } + + #[test] + fn direct_rgb8_matches_final_rgba8_output_with_transforms() { + let image = synthetic_yuv_image(super::HeicPixelLayout::Yuv420); + let transforms = [ + isobmff::PrimaryItemTransformProperty::CleanAperture(synthetic_clean_aperture(-1, -1)), + isobmff::PrimaryItemTransformProperty::Mirror(isobmff::ImageMirrorProperty { + direction: isobmff::ImageMirrorDirection::Horizontal, + }), + isobmff::PrimaryItemTransformProperty::Rotation(isobmff::ImageRotationProperty { + rotation_ccw_degrees: 90, + }), + ]; + let rgba = super::decoded_heic_to_rgba_image(image.clone(), &transforms, None, None) + .expect("RGBA decode should succeed"); + let direct = super::decoded_heic_to_rgb8_image(image, &transforms, None) + .expect("direct RGB8 decode should succeed"); + let super::DecodedRgbaPixels::U8(rgba_pixels) = rgba.pixels else { + panic!("eight-bit input should produce RGBA8"); + }; + + assert_eq!((direct.width, direct.height), (rgba.width, rgba.height)); + assert_eq!(direct.pixels, rgba8_to_rgb8(&rgba_pixels)); + } + + #[test] + fn direct_rgb8_matches_rgba16_to_rgb8_rounding() { + let mut image = synthetic_yuv_image(super::HeicPixelLayout::Yuv444); + image.bit_depth_luma = 10; + image.bit_depth_chroma = 10; + for sample in &mut image.y_plane.samples { + *sample *= 4; + } + for sample in image.u_plane.as_mut().expect("U plane").samples.iter_mut() { + *sample *= 4; + } + for sample in image.v_plane.as_mut().expect("V plane").samples.iter_mut() { + *sample *= 4; + } + let transforms = [isobmff::PrimaryItemTransformProperty::Rotation( + isobmff::ImageRotationProperty { + rotation_ccw_degrees: 270, + }, + )]; + let rgba = super::decoded_heic_to_rgba_image(image.clone(), &transforms, None, None) + .expect("RGBA16 decode should succeed"); + let direct = super::decoded_heic_to_rgb8_image(image, &transforms, None) + .expect("direct RGB8 decode should succeed"); + let super::DecodedRgbaPixels::U16(rgba_pixels) = rgba.pixels else { + panic!("ten-bit input should produce RGBA16"); + }; + let expected: Vec = rgba_pixels + .chunks_exact(4) + .flat_map(|pixel| { + pixel[..3] + .iter() + .map(|sample| ((u32::from(*sample) + 128) / 257) as u8) + }) + .collect(); + + assert_eq!((direct.width, direct.height), (rgba.width, rgba.height)); + assert_eq!(direct.pixels, expected); + } + #[test] fn odd_clean_aperture_crop_falls_back_to_rgba_transform_path() { let image = synthetic_yuv_image(super::HeicPixelLayout::Yuv420);