Skip to content

Commit 41e4280

Browse files
matmenclaude
matmen
and
claude
committed
JPEG XL: faster large DCTs, blending and Modular; lighter tests
Cache the default dequantization matrices instead of computing them per frame, let the inverse DCT recursion end in the unrolled 16/32-point ones, blend rows instead of pixels and compute only the MA tree properties a tree uses. Same output. The tests drop two large conformance files and check only the first frame of the spline animation, so they fit the coverage run's time limit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Lfv41PqHEGbU7SLg5onpG1
1 parent dad4a48 commit 41e4280

15 files changed

Lines changed: 152 additions & 144 deletions

‎THIRD_PARTY_NOTICES.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ following projects, under their own licenses. Each derived file says so in its h
1717
| [libavif](https://github.com/AOMediaCodec/libavif) by Joe Drago | `src/codecs/avif.ts` (YUV to RGB conversion with bilinear chroma upsampling, alpha range expansion and unpremultiplication) | BSD-2-Clause |
1818
| [jxl-oxide](https://github.com/tirr-c/jxl-oxide) by Wonwoo Choi | `src/codecs/jxl/*` (the JPEG XL decoder: bit reader, header and frame header fields, entropy decoding, VarDCT, dequantization matrices, transforms, restoration filters, upsampling, the parsing of patches, splines and noise) | MIT (or Apache-2.0) |
1919
| [libjxl](https://github.com/libjxl/libjxl) by the JPEG XL Project Authors | `src/codecs/jxl/modular.ts` (Modular decoding: MA trees, predictors, the self-correcting predictor, transforms), `src/codecs/jxl/frame.ts` (the sections of frames, the order of the rendering stages, reference frames), `src/codecs/jxl/blend.ts`, `src/codecs/jxl/color.ts` (XYB to RGB), `src/codecs/jxl/features.ts` (rendering patches, splines and noise) | BSD-3-Clause |
20-
| [JPEG XL conformance](https://github.com/libjxl/conformance) by the JPEG XL Project Authors | `tests/jxl/alpha-*.jxl`, `animation-spline.jxl`, `blendmodes.jxl`, `mul-no-extra-channels.jxl`, `sunset-logo.jxl` (test files) | BSD-3-Clause |
20+
| [JPEG XL conformance](https://github.com/libjxl/conformance) by the JPEG XL Project Authors | `tests/jxl/alpha-*.jxl`, `animation-spline.jxl`, `blendmodes.jxl`, `mul-no-extra-channels.jxl` (test files) | BSD-3-Clause |
2121
| [FreeType](https://freetype.org) | `src/render/hinting.ts` (behavior and parts of the logic of the TrueType interpreter) | FreeType License |
2222
| [QOI](https://github.com/phoboslab/qoi) by Dominic Szablewski | `src/codecs/qoi.ts` (the reference encoder's choices of operations) | MIT |
2323
| [Pillow](https://github.com/python-pillow/Pillow) | `src/util/resample.ts` (filter coefficients, box, bilinear and bicubic filters) | MIT-CMU |

‎scripts/jxl-fixtures.py‎

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
# Writes the JPEG XLs of tests/jxl and their references: parts of tests/targets encoded by cjxl (libjxl) with its coding
22
# tools (lossless efforts, palettes, VarDCT, lossy Modular, alpha, progressive passes and LF frames, resampling, photon
33
# noise, patches, JPEG recompression, animations, the container), a few conformance files (github.com/libjxl/conformance,
4-
# BSD-3-Clause, for what cjxl doesn't write: splines, blend modes, small trees, premultiplied alpha), and what djxl decodes them to as 8-bit
5-
# sRGB PNGs (APNGs for animations; premultiplied alpha as djxl writes it). Needs Pillow, cjxl and djxl:
4+
# BSD-3-Clause, for what cjxl doesn't write: splines, blend modes, small trees, premultiplied alpha), and what djxl
5+
# decodes them to as 8-bit sRGB PNGs (APNGs for cjxl's animations; premultiplied alpha as djxl writes it). Needs Pillow,
6+
# cjxl and djxl:
67
# CJXL=path/to/cjxl DJXL=path/to/djxl CONFORMANCE=path/to/conformance python scripts/jxl-fixtures.py
78
import json, os, shutil, struct, subprocess, tempfile
89
from pathlib import Path
@@ -82,28 +83,23 @@ def write_pfm(path, im):
8283
frames[0].save(tmp / 'anim.png', save_all=True, append_images=frames[1:], duration=[40, 80, 40, 120], loop=0)
8384
encode('animation', tmp / 'anim.png', '-d', '0')
8485
encode('animation-lossy', tmp / 'anim.png', '-d', '2')
85-
# conformance files: [the step of the reference's samples, its frames]
86-
conformances = {'alpha_nonpremultiplied': [4, [0]], 'alpha_premultiplied': [8, [0]], 'alpha_triangles': [4, [0]], 'blendmodes': [8, [0]], 'sunset_logo': [11, [0]],
87-
'mul_no_extra_channels': [1, [0]], 'animation_spline': [4, [0, 17, 59]]}
86+
# conformance files: the step of the reference's samples
87+
conformances = {'alpha_nonpremultiplied': 4, 'alpha_premultiplied': 8, 'blendmodes': 8, 'mul_no_extra_channels': 1,
88+
'animation_spline': 4}
8889
for name in conformances:
8990
shutil.copy(conformance / 'testcases' / name / 'input.jxl', out / f'{name.replace("_", "-")}.jxl')
9091

91-
# the references (the large conformance ones as every step-th sample of some frames, stacked)
92+
# the references (of the conformance files every step-th sample of the first frame)
9293
manifest = {}
9394
for jxl in sorted(out.glob('*.jxl')):
9495
png = jxl.with_suffix('.png')
9596
subprocess.run([DJXL, str(jxl), str(png), '--bits_per_sample=8', '--color_space=RGB_D65_SRG_Rel_SRG', '--quiet'],
9697
check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
97-
sampled = conformances.get(jxl.stem.replace('-', '_'))
98-
if sampled:
99-
step, frames = sampled
100-
im = Image.open(png)
101-
parts = []
102-
for f in frames:
103-
im.seek(f)
104-
parts.append(np.asarray(im.convert('RGBA'))[::step, ::step])
105-
im.close()
106-
Image.fromarray(np.concatenate(parts)).save(png, optimize=True)
107-
manifest[jxl.stem] = {'step': step, 'frames': frames}
98+
step = conformances.get(jxl.stem.replace('-', '_'))
99+
if step:
100+
with Image.open(png) as im:
101+
sample = np.asarray(im.convert('RGBA'))[::step, ::step]
102+
Image.fromarray(sample).save(png, optimize=True)
103+
manifest[jxl.stem] = {'step': step}
108104
print(jxl.name, jxl.stat().st_size, png.stat().st_size)
109105
(out / 'sampled.json').write_text(json.dumps(manifest, indent=4) + '\n')

‎scripts/jxl-idct.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ def idct(vals, n, secname):
5151
vals = [f'v{i}' for i in range(n)]
5252
out = idct(vals, n, f'SEC{n}')
5353
print(f'// the inverse {n}-point DCT of buf[o + k * step] in place (generated: dct\'s recursion, unrolled)')
54-
print(f'export function idct{n}(buf: Float32Array, o: number, step: number): void {{')
54+
print(f'export function idct{n}(buf: Float32Array | Float64Array, o: number, step: number): void {{')
5555
for i in range(n):
5656
print(f' const v{i} = buf[o + {i} * step];' if i else ' const v0 = buf[o];')
5757
for l in lines: print(l)

‎src/codecs/jxl/blend.ts‎

Lines changed: 67 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@ export interface Blending {
1515
}
1616

1717
const clamp01 = (x: number): number => x < 0 ? 0 : x > 1 ? 1 : x;
18-
const tmp = new Float32Array(256 + 3);
18+
// (a row per channel: the extra channels are blended with the alpha from before blending)
19+
let rows: Float32Array[] = [];
1920

2021
/**
2122
* Blends n pixels of fg (channels at fi) onto bg (at bi) into out (at oi; may be bg); channels: numColor color ones,
@@ -24,63 +25,83 @@ const tmp = new Float32Array(256 + 3);
2425
export function performBlending(bg: Float32Array[], bi: number, fg: Float32Array[], fi: number, out: Float32Array[], oi: number, n: number,
2526
numColor: number, color: Blending, ec: Blending[], extra: ExtraChannel[]): void {
2627
const hasAlpha = extra.some(e => e.type === EXTRA_ALPHA), numEc = extra.length, nc = numColor + numEc;
27-
const t = nc <= tmp.length ? tmp : new Float32Array(nc);
28-
for (let k = 0; k < n; k++) {
29-
const b = bi + k, f = fi + k;
30-
// extra channels first, with the alpha before blending
31-
for (let i = 0; i < numEc; i++) {
32-
const c = numColor + i, e = ec[i], a = numColor + e.alpha, bgv = bg[c][b], fgv = fg[c][f];
33-
switch (e.mode) {
34-
case PB_ADD: t[c] = bgv + fgv; break;
35-
case PB_BLEND_ABOVE: t[c] = alphaBlend1(bgv, bg[a][b], fgv, fg[a][f], c === a, extra[e.alpha].alphaAssociated, e.clamp); break;
36-
case PB_BLEND_BELOW: t[c] = alphaBlend1(fgv, fg[a][f], bgv, bg[a][b], c === a, extra[e.alpha].alphaAssociated, e.clamp); break;
37-
case PB_ADD_ABOVE: t[c] = c === a ? bgv : bgv + fgv * (e.clamp ? clamp01(fg[a][f]) : fg[a][f]); break;
38-
case PB_ADD_BELOW: t[c] = c === a ? fgv : fgv + bgv * (e.clamp ? clamp01(bg[a][b]) : bg[a][b]); break;
39-
case PB_MUL: t[c] = bgv * (e.clamp ? clamp01(fgv) : fgv); break;
40-
case PB_REPLACE: t[c] = fgv; break;
41-
default: t[c] = bgv;
42-
}
43-
}
44-
const a = numColor + color.alpha;
45-
switch (color.mode) {
46-
case PB_ADD:
47-
for (let c = 0; c < numColor; c++) t[c] = bg[c][b] + fg[c][f];
28+
if (rows.length < nc || rows[0].length < n) rows = Array.from({ length: Math.max(nc, rows.length) }, () => new Float32Array(Math.max(n, rows[0]?.length ?? 0)));
29+
for (let i = 0; i < numEc; i++) {
30+
const c = numColor + i, e = ec[i], a = numColor + e.alpha, t = rows[c], B = bg[c], F = fg[c], BA = bg[a], FA = fg[a];
31+
const isAlpha = c === a, premultiplied = extra[e.alpha].alphaAssociated;
32+
switch (e.mode) {
33+
case PB_ADD: for (let k = 0; k < n; k++) t[k] = B[bi + k] + F[fi + k]; break;
34+
case PB_BLEND_ABOVE:
35+
for (let k = 0; k < n; k++) t[k] = alphaBlend1(B[bi + k], BA[bi + k], F[fi + k], FA[fi + k], isAlpha, premultiplied, e.clamp);
36+
break;
37+
case PB_BLEND_BELOW:
38+
for (let k = 0; k < n; k++) t[k] = alphaBlend1(F[fi + k], FA[fi + k], B[bi + k], BA[bi + k], isAlpha, premultiplied, e.clamp);
4839
break;
4940
case PB_ADD_ABOVE:
50-
case PB_ADD_BELOW: {
51-
const [lo, li, hi, hiIdx] = color.mode === PB_ADD_ABOVE ? [bg, b, fg, f] : [fg, f, bg, b];
52-
for (let c = 0; c < numColor; c++) {
53-
t[c] = hasAlpha ? lo[c][li] + hi[c][hiIdx] * (color.clamp ? clamp01(hi[a][hiIdx]) : hi[a][hiIdx]) : bg[c][b] + fg[c][f];
54-
}
41+
for (let k = 0; k < n; k++) t[k] = isAlpha ? B[bi + k] : B[bi + k] + F[fi + k] * (e.clamp ? clamp01(FA[fi + k]) : FA[fi + k]);
5542
break;
43+
case PB_ADD_BELOW:
44+
for (let k = 0; k < n; k++) t[k] = isAlpha ? F[fi + k] : F[fi + k] + B[bi + k] * (e.clamp ? clamp01(BA[bi + k]) : BA[bi + k]);
45+
break;
46+
case PB_MUL: for (let k = 0; k < n; k++) t[k] = B[bi + k] * (e.clamp ? clamp01(F[fi + k]) : F[fi + k]); break;
47+
case PB_REPLACE: t.set(F.subarray(fi, fi + n)); break;
48+
default: t.set(B.subarray(bi, bi + n));
49+
}
50+
}
51+
const a = numColor + color.alpha;
52+
switch (color.mode) {
53+
case PB_ADD:
54+
for (let c = 0; c < numColor; c++) {
55+
const t = rows[c], B = bg[c], F = fg[c];
56+
for (let k = 0; k < n; k++) t[k] = B[bi + k] + F[fi + k];
5657
}
57-
case PB_BLEND_ABOVE:
58-
case PB_BLEND_BELOW: {
58+
break;
59+
case PB_ADD_ABOVE:
60+
case PB_ADD_BELOW: {
61+
const [lo, li, hi, hj] = color.mode === PB_ADD_ABOVE ? [bg, bi, fg, fi] : [fg, fi, bg, bi];
62+
for (let c = 0; c < numColor; c++) {
63+
const t = rows[c], L = lo[c], H = hi[c];
5964
if (!hasAlpha) {
60-
for (let c = 0; c < numColor; c++) t[c] = fg[c][f];
61-
break;
65+
for (let k = 0; k < n; k++) t[k] = bg[c][bi + k] + fg[c][fi + k];
66+
continue;
6267
}
63-
const [lo, li, hi, hiIdx] = color.mode === PB_BLEND_ABOVE ? [bg, b, fg, f] : [fg, f, bg, b];
64-
const fa = color.clamp ? clamp01(hi[a][hiIdx]) : hi[a][hiIdx], ba = lo[a][li], na = 1 - (1 - fa) * (1 - ba);
65-
if (extra[color.alpha].alphaAssociated) for (let c = 0; c < numColor; c++) t[c] = hi[c][hiIdx] + lo[c][li] * (1 - fa);
68+
const HA = hi[a];
69+
for (let k = 0; k < n; k++) t[k] = L[li + k] + H[hj + k] * (color.clamp ? clamp01(HA[hj + k]) : HA[hj + k]);
70+
}
71+
break;
72+
}
73+
case PB_BLEND_ABOVE:
74+
case PB_BLEND_BELOW: {
75+
if (!hasAlpha) {
76+
for (let c = 0; c < numColor; c++) rows[c].set(fg[c].subarray(fi, fi + n));
77+
break;
78+
}
79+
const [lo, li, hi, hj] = color.mode === PB_BLEND_ABOVE ? [bg, bi, fg, fi] : [fg, fi, bg, bi];
80+
const premultiplied = extra[color.alpha].alphaAssociated, HA = hi[a], LA = lo[a], ta = rows[a];
81+
for (let k = 0; k < n; k++) {
82+
const fa = color.clamp ? clamp01(HA[hj + k]) : HA[hj + k], ba = LA[li + k], na = 1 - (1 - fa) * (1 - ba);
83+
if (premultiplied) for (let c = 0; c < numColor; c++) rows[c][k] = hi[c][hj + k] + lo[c][li + k] * (1 - fa);
6684
else {
6785
const rna = na > 0 ? 1 / na : 0;
68-
for (let c = 0; c < numColor; c++) t[c] = (hi[c][hiIdx] * fa + lo[c][li] * ba * (1 - fa)) * rna;
86+
for (let c = 0; c < numColor; c++) rows[c][k] = (hi[c][hj + k] * fa + lo[c][li + k] * ba * (1 - fa)) * rna;
6987
}
70-
t[a] = na;
71-
break;
88+
ta[k] = na;
7289
}
73-
case PB_MUL:
74-
for (let c = 0; c < numColor; c++) t[c] = bg[c][b] * (color.clamp ? clamp01(fg[c][f]) : fg[c][f]);
75-
break;
76-
case PB_REPLACE:
77-
for (let c = 0; c < numColor; c++) t[c] = fg[c][f];
78-
break;
79-
default:
80-
for (let c = 0; c < numColor; c++) t[c] = bg[c][b];
90+
break;
8191
}
82-
for (let c = 0; c < nc; c++) out[c][oi + k] = t[c];
92+
case PB_MUL:
93+
for (let c = 0; c < numColor; c++) {
94+
const t = rows[c], B = bg[c], F = fg[c];
95+
for (let k = 0; k < n; k++) t[k] = B[bi + k] * (color.clamp ? clamp01(F[fi + k]) : F[fi + k]);
96+
}
97+
break;
98+
case PB_REPLACE:
99+
for (let c = 0; c < numColor; c++) rows[c].set(fg[c].subarray(fi, fi + n));
100+
break;
101+
default:
102+
for (let c = 0; c < numColor; c++) rows[c].set(bg[c].subarray(bi, bi + n));
83103
}
104+
for (let c = 0; c < nc; c++) out[c].set(rows[c].subarray(0, n), oi);
84105
}
85106

86107
// alpha blending of a single channel (libjxl's clamps when not asked to)

‎src/codecs/jxl/dct.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,10 @@ export function dct(io: Float64Array, o: number, n: number, scratch: Float64Arra
6262
dct4(io, o, inverse);
6363
return;
6464
}
65+
if (inverse && (n === 16 || n === 32)) {
66+
(n === 16 ? idct16 : idct32)(io, o, 1);
67+
return;
68+
}
6569
const half = n / 2, sec = secHalf(n);
6670
if (!inverse) {
6771
for (let i = 0; i < half; i++) {

‎src/codecs/jxl/idct.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ const secants = (n: number): Float64Array => Float64Array.from({ length: n / 2 }
1010
const SEC8 = secants(8), SEC16 = secants(16), SEC32 = secants(32);
1111

1212
// the inverse 16-point DCT of buf[o + k * step] in place (generated: dct's recursion, unrolled)
13-
export function idct16(buf: Float32Array, o: number, step: number): void {
13+
export function idct16(buf: Float32Array | Float64Array, o: number, step: number): void {
1414
const v0 = buf[o];
1515
const v1 = buf[o + 1 * step];
1616
const v2 = buf[o + 2 * step];
@@ -150,7 +150,7 @@ export function idct16(buf: Float32Array, o: number, step: number): void {
150150
}
151151

152152
// the inverse 32-point DCT of buf[o + k * step] in place (generated: dct's recursion, unrolled)
153-
export function idct32(buf: Float32Array, o: number, step: number): void {
153+
export function idct32(buf: Float32Array | Float64Array, o: number, step: number): void {
154154
const v0 = buf[o];
155155
const v1 = buf[o + 1 * step];
156156
const v2 = buf[o + 2 * step];

‎src/codecs/jxl/modular.ts‎

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -630,10 +630,14 @@ function decodeChannel(r: BitReader, dec: EntropyDecoder, globalTree: Tree, imag
630630
const clusters = dec.clusters;
631631
// what the tree needs
632632
let maxProp = -1, usesWP = false;
633+
const uses = new Uint8Array(16);
633634
for (let i = 0; i < t.property.length; i++) {
634635
if (t.property[i] > maxProp) maxProp = t.property[i];
636+
if (t.property[i] >= 0 && t.property[i] < 16) uses[t.property[i]] = 1;
635637
if (t.property[i] === 15 || (t.property[i] < 0 && t.predictor[i] === 6)) usesWP = true;
636638
}
639+
// (only the properties the tree asks for are computed; 8 needs 9 of the pixel before)
640+
const neighborProps = uses.subarray(4).some(u => u === 1), gradient = uses[8] === 1 || uses[9] === 1;
637641
if (t.property[0] < 0) {
638642
// a single leaf: no properties
639643
const pred = t.predictor[0], offset = t.offset[0], mult = t.multiplier[0], cluster = clusters[t.ctx[0]];
@@ -677,17 +681,21 @@ function decodeChannel(r: BitReader, dec: EntropyDecoder, globalTree: Tree, imag
677681
const leftleft = x > 1 ? data[i - 2] : left, toptop = y > 1 ? data[i - 2 * w] : top;
678682
const toprightright = x + 2 < w && y ? data[i - w + 2] : topright;
679683
props[3] = x;
680-
props[4] = Math.abs(top);
681-
props[5] = Math.abs(left);
682-
props[6] = top;
683-
props[7] = left;
684-
props[8] = left - props[9];
685-
props[9] = left + top - topleft;
686-
props[10] = left - topleft;
687-
props[11] = topleft - top;
688-
props[12] = top - topright;
689-
props[13] = top - toptop;
690-
props[14] = left - leftleft;
684+
if (neighborProps) {
685+
props[4] = Math.abs(top);
686+
props[5] = Math.abs(left);
687+
props[6] = top;
688+
props[7] = left;
689+
if (gradient) {
690+
props[8] = left - props[9];
691+
props[9] = left + top - topleft;
692+
}
693+
props[10] = left - topleft;
694+
props[11] = topleft - top;
695+
props[12] = top - topright;
696+
props[13] = top - toptop;
697+
props[14] = left - leftleft;
698+
}
691699
const wpPred = wp ? wp.predict(x, y, top, left, topright, topleft, toptop) : 0;
692700
if (wp) props[15] = wp.maxError;
693701
for (let k = 0; k < refs.length; k++) {

0 commit comments

Comments
 (0)