Skip to content

Commit e8888eb

Browse files
matmenclaude
andcommitted
JXL: lossless frames straight to bytes, in threads
Frames of integer samples shown as they are (not XYB, no filters, features or upsampling, straight alpha) become RGBA bytes right away instead of float planes that toRGBA converts again: bands of 64 rows in the threads (one pass: 32-bit words from a table of the bytes of the same Float32 values), with an RCT applied first to the colors undone in the same pass (ModularImage.undo(keep), rctSamples/rctOrder). The same pixels on 98 files (tests, conformance suite, jxl-oxide, cjxl -e 1-9 grey/16-bit/alpha/palette) with 1 and 4 threads. 3840x2160 from encodeJXL(), Ryzen 7 5800X: 379 -> 304 ms with 8 threads, 581 -> 548 ms with one; cjxl -e 7 (RCTs per group) 1113 -> 1073 ms. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BsU2GAqbnFSNaYDBGD1Km2
1 parent c968009 commit e8888eb

4 files changed

Lines changed: 172 additions & 47 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,10 @@
2424
instead of 640 ms lossy and 328 ms instead of 671 ms lossless, a 5120x2880 one in 568 ms instead of 1121 ms and
2525
487 ms instead of 1122 ms. Lossy frames no longer keep the coefficients of the whole frame (about 100 MB at 4K): each
2626
group is rendered once its passes are read.
27+
- Lossless JPEG XLs shown as they are (integer samples, no filters or upsampling) become RGBA bytes right away instead
28+
of through float planes, in threads from 1 megapixel on, with an RCT of the colors undone in the same pass: the same
29+
pixels. A 3840x2160 image from `encodeJXL()` decodes in 304 ms instead of 379 ms with 8 threads, 548 ms instead of
30+
581 ms with one.
2731

2832
### Tests
2933

‎src/codecs/jxl/color.ts‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import { toSRGB, hdrToLinear, ToneMapping, REFERENCE_WHITE } from '../hdr.ts';
88
import type { Planes } from './frame.ts';
99

1010
// a sample (0-1) as a byte
11-
const byte = (v: number): number => v <= 0 ? 0 : v >= 1 ? 255 : Math.round(v * 255);
11+
export const byte = (v: number): number => v <= 0 ? 0 : v >= 1 ? 255 : Math.round(v * 255);
1212

1313
// linear light to sRGB like libjxl (TF_SRGB::EncodedFromDisplay, lib/jxl/cms/transfer_functions-inl.h): a rational
1414
// polynomial of the square root instead of the power, within 5e-7
@@ -87,6 +87,8 @@ export class ColorOutput {
8787
* them (straight alpha, no black or spot channels), then XYB ones become interleaved sRGB bytes right away
8888
*/
8989
transform(p: Planes, ycbcr: boolean, final = false): Planes {
90+
// (bytes already: see modularToBytes)
91+
if (p.rgba) return p;
9092
const image = this.image, n = p.width * p.height, ch = p.channels.slice();
9193
if (ycbcr) {
9294
const [cb, y, cr] = ch, r = new Float32Array(n), g = new Float32Array(n), b = new Float32Array(n);
@@ -125,6 +127,11 @@ export class ColorOutput {
125127
return { width: p.width, height: p.height, channels: ch };
126128
}
127129

130+
/** Whether toRGBA makes bytes of the samples as they are (not XYB, not PQ or HLG) */
131+
bytesOfSamples(): boolean {
132+
return !this.image.xybEncoded && !this.hdrSignal;
133+
}
134+
128135
/** What transform converts XYB to bytes with right away (final, not YCbCr; see xybToBytes), null if it doesn't */
129136
bytesParams(ycbcr: boolean, final: boolean): { m: number[], bias: number[], k: number } | null {
130137
const image = this.image;

‎src/codecs/jxl/frame.ts‎

Lines changed: 114 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,14 @@ import {
1010
FrameHeader, TOC, frameHeader, toc, passShifts, FRAME_LF, FRAME_REGULAR, FRAME_SKIP_PROGRESSIVE, BLEND_REPLACE, BLEND_ADD,
1111
BLEND_BLEND, BLEND_MULADD, BLEND_MUL, FLAG_PATCHES, FLAG_SPLINES, FLAG_NOISE, FLAG_USE_LF_FRAME, type EPF
1212
} from './frame-header.ts';
13-
import { ModularImage, Channel, MAConfig, channel, readTree, decodeModular } from './modular.ts';
13+
import { ModularImage, Channel, MAConfig, TR_RCT, channel, readTree, decodeModular, rctOrder, rctSamples } from './modular.ts';
1414
import { upsample } from './upsample.ts';
1515
import { gaborish, gaborishRows, epf, epfStrip, epfScratch, epfSteps, epfStrips } from './filters.ts';
1616
import { VarDCT, VarDCTGlobal, readVarDCTGlobal, type VarDCTShared } from './vardct.ts';
1717
import { runChunks, shared, sharedOf, type ChunkTask, type Thread } from '../../util/threads.ts';
1818
import { Patch, Splines, readPatches, applyPatches, readSplines, renderSplines, readNoise, renderNoise } from './features.ts';
19-
import { ColorOutput, xybToBytes } from './color.ts';
19+
import { byte, ColorOutput, xybToBytes } from './color.ts';
20+
import { littleEndian } from '../../util/alpha.ts';
2021
import { Blending, performBlending, PB_REPLACE, PB_ADD, PB_BLEND_ABOVE, PB_ADD_ABOVE, PB_MUL } from './blend.ts';
2122

2223
/** Float planes: the 3 color channels (or 1), then the extra channels */
@@ -149,8 +150,15 @@ export async function decodeFrame(d: Uint8Array, r: BitReader, state: DecoderSta
149150
}
150151
}
151152
if (!groupsDone) for (let g = 0; g < f.numGroups; g++) local(g);
152-
lf.modular.undo();
153-
let planes = modularToPlanes(lf, f, image);
153+
// (a lossless frame shown as it is: its samples to bytes right away)
154+
const bytes = !vardct && f.modular && !f.doYCbCr && !f.gab && !f.epf && !lf.patches && !lf.splines && !lf.noise && f.upsampling === 1
155+
&& f.ecUpsampling.every(u => u === 1) && shownAsIs(f, image) && state.output.bytesOfSamples() && !image.bitDepth.float;
156+
let planes: Planes;
157+
if (bytes) planes = await modularToBytes(lf, f, image, pool);
158+
else {
159+
lf.modular.undo();
160+
planes = modularToPlanes(lf, f, image);
161+
}
154162
let sigma: Float32Array | null = null;
155163
if (vardct) {
156164
const xyb = vardct.rendered;
@@ -313,18 +321,21 @@ function decodeModularGroup(r: BitReader, lf: LfGlobal, f: FrameHeader, x0: numb
313321
}
314322

315323
// a Modular image's channels as floats: [0, 1] for integers, cast for float samples; XYB scaled by the LF dequantization
324+
// a channel's samples as floats (integers scaled to 0 to 1)
325+
function samplesToFloat(c: Channel, depth: ImageHeader['bitDepth']): Float32Array {
326+
const o = new Float32Array(c.w * c.h);
327+
if (depth.float) for (let i = 0; i < o.length; i++) o[i] = floatSample(c.data[i], depth.bits, depth.expBits);
328+
else {
329+
const scale = 1 / (2 ** depth.bits - 1);
330+
for (let i = 0; i < o.length; i++) o[i] = c.data[i] * scale;
331+
}
332+
return o;
333+
}
334+
316335
function modularToPlanes(lf: LfGlobal, f: FrameHeader, image: ImageHeader): Planes {
317336
const ch = lf.modular.channels.slice(lf.modular.nbMeta);
318337
const out: Float32Array[] = [];
319-
const toFloat = (c: Channel, depth = image.bitDepth): Float32Array => {
320-
const o = new Float32Array(c.w * c.h);
321-
if (depth.float) for (let i = 0; i < o.length; i++) o[i] = floatSample(c.data[i], depth.bits, depth.expBits);
322-
else {
323-
const scale = 1 / (2 ** depth.bits - 1);
324-
for (let i = 0; i < o.length; i++) o[i] = c.data[i] * scale;
325-
}
326-
return o;
327-
};
338+
const toFloat = (c: Channel, depth = image.bitDepth): Float32Array => samplesToFloat(c, depth);
328339
if (!f.modular) {
329340
// (the extra channels only)
330341
image.extraChannels.forEach((ec, i) => out.push(toFloat(ch[i], ec.bitDepth)));
@@ -347,6 +358,74 @@ function modularToPlanes(lf: LfGlobal, f: FrameHeader, image: ImageHeader): Plan
347358
return { width: f.sampleWidth, height: f.sampleHeight, channels: out };
348359
}
349360

361+
/**
362+
* The color channels of a lossless frame's Modular image (integer samples, not XYB, 1 or 3 channels) as interleaved
363+
* RGBA bytes like toRGBA makes of modularToPlanes' (the alpha of 255; the extra channels' planes as there), in threads
364+
* if there are; the transforms undone, an RCT of the colors applied first with the conversion
365+
*/
366+
async function modularToBytes(lf: LfGlobal, f: FrameHeader, image: ImageHeader, pool: FrameThreads): Promise<Planes> {
367+
const m = lf.modular, numColor = f.encodedColorChannels, n = f.width * f.height, first = m.transforms[0];
368+
const rct = numColor === 3 && first?.id === TR_RCT && first.beginC === m.nbMeta ? first : null;
369+
m.undo(rct ? 1 : 0);
370+
const color = rct ? m.rctChannels(rct) : [0, 1, 2].map(c => m.channels[m.nbMeta + (numColor === 1 ? 0 : c)].data);
371+
const task = (threaded: boolean): PostTask => ({
372+
task: 'jxl-post', kind: 'samples', w: f.width, h: f.height, inputs: [], outputs: [], samples: threaded ? color.map(sharedOf) : color,
373+
rctType: rct ? rct.rctType : -1, bits: image.bitDepth.bits, rgba: new Uint8Array(threaded ? new SharedArrayBuffer(4 * n) : new ArrayBuffer(4 * n)),
374+
});
375+
const rgba = await threaded(pool, n, async threads => {
376+
const t = task(true);
377+
await runChunks(threads, t, Math.ceil(f.height / BAND), postChunks(t));
378+
// (the samples undone in shared memory; the pixels not: APIs like crypto.subtle and TextDecoder reject views of it)
379+
if (rct) for (let c = 0; c < 3; c++) m.channels[m.nbMeta + c].data = (t.samples as Int32Array[])[c];
380+
return new Uint8Array(t.rgba as Uint8Array);
381+
}, () => {
382+
const t = task(false);
383+
samplesToBytes(t.samples as Int32Array[], t.rctType as number, t.bits as number, t.rgba as Uint8Array, 0, n);
384+
return t.rgba as Uint8Array;
385+
});
386+
if (rct) {
387+
m.permuteRCT(rct);
388+
m.transforms.length = 0;
389+
}
390+
const ch = m.channels.slice(m.nbMeta);
391+
return { width: f.width, height: f.height, channels: image.extraChannels.map((ec, i) => samplesToFloat(ch[numColor + i], ec.bitDepth)), rgba };
392+
}
393+
394+
/**
395+
* Samples i0 to i1 of 3 color channels as RGBA bytes (see modularToBytes): first undoing an RCT (of rctType, -1 for
396+
* none), the samples bits deep
397+
*/
398+
function samplesToBytes(samples: Int32Array[], rctType: number, bits: number, rgba: Uint8Array, i0: number, i1: number): void {
399+
if (rctType >= 0) rctSamples(samples, rctType, i0, i1);
400+
// (the RCT's channels go to their places)
401+
const order = rctType >= 0 ? rctOrder(rctType) : [0, 1, 2], [r, g, b] = [0, 1, 2].map(c => samples[order.indexOf(c)]);
402+
const max = 2 ** bits - 1, scale = 1 / max;
403+
// (the bytes of the samples 0 to max: through the same Float32 value as the planes)
404+
const table = bits <= 16 ? byteTable(bits) : null;
405+
if (table && littleEndian) {
406+
// (a pixel per 32-bit word)
407+
const words = new Uint32Array(rgba.buffer, rgba.byteOffset, rgba.length >> 2);
408+
const at = (v: number): number => v >>> 0 <= max ? table[v] : v < 0 ? 0 : 255;
409+
for (let i = i0; i < i1; i++) words[i] = at(r[i]) | at(g[i]) << 8 | at(b[i]) << 16 | 0xff000000;
410+
} else {
411+
[r, g, b].forEach((data, c) => {
412+
for (let i = i0, o = 4 * i0 + c; i < i1; i++, o += 4) rgba[o] = table ? (data[i] >>> 0 <= max ? table[data[i]] : data[i] < 0 ? 0 : 255) : byte(Math.fround(data[i] * scale));
413+
});
414+
for (let o = 4 * i0 + 3; o < 4 * i1; o += 4) rgba[o] = 255;
415+
}
416+
}
417+
418+
const byteTables = new Map<number, Uint8Array>();
419+
// the bytes of the samples 0 to 2^bits - 1 (see samplesToBytes)
420+
function byteTable(bits: number): Uint8Array {
421+
let table = byteTables.get(bits);
422+
if (!table) {
423+
const scale = 1 / (2 ** bits - 1);
424+
byteTables.set(bits, table = Uint8Array.from({ length: 2 ** bits }, (_, v) => byte(Math.fround(v * scale))));
425+
}
426+
return table;
427+
}
428+
350429
/** The restoration filters (sigma: per 8x8 block for VarDCT, blocks sigmaStride per row) */
351430
function filter(p: Planes, f: FrameHeader, image: ImageHeader, sigma: Float32Array | null, sigmaStride = 0): Planes {
352431
if (!f.gab && !f.epf) return p;
@@ -371,13 +450,16 @@ async function threaded<T>(pool: FrameThreads, pixels: number, inThreads: (threa
371450
return here();
372451
}
373452

374-
// a typed array in shared memory (itself if it is)
453+
// planes in shared memory
375454
const sharedPlanes = (n: number, w: number, h: number): Float32Array[] => Array.from({ length: n }, () => new Float32Array(new SharedArrayBuffer(4 * w * h)));
376455

377-
/** The rows of the steps after decoding for threads (see postChunks): Gaborish, an EPF pass, XYB to bytes */
456+
/**
457+
* The rows of the steps after decoding for threads (see postChunks): Gaborish, an EPF pass, XYB to bytes, integer
458+
* samples to bytes
459+
*/
378460
export interface PostTask extends ChunkTask {
379461
task: 'jxl-post';
380-
kind: 'gaborish' | 'epf' | 'xyb';
462+
kind: 'gaborish' | 'epf' | 'xyb' | 'samples';
381463
w: number;
382464
h: number;
383465
inputs: Float32Array[];
@@ -391,6 +473,9 @@ export interface PostTask extends ChunkTask {
391473
m?: number[];
392474
bias?: number[];
393475
k?: number;
476+
samples?: Int32Array[];
477+
rctType?: number;
478+
bits?: number;
394479
}
395480

396481
// (the rows of a chunk of Gaborish and the XYB conversion; EPF's are its strips)
@@ -404,6 +489,7 @@ export function postChunks(t: PostTask): (chunk: number) => void {
404489
const scratch = epfScratch(w);
405490
return chunk => epfStrip(inputs, outputs, w, h, t.epf as EPF, t.sigma ?? null, t.bw as number, t.step as number, chunk, scratch);
406491
}
492+
if (t.kind === 'samples') return chunk => samplesToBytes(t.samples as Int32Array[], t.rctType as number, t.bits as number, t.rgba as Uint8Array, chunk * BAND * w, Math.min(h, (chunk + 1) * BAND) * w);
407493
const [x, y, b] = inputs;
408494
return chunk => xybToBytes(x, y, b, t.m as number[], t.bias as number[], t.k as number, t.rgba as Uint8Array, chunk * BAND * w, Math.min(h, (chunk + 1) * BAND) * w);
409495
}
@@ -431,6 +517,7 @@ async function filterThreaded(p: Planes, f: FrameHeader, image: ImageHeader, sig
431517

432518
/** Upsamples the color channels and the extra channels to the frame's size */
433519
function upsamplePlanes(p: Planes, f: FrameHeader, image: ImageHeader): Planes {
520+
if (f.upsampling === 1 && f.ecUpsampling.every(u => u === 1)) return p;
434521
const numColor = p.channels.length - image.extraChannels.length, W = f.upsampledWidth, H = f.upsampledHeight;
435522
const channels = p.channels.map((c, i) => {
436523
const up = i < numColor ? f.upsampling : f.ecUpsampling[i - numColor];
@@ -459,22 +546,29 @@ function floatSample(v: number, bits: number, expBits: number): number {
459546
// ---- rendering: filters, features, upsampling, the color transform, blending, references
460547

461548
// what follows the features: storing LF frames and references, the color transform, blending
549+
// whether a frame can be saved as a reference
550+
const referenced = (f: FrameHeader): boolean => !f.isLast && (f.duration === 0 || f.saveAsReference !== 0) && f.type !== FRAME_LF;
551+
552+
// whether only toRGBA uses the planes of a frame (if of the image's size): shown as it is, replacing all, not
553+
// referenced; with straight alpha only
554+
const shownAsIs = (f: FrameHeader, image: ImageHeader): boolean => (f.type === FRAME_REGULAR || f.type === FRAME_SKIP_PROGRESSIVE) && !f.haveCrop
555+
&& f.blending.mode === BLEND_REPLACE && f.ecBlending.every(b => b.mode === BLEND_REPLACE) && f.width === image.width && f.height === image.height
556+
&& !referenced(f) && image.extraChannels.every(e => e.type === EXTRA_ALPHA && !e.alphaAssociated);
557+
462558
async function render(p: Planes, f: FrameHeader, state: DecoderState, pool: FrameThreads): Promise<Frame> {
463559
const image = state.image;
464560
let planes = p;
465561
if (f.type === FRAME_LF) {
466562
state.lfFrames[f.lfLevel - 1] = planes;
467563
return { header: f, canvas: null };
468564
}
469-
const canReference = !f.isLast && (f.duration === 0 || f.saveAsReference !== 0) && f.type !== FRAME_LF;
565+
const canReference = referenced(f);
470566
if (canReference && f.saveBeforeCT) {
471567
state.references[f.saveAsReference] = planes;
472568
state.referenceIsXYB[f.saveAsReference] = image.xybEncoded;
473569
}
474570
const normal = f.type === FRAME_REGULAR || f.type === FRAME_SKIP_PROGRESSIVE;
475-
// (only toRGBA uses the planes of a frame shown as it is: replacing all, not referenced; with straight alpha only)
476-
const replaces = !f.haveCrop && f.blending.mode === BLEND_REPLACE && f.ecBlending.every(b => b.mode === BLEND_REPLACE) && p.width === image.width && p.height === image.height;
477-
const final = normal && replaces && !canReference && image.extraChannels.every(e => e.type === EXTRA_ALPHA && !e.alphaAssociated);
571+
const final = shownAsIs(f, image) && p.width === image.width && p.height === image.height;
478572
// (XYB to bytes in threads)
479573
const bytes = state.output.bytesParams(f.doYCbCr, final), input = planes;
480574
planes = bytes ? await threaded(pool, planes.width * planes.height, async threads => {

‎src/codecs/jxl/modular.ts‎

Lines changed: 46 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -336,6 +336,35 @@ function transform(r: BitReader): Transform {
336336
}
337337

338338
/** A Modular image: channels (the first nbMeta of them meta channels) and the transforms to undo */
339+
/** Where an RCT's 3 channels go after it's undone (offsets from its first) */
340+
export function rctOrder(rctType: number): number[] {
341+
const perm = Math.floor(rctType / 7);
342+
return [perm % 3, (perm + 1 + Math.floor(perm / 3)) % 3, (perm + 2 - Math.floor(perm / 3)) % 3];
343+
}
344+
345+
/** Samples i0 to i1 of an inverse RCT (of the channels a, b, c, in place; their permutation is rctOrder's) */
346+
export function rctSamples([a, b, c]: Int32Array[], rctType: number, i0: number, i1: number): void {
347+
const type = rctType % 7;
348+
if (type === 6) {
349+
// YCoCg
350+
for (let i = i0; i < i1; i++) {
351+
const Y = a[i], Co = b[i], Cg = c[i];
352+
const tmp = (Y - (Cg >> 1)) | 0, G = (Cg + tmp) | 0, B = (tmp - (Co >> 1)) | 0;
353+
a[i] = (B + Co) | 0;
354+
b[i] = G;
355+
c[i] = B;
356+
}
357+
} else if (type !== 0) {
358+
const second = type >> 1, third = type & 1;
359+
for (let i = i0; i < i1; i++) {
360+
const first = a[i];
361+
if (third) c[i] = (c[i] + first) | 0;
362+
if (second === 1) b[i] = (b[i] + first) | 0;
363+
else if (second === 2) b[i] = (b[i] + (((first + c[i]) | 0) >> 1)) | 0;
364+
}
365+
}
366+
}
367+
339368
export class ModularImage {
340369
nbMeta = 0;
341370
transforms: Transform[] = [];
@@ -423,8 +452,9 @@ export class ModularImage {
423452
}
424453

425454
/** undoes the transforms */
426-
undo(): void {
427-
while (this.transforms.length) {
455+
/** Undoes the transforms but the first keep ones */
456+
undo(keep = 0): void {
457+
while (this.transforms.length > keep) {
428458
const t = this.transforms.pop() as Transform;
429459
if (t.id === TR_RCT) this.inverseRCT(t);
430460
else if (t.id === TR_PALETTE) this.inversePalette(t);
@@ -433,31 +463,21 @@ export class ModularImage {
433463
}
434464

435465
private inverseRCT(t: Transform): void {
466+
const c = this.channels[t.beginC];
467+
rctSamples(this.rctChannels(t), t.rctType, 0, c.w * c.h);
468+
this.permuteRCT(t);
469+
}
470+
471+
/** The channels of an RCT (of the same size), as they are before it's undone */
472+
rctChannels(t: Transform): Int32Array[] {
436473
this.checkEqual(t.beginC, t.beginC + 2);
437-
const m = t.beginC, perm = Math.floor(t.rctType / 7), type = t.rctType % 7;
438-
const ch = this.channels, c0 = ch[m], c1 = ch[m + 1], c2 = ch[m + 2];
439-
const o0 = m + perm % 3, o1 = m + (perm + 1 + Math.floor(perm / 3)) % 3, o2 = m + (perm + 2 - Math.floor(perm / 3)) % 3;
440-
if (type !== 0) {
441-
const a = c0.data, b = c1.data, c = c2.data, n = c0.w * c0.h;
442-
if (type === 6) {
443-
// YCoCg
444-
for (let i = 0; i < n; i++) {
445-
const Y = a[i], Co = b[i], Cg = c[i];
446-
const tmp = (Y - (Cg >> 1)) | 0, G = (Cg + tmp) | 0, B = (tmp - (Co >> 1)) | 0;
447-
a[i] = (B + Co) | 0;
448-
b[i] = G;
449-
c[i] = B;
450-
}
451-
} else {
452-
const second = type >> 1, third = type & 1;
453-
for (let i = 0; i < n; i++) {
454-
const first = a[i];
455-
if (third) c[i] = (c[i] + first) | 0;
456-
if (second === 1) b[i] = (b[i] + first) | 0;
457-
else if (second === 2) b[i] = (b[i] + (((first + c[i]) | 0) >> 1)) | 0;
458-
}
459-
}
460-
}
474+
return [0, 1, 2].map(i => this.channels[t.beginC + i].data);
475+
}
476+
477+
/** The permutation of an RCT's channels (after its samples are undone, see rctSamples) */
478+
permuteRCT(t: Transform): void {
479+
const m = t.beginC, ch = this.channels, [c0, c1, c2] = ch.slice(m, m + 3);
480+
const [o0, o1, o2] = rctOrder(t.rctType).map(o => m + o);
461481
ch[o0] = c0;
462482
ch[o1] = c1;
463483
ch[o2] = c2;

0 commit comments

Comments
 (0)