Skip to content

Commit c09009f

Browse files
matmenclaude
andcommitted
JXL: varblock search in worker threads
Rows of 64x64 areas are searched in parallel (runChunks: each thread takes the next row, the types in a SharedArrayBuffer): the same files. 1920x1080 photo, 8 threads: effort 1 2.4 -> 1.1 s, effort 2 5.0 -> 1.8 s. New option JXLEncodeOptions.threads; effort 0 and small images stay on the calling thread. runChunks/shared in src/util/threads.ts are the generic part for chunked tasks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Lfv41PqHEGbU7SLg5onpG1
1 parent 1827e49 commit c09009f

10 files changed

Lines changed: 146 additions & 18 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,10 @@
1212
Deno, Bun and browsers (Web Workers; bundlers like Vite pick up the worker's module). Small images (under about
1313
640x640 pixels at effort 0, 448x448 at 1, 320x320 at 2), browser pages without cross-origin isolation (no
1414
SharedArrayBuffer), runtimes without workers, and workers that fail to start encode on the calling thread as before.
15+
- Lossy `encodeJXL()` searches the varblocks in worker threads (`JXLEncodeOptions.threads`, the same default), rows of
16+
64x64 areas in parallel like libjxl's: the same files. A 1920x1080 photo takes 1.1 s instead of 2.4 s at effort 1 and
17+
1.8 s instead of 5.0 s at effort 2 (8 threads). Effort 0 (which has no search) and images under about 512x512 pixels
18+
at effort 1, 384x384 at 2, encode without threads.
1519

1620
### Tests
1721

‎src/codecs/jxl.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,10 @@ export function isJXL(d: Uint8Array): boolean {
2525

2626
/**
2727
* Encodes RGBA pixels as a JPEG XL, lossless or lossy (with a Butteraugli distance above 0, and an effort 0-2): the
28-
* codestream, in the container if there is Exif data (a TIFF header and IFDs)
28+
* codestream, in the container if there is Exif data (a TIFF header and IFDs); lossy with up to threads worker threads
2929
*/
30-
export function encode(rgba: Uint8Array, width: number, height: number, exif?: Uint8Array, distance = 0, effort = 1): Uint8Array {
31-
const stream = distance > 0 ? encodeVarDCT(rgba, width, height, { distance, effort }) : encodeRGBA(rgba, width, height);
30+
export async function encode(rgba: Uint8Array, width: number, height: number, exif?: Uint8Array, distance = 0, effort = 1, threads = 1): Promise<Uint8Array> {
31+
const stream = distance > 0 ? await encodeVarDCT(rgba, width, height, { distance, effort, threads }) : encodeRGBA(rgba, width, height);
3232
if (!exif?.length) return stream;
3333
// (ftyp: brand jxl, minor version 0, compatible jxl; Exif: the offset of the TIFF header, 0, then the TIFF data)
3434
const jxl = [0x6a, 0x78, 0x6c, 0x20];

‎src/codecs/jxl/acs-enc.ts‎

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
* area: the best 8x8 transform of each block, merged into larger DCTs where they are estimated to cost less.
55
* Copyright (c) the JPEG XL Project Authors, BSD-3-Clause License (see THIRD_PARTY_NOTICES.md)
66
*/
7+
import { runChunks, shared, type ChunkTask, type Thread } from '../../util/threads.ts';
78
import { forwardTransform } from './transforms-enc.ts';
89
import { inverseTransform } from './dct.ts';
910
import { defaultDequant, MATRIX_INDEX } from './quant.ts';
@@ -117,8 +118,8 @@ function best8x8(x: number, y: number, speed: number, cfg: Config): [number, num
117118
class Strategy {
118119
readonly types: Int8Array;
119120

120-
constructor(readonly bw: number, readonly bh: number) {
121-
this.types = new Int8Array(bw * bh);
121+
constructor(readonly bw: number, readonly bh: number, types?: Int8Array) {
122+
this.types = types ?? new Int8Array(bw * bh);
122123
}
123124

124125
set(bx: number, by: number, type: number): void {
@@ -285,15 +286,42 @@ function processRect(bx: number, by: number, w: number, h: number, speed: number
285286
* The varblock types of an image (XYB planes of whole blocks, rows stride apart): libjxl's search at a speed tier (5:
286287
* its effort 5, 4: also the 4x8 and AFV transforms)
287288
*/
288-
export function chooseVarblocks(planes: Float32Array[], stride: number, bw: number, bh: number, field: QuantField, distance: number, speed: number): Int8Array {
289+
export async function chooseVarblocks(planes: Float32Array[], stride: number, bw: number, bh: number, field: QuantField, distance: number, speed: number, threads: Thread[] | null): Promise<Int8Array> {
290+
const types = new Int8Array(threads ? new SharedArrayBuffer(bw * bh) : new ArrayBuffer(bw * bh));
291+
const task: VarblockTask = {
292+
task: 'jxl-varblocks', stride, bw, bh, distance, speed, types,
293+
planes: threads ? planes.map(shared) : planes,
294+
field: threads ? { quant: shared(field.quant), mask: shared(field.mask), mask1x1: shared(field.mask1x1) } : field,
295+
};
296+
// (a row of 64x64 areas each, independent of the others)
297+
await runChunks(threads, task, Math.ceil(bh / 8), varblockChunks(task));
298+
return types;
299+
}
300+
301+
/** The varblock search of rows of 64x64 areas (see chooseVarblocks) */
302+
export interface VarblockTask extends ChunkTask {
303+
planes: Float32Array[];
304+
stride: number;
305+
bw: number;
306+
bh: number;
307+
field: QuantField;
308+
distance: number;
309+
speed: number;
310+
types: Int8Array;
311+
}
312+
313+
/** The search of a row of 64x64 areas (the chunk), its types set */
314+
export function varblockChunks({ planes, stride, bw, bh, field, distance, speed, types }: VarblockTask): (chunk: number) => void {
289315
const ratio = (distance + 0.13731742964354549) / (1 + 0.13731742964354549);
290316
const cfg: Config = {
291317
planes, stride, bw, field, distance,
292318
infoLossMul: 1.2 * ratio ** 0.33677806662454718,
293319
zerosMul: 9.3089059022677905 * ratio ** 0.50990926717963703,
294320
costDelta: 10.833273317067883 * ratio ** 0.36702940662370243
295321
};
296-
const s = new Strategy(bw, bh);
297-
for (let by = 0; by < bh; by += 8) for (let bx = 0; bx < bw; bx += 8) processRect(bx, by, Math.min(8, bw - bx), Math.min(8, bh - by), speed, cfg, s);
298-
return s.types;
322+
const s = new Strategy(bw, bh, types);
323+
return row => {
324+
const by = row * 8;
325+
for (let bx = 0; bx < bw; bx += 8) processRect(bx, by, Math.min(8, bw - bx), Math.min(8, bh - by), speed, cfg, s);
326+
};
299327
}

‎src/codecs/jxl/vardct-enc.ts‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
* contexts, chroma from luma per tile (libjxl's fast search), the edge preserving filter. Copyright (c) the JPEG XL
66
* Project Authors, BSD-3-Clause License (see THIRD_PARTY_NOTICES.md)
77
*/
8+
import { startThreads } from '../../util/threads.ts';
89
import { findPatches } from './patches-enc.ts';
910
import { BitWriter } from './bit-writer.ts';
1011
import { imageHeader, u32 } from './headers-enc.ts';
@@ -22,8 +23,14 @@ export interface VarDCTOptions {
2223
distance: number;
2324
/** 0: 8x8 DCTs only (like libjxl's effort 3), 1: libjxl's varblock search (its effort 5), 2: also the 4x8 and AFV transforms */
2425
effort: number;
26+
/** at most this many worker threads (1: none) */
27+
threads: number;
2528
}
2629

30+
// (images below this many pixels encode faster without threads, which take about 0.1 s to start: by effort; effort 0
31+
// has no varblock search, what the threads run)
32+
const MIN_THREADED_PIXELS = [Infinity, 512 * 512, 384 * 384];
33+
2734
// libjxl's opsin absorbance matrix and bias (lib/jxl/cms/opsin_params.h)
2835
const OPSIN = [
2936
0.30, 0.622, 0.078,
@@ -167,7 +174,7 @@ function chromaFromLuma(raw: Float32Array[], types: Int8Array, muls: Int32Array,
167174
}
168175

169176
/** Encodes RGBA pixels as a lossy JPEG XL codestream, alpha (if the image has transparency) losslessly */
170-
export function encodeVarDCT(rgba: Uint8Array, width: number, height: number, { distance, effort }: VarDCTOptions): Uint8Array {
177+
export async function encodeVarDCT(rgba: Uint8Array, width: number, height: number, { distance, effort, threads: maxThreads }: VarDCTOptions): Promise<Uint8Array> {
171178
if (width < 1 || height < 1) throw new RangeError('JPEG XL: dimensions must be at least 1');
172179
const bw = Math.ceil(width / 8), bh = Math.ceil(height / 8), pw = bw * 8;
173180
const planes = toXYB(rgba, width, height, pw, bh * 8);
@@ -191,7 +198,24 @@ export function encodeVarDCT(rgba: Uint8Array, width: number, height: number, {
191198
// HF multipliers
192199
// patches at effort 2: text-like features of screenshots stored once (subtracted from the planes here)
193200
const patches = effort >= 2 ? findPatches(planes, pw, width, height, distance, alpha ? 1 : 0) : null;
194-
const types = effort ? chooseVarblocks(planes, pw, bw, bh, initialQuantField(distance * 0.62, planes, pw, bh * 8, 0.725 / (distance * 0.62)), distance, effort === 1 ? 5 : 4) : new Int8Array(bw * bh);
201+
// (the search in threads, if they start: the same types)
202+
let types = new Int8Array(bw * bh);
203+
if (effort) {
204+
const field = initialQuantField(distance * 0.62, planes, pw, bh * 8, 0.725 / (distance * 0.62));
205+
// (at most one per 256x256 pixels: more take longer to start than they save)
206+
const n = Math.min(maxThreads, Math.ceil(bh / 8), Math.floor(width * height / 65536));
207+
let threads = n > 1 && width * height >= MIN_THREADED_PIXELS[effort] ? startThreads(n) : null;
208+
try {
209+
types = await chooseVarblocks(planes, pw, bw, bh, field, distance, effort === 1 ? 5 : 4, threads);
210+
} catch (e) {
211+
if (!threads) throw e;
212+
for (const t of threads) t.terminate();
213+
threads = null;
214+
types = await chooseVarblocks(planes, pw, bw, bh, field, distance, effort === 1 ? 5 : 4, null);
215+
} finally {
216+
if (threads) for (const t of threads) t.terminate();
217+
}
218+
}
195219
const muls = new Int32Array(bw * bh).fill(mul);
196220

197221
// the varblocks' coefficients (in the planes' layout) and the LF of their blocks

‎src/image.ts‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,12 @@ import * as hdr from './codecs/hdr.ts';
1212
import * as avif from './codecs/avif.ts';
1313
import * as avifEnc from './codecs/avif-enc.ts';
1414
import { availableThreads } from './util/threads.ts';
15+
16+
// (worker threads: the processors available, at most 8)
17+
const defaultThreads = (): number => Math.max(1, Math.min(availableThreads(), 8));
18+
const checkThreads = (threads: number): void => {
19+
if (!Number.isInteger(threads) || threads < 1) throw new RangeError(`threads must be a positive integer, not ${threads}`);
20+
};
1521
import * as jxl from './codecs/jxl.ts';
1622
import type { ToneMapping } from './codecs/hdr.ts';
1723
import * as quantize from './codecs/quantize.ts';
@@ -221,6 +227,13 @@ export interface JXLEncodeOptions {
221227
* text-like features that repeat (the glyphs of screenshots) once, as patches like `cjxl -e 7`. Default: `1`
222228
*/
223229
effort?: number;
230+
/**
231+
* At most this many worker threads search the transforms of lossy compression in parallel (the same file, faster),
232+
* `1` for none. Effort 0 (no search), small images (under about 512x512 pixels at effort 1, 384x384 at 2), and
233+
* runtimes whose workers can't share memory (browsers without cross-origin isolation) encode without them.
234+
* Default: the processors available, at most 8
235+
*/
236+
threads?: number;
224237
/** Metadata to include, see {@link Metadata} */
225238
metadata?: Metadata;
226239
}
@@ -2176,8 +2189,8 @@ export class Image {
21762189
* Encodes the image as AVIF (AV1, with an alpha image if the image has transparency)
21772190
* @param options See {@link AVIFEncodeOptions}
21782191
*/
2179-
async encodeAVIF({ quality = 60, alphaQuality = 100, subsampling = '4:4:4', effort = 1, threads = Math.max(1, Math.min(availableThreads(), 8)), metadata = {} }: AVIFEncodeOptions = {}): Promise<Uint8Array> {
2180-
if (!Number.isInteger(threads) || threads < 1) throw new RangeError(`threads must be a positive integer, not ${threads}`);
2192+
async encodeAVIF({ quality = 60, alphaQuality = 100, subsampling = '4:4:4', effort = 1, threads = defaultThreads(), metadata = {} }: AVIFEncodeOptions = {}): Promise<Uint8Array> {
2193+
checkThreads(threads);
21812194
const exif = exifData(resolveMetadata(metadata));
21822195
return avifEnc.encode(this.__bytes__(), this.width, this.height, { quality: finite(quality, 'quality'), alphaQuality: finite(alphaQuality, 'alphaQuality'), subsampling, effort, exif, threads });
21832196
}
@@ -2190,16 +2203,17 @@ export class Image {
21902203
* transparency) is lossless either way. The codestream is in the container, with the metadata as Exif
21912204
* @param options See {@link JXLEncodeOptions}
21922205
*/
2193-
async encodeJXL({ quality, effort = 1, metadata = {} }: JXLEncodeOptions = {}): Promise<Uint8Array> {
2206+
async encodeJXL({ quality, effort = 1, threads = defaultThreads(), metadata = {} }: JXLEncodeOptions = {}): Promise<Uint8Array> {
21942207
if (!Number.isInteger(effort) || effort < 0 || effort > 2) throw new RangeError(`effort must be an integer 0-2, not ${effort}`);
2208+
checkThreads(threads);
21952209
let distance = 0;
21962210
if (quality !== undefined) {
21972211
const q = finite(quality, 'quality');
21982212
if (q < 0 || q > 100) throw new RangeError(`quality must be 0-100, not ${q}`);
21992213
// (cjxl's distance of a quality)
22002214
distance = q >= 100 ? 0 : q >= 30 ? 0.1 + (100 - q) * 0.09 : 53 / 3000 * q * q - 23 / 20 * q + 25;
22012215
}
2202-
return jxl.encode(this.__bytes__(), this.width, this.height, exifData(resolveMetadata(metadata)), distance, effort);
2216+
return jxl.encode(this.__bytes__(), this.width, this.height, exifData(resolveMetadata(metadata)), distance, effort, threads);
22032217
}
22042218

22052219
/**

‎src/util/threads.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,3 +129,37 @@ export function startThreads(n: number): Thread[] | null {
129129
return null;
130130
}
131131
}
132+
133+
/** A task split into chunks, for runChunks: the worker task's name and its data (typed arrays in shared memory) */
134+
export interface ChunkTask {
135+
task: string;
136+
/** the next chunk to take (an Int32Array of 1 in shared memory, set by runChunks) */
137+
counter?: Int32Array;
138+
chunks?: number;
139+
}
140+
141+
/**
142+
* Runs chunks 0 to count - 1 of a task, each threads' next one (worker.ts runs a chunk task's chunks with the same
143+
* function as local, see chunkTasks there), or all of them here without threads: their results
144+
*/
145+
export async function runChunks<R>(threads: Thread[] | null, task: ChunkTask, count: number, local: (chunk: number) => R): Promise<R[]> {
146+
const results = new Array<R>(count);
147+
if (!threads || count < 2) {
148+
for (let i = 0; i < count; i++) results[i] = local(i);
149+
return results;
150+
}
151+
const counter = new Int32Array(new SharedArrayBuffer(4));
152+
await Promise.all(threads.slice(0, count).map(t => t.run({ ...task, counter, chunks: count }, m => {
153+
const { chunk, result } = m as { chunk: number, result: R };
154+
results[chunk] = result;
155+
})));
156+
return results;
157+
}
158+
159+
/** A copy of a typed array in shared memory (for threads) */
160+
export function shared<T extends Float32Array | Int32Array | Int8Array | Uint8Array | Uint16Array>(array: T): T {
161+
const Type = array.constructor as new (buffer: SharedArrayBuffer) => T;
162+
const copy = new Type(new SharedArrayBuffer(array.byteLength));
163+
copy.set(array);
164+
return copy;
165+
}

‎src/util/worker.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,12 @@
11
// the module worker threads run (see threads.ts): tasks posted to it, their results posted back
22
import { searchRows, filterTrials, type SearchTask, type FilterTask } from '../codecs/av1/encode.ts';
3+
import { varblockChunks } from '../codecs/jxl/acs-enc.ts';
4+
import type { ChunkTask } from './threads.ts';
5+
6+
// the chunk tasks (see runChunks): a task's function of its chunks
7+
const chunkTasks: Record<string, (task: never) => (chunk: number) => unknown> = {
8+
'jxl-varblocks': varblockChunks,
9+
};
310

411
interface Port {
512
post(message: unknown): void;
@@ -19,6 +26,10 @@ function handle(port: Port, task: { task: string }): void {
1926
let result: unknown;
2027
if (task.task === 'av1-search') result = searchRows(task as SearchTask, message);
2128
else if (task.task === 'av1-filter') result = filterTrials(task as FilterTask);
29+
else if (task.task in chunkTasks) {
30+
const { counter, chunks } = task as ChunkTask, run = chunkTasks[task.task](task as never);
31+
for (let chunk; (chunk = Atomics.add(counter as Int32Array, 0, 1)) < (chunks as number);) message({ chunk, result: run(chunk) });
32+
}
2233
else throw new Error(`unknown task ${task.task}`);
2334
port.post({ result });
2435
} catch (e) {

‎tests/browser/page.mjs‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
1-
// the browser test's page script: AVIF with and without threads (the same file), the codecs' smoke test
1+
// the browser test's page script: AVIF and JXL with and without threads (the same files), the codecs' smoke test
22
import { Image } from 'imagescript';
33
import { smoke } from '../edge/smoke.mjs';
44

55
const digest = async data => [...new Uint8Array(await crypto.subtle.digest('SHA-256', data))].map(b => b.toString(16).padStart(2, '0')).join('');
66
(async () => {
77
const results = { isolated: self.crossOriginIsolated };
88
try {
9-
const image = new Image(640, 480);
10-
for (let y = 0; y < 480; y++) for (let x = 0; x < 640; x++) image.setPixelAt(x, y, Image.rgbaToColor(x * 3 & 255, y * 5 & 255, (x ^ y) & 255, 255));
9+
const image = new Image(640, 512);
10+
for (let y = 0; y < 512; y++) for (let x = 0; x < 640; x++) image.setPixelAt(x, y, Image.rgbaToColor(x * 3 & 255, y * 5 & 255, (x ^ y) & 255, 255));
1111
results.single = await digest(await image.encodeAVIF({ effort: 1, threads: 1, metadata: { creationTime: 0 } }));
1212
results.threaded = await digest(await image.encodeAVIF({ effort: 1, threads: 4, metadata: { creationTime: 0 } }));
13+
results.jxlSingle = await digest(await image.encodeJXL({ quality: 85, threads: 1, metadata: { creationTime: 0 } }));
14+
results.jxlThreaded = await digest(await image.encodeJXL({ quality: 85, threads: 4, metadata: { creationTime: 0 } }));
1315
results.smoke = await smoke();
1416
} catch (e) {
1517
results.error = String(e && e.stack || e);

‎tests/browser/run.mjs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ try {
5050
if (r.error) problems.push(r.error);
5151
if (r.isolated !== isolated) problems.push(`crossOriginIsolated ${r.isolated}`);
5252
if (r.single !== r.threaded) problems.push('threads wrote a different file');
53+
if (r.jxlSingle !== r.jxlThreaded) problems.push('threads wrote a different JPEG XL');
5354
if (isolated ? workers < 2 : workers) problems.push(`${workers} workers started`);
5455
for (const [step, result] of Object.entries(r.smoke ?? {})) if (String(result).startsWith('FAIL')) problems.push(`${step}: ${result}`);
5556
console.log(`${isolated ? 'cross-origin isolated' : 'not isolated'}: ${workers} workers, ${problems.length ? problems.join('\n ') : 'ok'}`);

‎tests/jxl-encode.js‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,16 @@ const codestream = jxl => {
113113
else if (jxl.length > effort0 * 1.02) panic(`effort ${effort}: ${jxl.length} B, effort 0 ${effort0} B`);
114114
}
115115
for (const effort of [-1, 3, 1.5]) if (!await shot.encodeJXL({quality: 90, effort}).then(() => false, e => e instanceof RangeError)) panic(`effort ${effort} should throw`);
116+
// threads (the varblock search of rows of 64x64 areas in parallel): the same files as without
117+
{
118+
const big = photo.clone().resize(640, Image.RESIZE_AUTO);
119+
for (const effort of [1, 2]) {
120+
const one = await big.encodeJXL({quality: 85, effort, threads: 1, metadata: {creationTime: 0}});
121+
const four = await big.encodeJXL({quality: 85, effort, threads: 4, metadata: {creationTime: 0}});
122+
if (Buffer.compare(Buffer.from(one), Buffer.from(four))) panic(`threads, effort ${effort}: not the same file`);
123+
}
124+
for (const threads of [0, 2.5]) if (!await big.encodeJXL({quality: 85, threads}).then(() => false, e => e instanceof RangeError)) panic(`threads ${threads} should throw`);
125+
}
116126
// patches (effort 2): repeated glyphs of text on a flat background are found, stored in a reference frame and
117127
// added back
118128
{

0 commit comments

Comments
 (0)