Carry JPEG tokens and WebP bits at their natural width - #320
Conversation
Second-pass memory audit of the encode paths, measured with per-workload peak RSS and tracemalloc probes in fresh subprocesses. Both changes are byte-identical by construction and validated as such. JPEG. `_component_tokens` was the single largest allocation in the engine: a 3200x2400 photographic encode peaked at 731 MB, of which the token stage alone was 742 MB of the traced high-water. It held roughly twenty parallel int64 arrays over the nonzero coefficients — but AC positions top out at 62, run lengths at 62, magnitude categories at 12, and symbols are a byte by definition. The coefficients themselves fit int16: `_DCT` is orthonormal, so a level-shifted 8-bit block transforms to at most sqrt(64) * 128 = 1024 in magnitude, and quantizers are >= 1. Carrying block indices as int32, amplitudes as int16 and everything else as uint8 — plus int32/uint8 Huffman tables, since both are gathered once per token — cuts that encode to 200 MB and makes it ~3% faster. The MCU sort key is explicitly widened to int64 on the way in: a maximal 65535x65535 frame has 67M blocks, and `block * 3 << 8` leaves int32 well before that. The RGB->YCbCr transform also promoted the whole frame to interleaved float before splitting it into planes — three floats per pixel on top of the three planes, 92 MB of pure transient at that size. It now runs in row bands, which is elementwise-identical. WebP. `_pack_lsb` was the entire peak of a WebP export (216.2 MB of 216.2 MB). It ran one masked pass per bit position over the whole token stream, and each pass built a mask plus three compacted gathers — about five machine words per entry, forty times over. The passes now run over bounded entry blocks, each iterating only as far as its own widest entry; every entry still scatters to the same absolute bit position, so the output cannot move. The header also rides in the token buffer instead of being concatenated onto the front of it, which was a second full copy of the stream. 216 MB -> 135 MB at unchanged speed. Measured on the 20-workload memory suite, warm __pycache__, native core held constant: jpeg_export_photo 417.7 -> 205.1 MB traced peak (845.6 -> 431.4 MB RSS), webp_export 245.5 -> 164.1 MB (397.8 -> 289.8 MB RSS). Every other workload is identical to the hundredth of a megabyte. Equivalence: 117 JPEG encodes (13 images x 9 qualities, including flat, 1x1, single row/column, RGB-vs-RGBA, pixel checkerboards, ZRL-heavy sparse runs and an unaligned photographic case) and 10 WebP encodes are byte-identical to the pristine encoders, as are all eight JPEG error paths. New tests pin the coefficient bound, the token field widths, and the transparency of both new chunk sizes against one-shot runs.
`_flatten_alpha` is the JPEG determinism backstop: it composites leftover alpha over white before encoding. It did so for the whole frame at once, in uint16 — six full-frame temporaries — and returned a four-channel image whose alpha plane the encoder immediately discards (`_jpeg.encode` reads `rgba[..., :3]`, and both call sites feed it directly). Returning (h, w, 3) and compositing in row bands takes a 3200x2400 export from 117.2 MB to 27.4 MB of traced peak and makes it 30.8% faster. The arithmetic is unchanged and elementwise: `rgb * a + 255 * (255 - a) + 127` peaks at 65152 for any 8-bit input, comfortably inside uint16, so the band size cannot move a byte. Verified against the previous implementation over all 65,536 (value, alpha) pairs plus opaque, fully transparent, graded-alpha, 1x1 and single-row images, comparing both the composited RGB and the encoded JPEG.
|
Warning Review limit reached
Next review available in: 31 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (6)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Merging this PR will not alter performance
Comparing 
Footnotes
|
1 participant
Second pass over engine memory, following #309. Same method: peak RSS and
tracemalloc.reset_peak()probes per workload, each in a fresh subprocess, withbyte-identical output as the gate. This PR is the encode paths; four sibling PRs
carry the rest.
Stage attribution found the whole cost in two functions rather than spread
around, so the fixes are narrow.
JPEG: the token pipeline was int64 everywhere
_component_tokenswas the single largest allocation in the engine. Replayingencodewithreset_peak()between stages, a 3200x2400 photographic image:The token stage held ~20 parallel int64 arrays over the nonzero coefficients.
None of those values need eight bytes: AC positions top out at 62, run lengths
at 62, magnitude categories at 12, and symbols are a byte by definition. The
coefficients themselves fit int16 —
_DCTis orthonormal, so a level-shifted8-bit block transforms to at most
sqrt(64) * 128 = 1024in magnitude, andquantizers are always >= 1.
Carrying block indices as int32, amplitudes as int16 and everything else as
uint8 — plus int32/uint8 Huffman tables, since both are gathered once per
token — takes that encode from 730.7 MB to 199.7 MB and makes it 2.6%
faster. The MCU sort key is explicitly widened to int64 on the way in: a
maximal 65535x65535 frame has 67M blocks and
block * 3 << 8leaves int32 wellbefore that.
The RGB→YCbCr transform also promoted the whole frame to interleaved float
before splitting it into planes — three floats per pixel on top of the three
planes, 92 MB of pure transient at that size. It now runs in row bands, which is
elementwise-identical.
WebP:
_pack_lsbwas the entire peak216.2 MB of a 216.2 MB export. It ran one masked pass per bit position over the
whole token stream, and each pass built a boolean mask plus three compacted
gathers — about five machine words per entry, forty times over. The passes now
run over bounded entry blocks, each iterating only as far as its own widest
entry. Every entry still scatters to the same absolute bit position, so the
output cannot move. The header also rides in the token buffer instead of being
concatenated onto the front of it, which was a second full copy of the stream.
216 MB -> 135 MB at unchanged speed.
_flatten_alphareturned a plane JPEG never readsThe JPEG determinism backstop composited leftover alpha over white for the whole
frame at once, in uint16 — six full-frame temporaries — and returned a
four-channel image whose alpha the encoder immediately discards (
_jpeg.encodereads
rgba[..., :3], and both call sites feed it directly). Returning(h, w, 3)and compositing in row bands: 117.2 MB -> 27.4 MB, 30.8% faster.rgb * a + 255 * (255 - a) + 127peaks at 65152 for any 8-bit input, so theuint16 intermediates stay exact and the band size cannot move a byte.
Measured
20-workload memory suite, warm
__pycache__, native core held constant:jpeg_export_phototraced peakjpeg_export_photopeak RSSwebp_exporttraced peakwebp_exportpeak RSSEvery other workload is identical to the hundredth of a megabyte.
Equivalence
RGB-vs-RGBA, pixel and 8x8 checkerboards, ZRL-heavy sparse runs, smooth
gradients, an unaligned photographic case) byte-identical to the pristine
encoder, plus all 8 error paths.
and one crossing the new entry-block boundary.
selections through both zone-pruning branches, density views, append,
transition keys, pyplot) identical to main.
transparency of both new chunk sizes against one-shot runs.
Full suite: 2428 passed, 4 skipped. Ruff clean.
Perf
Full 103-region CodSpeed suite in walltime mode, interleaved A/B against main
with the native core held constant: median -0.96%, mean -0.91%, 5 regions more
than 5% better. One region read +20% (
factorize_fixed_categorical) — it has52% intra-arm spread in the reference arm (0.90–1.37 ms) and shares the same
dylib in both arms, so it cannot be caused by this change.