From 4e7729773e3af72217458ea65d67f4590fc10960 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 01:03:27 -0700 Subject: [PATCH 1/2] Carry JPEG tokens and WebP bits at their natural width MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- CHANGELOG.md | 14 ++++ python/xy/_jpeg.py | 156 ++++++++++++++++++++++++++++++++------------- python/xy/_webp.py | 45 +++++++++---- tests/test_jpeg.py | 78 +++++++++++++++++++++++ tests/test_webp.py | 44 +++++++++++++ 5 files changed, 279 insertions(+), 58 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 82f041fd..310560fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -73,6 +73,20 @@ in the README). - The native rasterizer's serial point and segment passes no longer materialize an `0..n` index vector to read back in order (4 bytes per mark), and the per-point density quantizer keeps one temporary instead of three. + - The JPEG token pipeline carries every field at its natural width. It held + ~20 parallel `int64` arrays over the nonzero coefficients — positions that + top out at 62, categories at 12, symbols that are a byte by definition — + while the coefficients themselves fit `int16`, because an orthonormal 8×8 + DCT of level-shifted 8-bit samples is bounded by 1024 and quantizers are + ≥ 1. The RGB→YCbCr transform also promoted the whole frame to interleaved + float before splitting it into planes. A 3200×2400 photographic encode now + peaks at **200 MB instead of 731 MB** and is ~3% faster; the chart-shaped + export in the memory suite drops 418 MB → 205 MB. + - The lossless WebP packer scatters bits in bounded entry blocks instead of + one masked pass per bit position over the whole token stream (five machine + words per entry, per position), and the header now rides in the token + buffer rather than being concatenated onto the front of it. A 3200×2400 + export peaks at 135 MB instead of 216 MB at unchanged speed. - `memory_report()` reports `capacity_bytes` per column and `canonical_capacity_bytes` per store, and builds `resident_array_bytes` from the capacity total. A streamed column's `values` is a prefix view of its diff --git a/python/xy/_jpeg.py b/python/xy/_jpeg.py index 222eaf0c..80dbbffd 100644 --- a/python/xy/_jpeg.py +++ b/python/xy/_jpeg.py @@ -128,11 +128,16 @@ def _zigzag_order() -> np.ndarray: def _huff_lookup(bits: bytes, values: bytes) -> tuple[np.ndarray, np.ndarray]: - """Canonical symbol→(code, length) arrays from a BITS/HUFFVAL spec.""" + """Canonical symbol→(code, length) arrays from a BITS/HUFFVAL spec. + + Codes are at most 16 bits and lengths at most 16, so the lookup tables are + int32/uint8 rather than int64: every token gathers one entry from each, and + the gathered arrays are the widest thing the entropy stage carries. + """ if len(values) != sum(bits): raise AssertionError("Huffman spec mismatch: HUFFVAL count != sum(BITS)") - codes = np.zeros(256, dtype=np.int64) - lens = np.zeros(256, dtype=np.int64) + codes = np.zeros(256, dtype=np.int32) + lens = np.zeros(256, dtype=np.uint8) code = 0 k = 0 for length in range(1, 17): @@ -175,7 +180,8 @@ def _dct_matrix() -> np.ndarray: def _bit_size(magnitude: np.ndarray) -> np.ndarray: - return np.searchsorted(_POW2, magnitude, side="right").astype(np.int64) + # Categories top out at 12, so the result is a byte per token, not 8. + return np.searchsorted(_POW2, magnitude, side="right").astype(np.uint8) def _scaled_quant(base: np.ndarray, quality: int) -> np.ndarray: @@ -192,18 +198,27 @@ def _component_tokens( Returns parallel arrays (block, seq, table, symbol, amplitude, amp_bits); `seq` orders tokens within a block (0 = DC, 255 = EOB) so a single sort on (block, component, seq) later interleaves the MCU stream. + + Every field is carried at its natural width rather than as an int64. A + photographic image mints millions of tokens and this pipeline holds ~20 + parallel arrays at once, so width is the whole cost: positions and runs are + at most 62, categories at most 12, symbols are a byte by definition, and + amplitudes fit int16 because an orthonormal 8x8 DCT of level-shifted 8-bit + samples is bounded by sqrt(64) * 128 = 1024 before quantization (and + quantizers are >= 1). Only the block index needs more than two bytes. """ n = coef.shape[0] - # DC is coded differentially along the component's block sequence. - diff = np.diff(coef[:, 0], prepend=np.int64(0)) + # DC is coded differentially along the component's block sequence. These are + # per-block, not per-token, so they stay small regardless of width. + diff = np.diff(coef[:, 0], prepend=np.int16(0)) dsize = _bit_size(np.abs(diff)) # T.81 amplitude coding: negatives are sent as v + 2**size - 1. - dampl = np.where(diff < 0, diff + (np.int64(1) << dsize) - 1, diff) + dampl = np.where(diff < 0, diff + (np.int16(1) << dsize) - 1, diff).astype(np.int16) dc_tok = ( - np.arange(n, dtype=np.int64), - np.zeros(n, dtype=np.int64), - np.full(n, dc_tbl, dtype=np.int64), + np.arange(n, dtype=np.int32), + np.zeros(n, dtype=np.uint8), + np.full(n, dc_tbl, dtype=np.uint8), dsize, dampl, dsize, @@ -211,50 +226,80 @@ def _component_tokens( ac = coef[:, 1:] blk, pos = np.nonzero(ac) # row-major: block-ascending, position-ascending - last = np.full(n, -1, dtype=np.int64) + last = np.full(n, -1, dtype=np.int16) if blk.size: last[blk] = pos # row-major order → last write per block wins val = ac[blk, pos] + # np.nonzero hands back two intp vectors; narrow both before anything + # else is derived from them, so the int64 pair is transient rather than + # the base width of the whole pipeline. + blk = blk.astype(np.int32) + pos = pos.astype(np.uint8) first = np.empty(blk.size, dtype=bool) first[0] = True np.not_equal(blk[1:], blk[:-1], out=first[1:]) - prev = np.where(first, np.int64(-1), np.concatenate((pos[:1] * 0 - 1, pos[:-1]))) - run = pos - prev - 1 + # Previous nonzero position within the block, -1 at a block's first. + # Built in place: the shift-and-mask form allocated a concatenate and a + # where, and an unsigned `pos[:1] * 0 - 1` would wrap to 255. + prev = np.empty(pos.size, dtype=np.int8) + prev[0] = -1 + prev[1:] = pos[:-1] + np.copyto(prev, np.int8(-1), where=first) + run = (pos - prev - np.int16(1)).astype(np.uint8) + del prev zrl = run >> 4 # each 16 zeros of run becomes a ZRL (0xF0) token asize = _bit_size(np.abs(val)) sym = ((run & 15) << 4) | asize - aampl = np.where(val < 0, val + (np.int64(1) << asize) - 1, val) + del run + aampl = np.where(val < 0, val + (np.int16(1) << asize) - 1, val).astype(np.int16) + del val # Expand each nonzero into its ZRL prefix + the coefficient token, # numbering tokens within their block via a segmented cumsum: `start` # is the global exclusive cumsum of token counts and `base` forward- # fills each block's opening value, so `start - base` restarts at 0. - tot = zrl + 1 - cum = np.cumsum(tot) + tot = zrl + np.uint8(1) + cum = np.cumsum(tot, dtype=np.int32) start = cum - tot - base = np.maximum.accumulate(np.where(first, start, 0)) - rep = np.repeat(np.arange(blk.size, dtype=np.int64), tot) - j = np.arange(cum[-1], dtype=np.int64) - np.repeat(start, tot) + ntok_ac = int(cum[-1]) + del cum + base = np.maximum.accumulate(np.where(first, start, np.int32(0))) + del first + rep = np.repeat(np.arange(blk.size, dtype=np.int32), tot) + # `j` only ever counts within one nonzero's ZRL run, so it is <= 3. + j = (np.arange(ntok_ac, dtype=np.int32) - np.repeat(start, tot)).astype(np.uint8) + del tot is_zrl = j < zrl[rep] + seq = (start - base)[rep] + del start, base + seq += 1 + seq += j ac_tok = ( - blk[rep].astype(np.int64), - 1 + (start - base)[rep] + j, - np.full(rep.size, ac_tbl, dtype=np.int64), - np.where(is_zrl, np.int64(0xF0), sym[rep]), - np.where(is_zrl, np.int64(0), aampl[rep]), - np.where(is_zrl, np.int64(0), asize[rep]), + blk[rep], + seq.astype(np.uint8), + np.full(rep.size, ac_tbl, dtype=np.uint8), + np.where(is_zrl, np.uint8(0xF0), sym[rep]), + np.where(is_zrl, np.int16(0), aampl[rep]), + np.where(is_zrl, np.uint8(0), asize[rep]), ) else: - ac_tok = tuple(np.empty(0, dtype=np.int64) for _ in range(6)) + ac_tok = ( + np.empty(0, dtype=np.int32), + np.empty(0, dtype=np.uint8), + np.empty(0, dtype=np.uint8), + np.empty(0, dtype=np.uint8), + np.empty(0, dtype=np.int16), + np.empty(0, dtype=np.uint8), + ) # EOB unless the block's final zigzag coefficient (AC position 62) is set. - eob = np.flatnonzero(last != 62).astype(np.int64) + eob = np.flatnonzero(last != 62).astype(np.int32) eob_tok = ( eob, - np.full(eob.size, 255, dtype=np.int64), - np.full(eob.size, ac_tbl, dtype=np.int64), - np.zeros(eob.size, dtype=np.int64), - np.zeros(eob.size, dtype=np.int64), - np.zeros(eob.size, dtype=np.int64), + np.full(eob.size, 255, dtype=np.uint8), + np.full(eob.size, ac_tbl, dtype=np.uint8), + np.zeros(eob.size, dtype=np.uint8), + np.zeros(eob.size, dtype=np.int16), + np.zeros(eob.size, dtype=np.uint8), ) merged = [np.concatenate(parts) for parts in zip(dc_tok, ac_tok, eob_tok, strict=True)] return merged[0], merged[1], merged[2], merged[3], merged[4], merged[5] @@ -267,6 +312,11 @@ def _component_tokens( #: leftover bits into the next one, which keeps the byte stream identical. _ENTROPY_BIT_CHUNK = 1 << 18 +#: Pixels per band of the RGB→YCbCr transform. The transform is elementwise, so +#: banding it is bit-identical; this bounds the interleaved f32 promotion that +#: would otherwise be three floats per pixel of the whole frame. +_YCBCR_PIXEL_CHUNK = 1 << 18 + def _pack_entropy(chunk: np.ndarray, nbits: np.ndarray) -> bytes: """MSB-first bit packing of (value, bit-count) chunks, with 1-padding to a @@ -278,7 +328,7 @@ def _pack_entropy(chunk: np.ndarray, nbits: np.ndarray) -> bytes: local, so it applies per group. """ total = int(nbits.sum()) - cum = np.cumsum(nbits) + cum = np.cumsum(nbits, dtype=np.int64) # Token index at each bit-budget boundary: a group is [t0, t1) tokens, and a # single token never spans groups (its bits stay contiguous, as before). edges = np.searchsorted( @@ -293,7 +343,9 @@ def _pack_entropy(chunk: np.ndarray, nbits: np.ndarray) -> bytes: values = chunk[t0:t1] widths = nbits[t0:t1] group_total = int(widths.sum()) - start = np.cumsum(widths) - widths + # Pinned to int64: bit widths are uint8, and an unsigned accumulator + # would make `arange(int64) - start` promote to float64 under NEP 50. + start = np.cumsum(widths, dtype=np.int64) - widths idx = np.repeat(np.arange(values.size, dtype=np.int64), widths) offset = np.arange(group_total, dtype=np.int64) - np.repeat(start, widths) bits = ((values[idx] >> (widths[idx] - 1 - offset)) & 1).astype(np.uint8) @@ -395,17 +447,26 @@ def encode(rgba: np.ndarray, *, quality: int = 90) -> bytes: # Alpha is ignored: the caller has already composited onto an opaque # background, so only the RGB planes carry information. - rgb = rgba[..., :3].astype(np.float32) - r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2] + # # JFIF BT.601 full-range transform. The −128 DCT level shift cancels the # +128 chroma offset, so Y is shifted here and Cb/Cr are left centered. - y = 0.299 * r + 0.587 * g + 0.114 * b - 128.0 - cb = -0.168736 * r - 0.331264 * g + 0.5 * b - cr = 0.5 * r - 0.418688 * g - 0.081312 * b - # The interleaved f32 source is dead once the planes exist, and it is three - # floats per pixel — hand it back before the per-component pipeline runs. - # (`r`/`g`/`b` are views into it, so all four references have to go.) - del rgb, r, g, b + # + # Done in row bands: the transform is elementwise, so a band produces + # bit-identical f32 output, but promoting the whole frame to interleaved + # float first costs three floats per pixel on top of the three planes — + # 92 MB of transient on a 3200x2400 export, more than the planes themselves. + y = np.empty((h, w), dtype=np.float32) + cb = np.empty((h, w), dtype=np.float32) + cr = np.empty((h, w), dtype=np.float32) + band = max(1, _YCBCR_PIXEL_CHUNK // w) + for r0 in range(0, h, band): + r1 = min(r0 + band, h) + rgb = rgba[r0:r1, :, :3].astype(np.float32) + r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2] + y[r0:r1] = 0.299 * r + 0.587 * g + 0.114 * b - 128.0 + cb[r0:r1] = -0.168736 * r - 0.331264 * g + 0.5 * b + cr[r0:r1] = 0.5 * r - 0.418688 * g - 0.081312 * b + del rgb, r, g, b qy = _scaled_quant(_QUANT_LUMA, quality) qc = _scaled_quant(_QUANT_CHROMA, quality) @@ -450,7 +511,10 @@ def encode(rgba: np.ndarray, *, quality: int = 90) -> bytes: scaled += bias del bias np.trunc(scaled, out=scaled) - quant = scaled.astype(np.int64) + # int16 is exact here and halves every array the token stage derives: + # an orthonormal DCT of level-shifted 8-bit samples is bounded by 1024 + # in magnitude and quantizers are >= 1, so |coefficient| <= 1024. + quant = scaled.astype(np.int16) del scaled # Exact-math AC magnitudes cap at 1020 (category 10); clamp is one-LSB # insurance against float rounding ever minting category 11, which the @@ -461,7 +525,9 @@ def encode(rgba: np.ndarray, *, quality: int = 90) -> bytes: del quant # 4:4:4 → one block per component per MCU, so the MCU-interleaved # order Y, Cb, Cr is a stable sort on (block, component, in-block seq). - keys.append(((toks[0] * 3 + comp) << 8) | toks[1]) + # The key is widened deliberately: the maximum 65535x65535 frame has 67M + # blocks, and `block * 3 << 8` leaves int32 well before that. + keys.append(((toks[0].astype(np.int64) * 3 + comp) << 8) | toks[1]) # Only the four emitted fields outlive the key: the block and in-block # sequence numbers are folded into it, and on a photographic image each # dropped field is an int64 per token. diff --git a/python/xy/_webp.py b/python/xy/_webp.py index 40d68cb6..54ab0082 100644 --- a/python/xy/_webp.py +++ b/python/xy/_webp.py @@ -208,18 +208,32 @@ def _write_prefix_code(put, hist: np.ndarray) -> tuple[np.ndarray, np.ndarray]: return emit_len, emit_code +#: Entries per bit-scatter pass. Each pass builds a boolean mask and three +#: compacted gathers over the entries it covers, so an unbounded pass costs +#: ~5 machine words per entry *per bit position* -- 216 MB of transient on a +#: 3200x2400 export. Bounding it changes nothing about which bits are written +#: or where, because every entry scatters to absolute positions in `bits`. +_PACK_ENTRY_CHUNK = 1 << 17 + + def _pack_lsb(values: np.ndarray, nbits: np.ndarray) -> bytes: """Pack (value, nbits) pairs into an LSB-first byte stream. Expands to one uint8 per bit and lets packbits fold them: one vectorized - pass per bit *position* (<= ~40) instead of a Python loop per symbol. + pass per bit *position* (<= ~40) instead of a Python loop per symbol. The + passes run over bounded entry blocks, and each block only iterates as far + as its own widest entry. """ - ends = np.cumsum(nbits, dtype=np.int64) - starts = ends - nbits - bits = np.zeros(int(ends[-1]), np.uint8) - for k in range(int(nbits.max())): - m = nbits > k - bits[starts[m] + k] = (values[m] >> np.uint64(k)) & np.uint64(1) + starts = np.cumsum(nbits, dtype=np.int64) + total = int(starts[-1]) if starts.size else 0 + starts -= nbits # exclusive scan, in place: the ends array is dead here + bits = np.zeros(total, np.uint8) + for i0 in range(0, values.size, _PACK_ENTRY_CHUNK): + i1 = min(i0 + _PACK_ENTRY_CHUNK, values.size) + vals, widths, offs = values[i0:i1], nbits[i0:i1], starts[i0:i1] + for k in range(int(widths.max())): + m = widths > k + bits[offs[m] + k] = (vals[m] >> np.uint64(k)) & np.uint64(1) return np.packbits(bits, bitorder="little").tobytes() @@ -293,10 +307,18 @@ def put(v: int, n: int) -> None: # --- token stream as (value, nbits) entries: two per literal (green+red, # blue+alpha packed pairwise, <=30 bits each) and one per reference # (length code + extra bits + distance code, <=40 bits). + # The header rides in the same buffer rather than being concatenated onto + # the front of it: `put` has finished by now, so its length is known, and + # joining afterwards would hold a second copy of the whole token stream. + nhead = len(head_vals) per_seg = 2 + nref - offsets = np.cumsum(per_seg) - per_seg - ev = np.zeros(int(per_seg.sum()), np.uint64) + offsets = np.cumsum(per_seg) + offsets -= per_seg + offsets += nhead + ev = np.zeros(nhead + int(per_seg.sum()), np.uint64) eb = np.zeros(ev.size, np.uint8) + ev[:nhead] = head_vals + eb[:nhead] = head_bits gi, ri, bi, ai = lit[:, 1], lit[:, 0], lit[:, 2], lit[:, 3] ev[offsets] = g_code[gi] | (r_code[ri] << g_len[gi].astype(np.uint64)) eb[offsets] = g_len[gi] + r_len[ri] @@ -312,10 +334,7 @@ def put(v: int, n: int) -> None: ) eb[pos] = g_len[ref_sym] + _LP_EBITS[run] + d_len[1] - payload = _pack_lsb( - np.concatenate([np.asarray(head_vals, np.uint64), ev]), - np.concatenate([np.asarray(head_bits, np.uint8), eb]), - ) + payload = _pack_lsb(ev, eb) chunk = b"VP8L" + struct.pack(" 4 * (1 << 18) + + +# --- narrow token widths and the YCbCr band size ----------------------------- +# +# The token pipeline carries every field at its natural width (uint8 symbols, +# int16 amplitudes, int32 block indices) rather than as int64. That is only +# sound because the coefficient range is bounded, and only invisible because +# the RGB->YCbCr band size cannot change the arithmetic. Both are pinned here. + + +def test_quantized_coefficients_fit_int16(): + """|DCT coefficient| <= 1024 for any 8-bit input, so int16 is exact. + + `_DCT` is orthonormal, so the 2-D transform preserves the Frobenius norm: + no coefficient can exceed the norm of a level-shifted 8x8 block, which is + at most sqrt(64) * 128 = 1024. Quantizers are >= 1, so quantization only + shrinks it. This is the invariant the int16 coefficient buffer rests on. + """ + d = _jpeg._DCT.astype(np.float64) + np.testing.assert_allclose(d @ d.T, np.eye(8), atol=1e-6) + + rng = np.random.default_rng(0) + extremes = [ + np.full((8, 8), -128.0), # maximal negative DC + np.full((8, 8), 127.0), # maximal positive DC + (((np.arange(8)[:, None] + np.arange(8)[None, :]) % 2) * 255.0) - 128.0, + np.where(np.arange(64).reshape(8, 8) % 2, 127.0, -128.0), + *(rng.integers(0, 256, (8, 8)).astype(np.float64) - 128.0 for _ in range(64)), + ] + blocks = np.stack(extremes).astype(np.float32) + coef = _jpeg._DCT @ blocks @ _jpeg._DCT.T + assert np.abs(coef).max() <= 1024.0 + assert np.iinfo(np.int16).min < -1024 and np.iinfo(np.int16).max > 1024 + + +def test_token_fields_keep_their_narrow_widths(): + """Token arrays are byte/short-wide; widening them silently costs ~4x.""" + rng = np.random.default_rng(3) + coef = np.zeros((64, 64), dtype=np.int16) + coef[:, 0] = rng.integers(-1024, 1025, 64) + coef[:, 1:] = rng.integers(-3, 4, (64, 63)) + block, seq, table, symbol, ampl, abits = _jpeg._component_tokens(coef, 0, 1) + assert block.dtype == np.int32 + assert seq.dtype == np.uint8 + assert table.dtype == np.uint8 + assert symbol.dtype == np.uint8 + assert ampl.dtype == np.int16 + assert abits.dtype == np.uint8 + # The Huffman tables are gathered once per token, so they are narrow too. + assert _jpeg._HUFF_CODES.dtype == np.int32 + assert _jpeg._HUFF_LENS.dtype == np.uint8 + # Sequence numbers must stay inside uint8: 0 = DC, 255 = EOB, AC between. + assert seq.max() <= 255 and (seq[table == 0] == 0).all() + + +@pytest.mark.parametrize("band_px", [1, 2, 7, 64, 1 << 12, 1 << 30]) +def test_ycbcr_band_size_is_transparent(monkeypatch, band_px): + """Row banding is elementwise, so the band size cannot move a byte.""" + img = with_alpha(chart_rgb(53, 67)) + monkeypatch.setattr(_jpeg, "_YCBCR_PIXEL_CHUNK", 1 << 30) + one_shot = _jpeg.encode(img, quality=90) + monkeypatch.setattr(_jpeg, "_YCBCR_PIXEL_CHUNK", band_px) + assert _jpeg.encode(img, quality=90) == one_shot + + +def test_dc_differential_survives_maximal_swings(): + """Alternating black/white 8x8 blocks maximize the DC differential. + + At quality 100 every quantizer is 1, so this is the widest DC difference + the encoder can ever see — the case that decides whether int16 coefficient + and amplitude buffers are wide enough. + """ + blk = (np.arange(128)[:, None] // 8 + np.arange(128)[None, :] // 8) % 2 + img = with_alpha(np.repeat((blk * 255).astype(np.uint8)[..., None], 3, axis=2)) + out = decode(_jpeg.encode(img, quality=100)) + assert out.shape == (128, 128, 3) + # Corners are flat 8x8 fills, so they survive quantization near-exactly. + assert out[:4, :4].mean() < 40 and out[:4, 8:12].mean() > 215 diff --git a/tests/test_webp.py b/tests/test_webp.py index 23846ba0..364acbe9 100644 --- a/tests/test_webp.py +++ b/tests/test_webp.py @@ -140,3 +140,47 @@ def test_single_color_compresses(): img = np.full((256, 256, 4), (10, 30, 200, 255), np.uint8) out = _assert_roundtrip(img) assert len(out) < 5000 # 256 KiB raw; runs + Huffman crush it + + +@pytest.mark.parametrize("chunk", [1, 2, 17, 512, 1 << 30]) +def test_pack_entry_chunk_is_transparent(monkeypatch, chunk): + """Every entry scatters to absolute bit positions, so blocking is a no-op. + + The bit-scatter passes run over bounded entry blocks to keep their masks + and gathers off the whole token stream; the block size must not be able to + change a single emitted byte. + """ + img = _chart_like(64, 96) + monkeypatch.setattr(_webp, "_PACK_ENTRY_CHUNK", 1 << 30) + one_shot = _webp.encode(img) + monkeypatch.setattr(_webp, "_PACK_ENTRY_CHUNK", chunk) + assert _webp.encode(img) == one_shot + + +def test_pack_lsb_matches_a_scalar_reference(monkeypatch): + """The blocked packer against a plain per-symbol LSB-first reference.""" + rng = np.random.default_rng(9) + nbits = rng.integers(1, 41, 5000).astype(np.uint8) + raw = rng.integers(0, 1 << 40, 5000).astype(np.uint64) + values = raw & ((np.uint64(1) << nbits.astype(np.uint64)) - np.uint64(1)) + acc, nacc, out = 0, 0, bytearray() + for v, n in zip(values.tolist(), nbits.tolist(), strict=True): + acc |= v << nacc + nacc += n + while nacc >= 8: + out.append(acc & 0xFF) + acc >>= 8 + nacc -= 8 + if nacc: + out.append(acc & 0xFF) + monkeypatch.setattr(_webp, "_PACK_ENTRY_CHUNK", 64) + assert _webp._pack_lsb(values, nbits) == bytes(out) + + +def test_header_shares_the_token_buffer(): + """The header rides in `ev`/`eb`, so no second copy of the stream exists.""" + img = _chart_like(48, 64) + data = _webp.encode(img) + # Framing still parses and the payload still decodes bit-exactly. + assert data[:4] == b"RIFF" and data[8:12] == b"WEBP" + np.testing.assert_array_equal(_decode(data), img) From a24521cbe5dde6ade9958641196981650b8111dd Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 01:12:03 -0700 Subject: [PATCH 2/2] Flatten alpha in bands, and stop carrying a plane JPEG never reads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_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. --- python/xy/export.py | 30 ++++++++++++++++++++++++------ 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/python/xy/export.py b/python/xy/export.py index f44ef604..dc015e09 100644 --- a/python/xy/export.py +++ b/python/xy/export.py @@ -921,15 +921,33 @@ def _export_dimensions( return w, h +#: Pixels per band of the alpha flatten. The composite is elementwise, so the +#: band size cannot change a byte; it bounds the six uint16 full-frame +#: temporaries the one-shot form held (117 MB on a 3200x2400 export). +_FLATTEN_PIXEL_CHUNK = 1 << 18 + + def _flatten_alpha(rgba: "Any") -> "Any": - """Composite leftover alpha over white — the JPEG determinism backstop.""" + """Composite leftover alpha over white — the JPEG determinism backstop. + + Returns `(h, w, 3)`: both callers hand the result straight to the JPEG + encoder, which ignores alpha, so a fourth opaque plane is a quarter of a + frame carried for nothing. Composited in row bands — `rgb * a + 255 * + (255 - a) + 127` peaks at 65152 for any 8-bit input, so the uint16 + intermediates are exact, and banding an elementwise expression is + bit-identical. + """ import numpy as np - alpha = rgba[..., 3:4].astype(np.uint16) - rgb = (rgba[..., :3].astype(np.uint16) * alpha + 255 * (255 - alpha) + 127) // 255 - out = np.empty_like(rgba) - out[..., :3] = rgb.astype(np.uint8) - out[..., 3] = 255 + h, w = rgba.shape[0], rgba.shape[1] + out = np.empty((h, w, 3), dtype=np.uint8) + band = max(1, _FLATTEN_PIXEL_CHUNK // max(w, 1)) + for r0 in range(0, h, band): + r1 = min(r0 + band, h) + chunk = rgba[r0:r1] + alpha = chunk[..., 3:4].astype(np.uint16) + rgb = (chunk[..., :3].astype(np.uint16) * alpha + 255 * (255 - alpha) + 127) // 255 + out[r0:r1] = rgb.astype(np.uint8) return out