JXL: encode the remaining stages in worker threads

Besides the varblock search, XYB (bands of 64 rows), the quantization
field's masking and blur, the coefficients (transforms, chroma from
luma, quantization per row of 64x64 areas, which hold whole varblocks)
and the HF tokens (per pass group) run in the threads, started once per
encode and kept until it ends: the same files (20 encodes at efforts
0-2 with 1, 4 and 8 threads identical to before). The coefficients of
the whole frame are no longer kept, only a row's.

Ryzen 7 5800X, 8 threads, 1920x1080 / 3840x2160 photo:
effort 0 0.28 -> 0.20 s / 0.89 -> 0.42 s (now threaded from 1024x1024,
one thread per 512K pixels), effort 1 0.98 -> 0.86 s / 2.9 -> 2.0 s,
effort 2 1.8 -> 1.7 s / 5.9 -> 4.8 s.

The AVIF thread test encodes the photo with alpha instead of separately:
it still took 54 s of its 60 under coverage.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BsU2GAqbnFSNaYDBGD1Km2

Author: matmen

CHANGELOG.md

@@ -12,10 +12,12 @@
   Deno, Bun and browsers (Web Workers; bundlers like Vite pick up the worker's module). Small images (under about
   640x640 pixels at effort 0, 448x448 at 1, 320x320 at 2), browser pages without cross-origin isolation (no
   SharedArrayBuffer), runtimes without workers, and workers that fail to start encode on the calling thread as before.
-- Lossy `encodeJXL()` searches the varblocks in worker threads (`JXLEncodeOptions.threads`, the same default), rows of
-  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
-  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
-  at effort 1, 384x384 at 2, encode without threads.
+- Lossy `encodeJXL()` encodes in worker threads (`JXLEncodeOptions.threads`, the same default): XYB, the quantization
+  field, the varblock search (rows of 64x64 areas, like libjxl's), the coefficients and their tokens run in parallel,
+  writing the same files. A 1920x1080 photo takes 0.20 s instead of 0.28 s at effort 0, 0.86 s instead of 2.2 s at
+  effort 1 and 1.7 s instead of 4.9 s at effort 2, a 3840x2160 one 0.42 s instead of 0.89 s, 2.0 s instead of 8.5 s
+  and 4.8 s instead of 19.2 s (8 threads). Images under about 1024x1024 pixels at effort 0, 512x512 at 1 and 384x384
+  at 2 encode without threads.
 - JPEG XL decoding uses worker threads (`DecodeOptions.threads` and `AnimationDecodeOptions.threads`, the same
   default), from 4 megapixels on (one thread per 2 megapixels at most): the groups (with all their passes), Gaborish, the
   EPF's strips and the conversion to sRGB run in parallel, with the same pixels. A 3840x2160 image decodes in 395 ms

src/codecs/jxl/acs-enc.ts

@@ -4,7 +4,7 @@
  * area: the best 8x8 transform of each block, merged into larger DCTs where they are estimated to cost less.
  * Copyright (c) the JPEG XL Project Authors, BSD-3-Clause License (see THIRD_PARTY_NOTICES.md)
  */
-import { runChunks, shared, type ChunkTask, type Thread } from '../../util/threads.ts';
+import { runChunks, shared, sharedOf, type ChunkTask, type Thread } from '../../util/threads.ts';
 import { forwardTransform } from './transforms-enc.ts';
 import { inverseTransform } from './dct.ts';
 import { defaultDequant, MATRIX_INDEX } from './quant.ts';
@@ -290,8 +290,8 @@ export async function chooseVarblocks(planes: Float32Array[], stride: number, bw
     const types = new Int8Array(threads ? new SharedArrayBuffer(bw * bh) : new ArrayBuffer(bw * bh));
     const task: VarblockTask = {
         task: 'jxl-varblocks', stride, bw, bh, distance, speed, types,
-        planes: threads ? planes.map(shared) : planes,
-        field: threads ? { quant: shared(field.quant), mask: shared(field.mask), mask1x1: shared(field.mask1x1) } : field,
+        planes: threads ? planes.map(sharedOf) : planes,
+        field: threads ? { quant: shared(field.quant), mask: shared(field.mask), mask1x1: sharedOf(field.mask1x1) } : field,
     };
     // (a row of 64x64 areas each, independent of the others)
     await runChunks(threads, task, Math.ceil(bh / 8), varblockChunks(task));

src/codecs/jxl/aq-enc.ts

@@ -4,6 +4,7 @@
  * PerBlockModulations, Blur1x1Masking), on the whole image instead of tiles. Copyright (c) the JPEG XL Project Authors, BSD-3-Clause
  * License (see THIRD_PARTY_NOTICES.md)
  */
+import { runChunks, type ChunkTask, type Thread } from '../../util/threads.ts';
 
 const SG_MUL = 226.77216153508914, SG_MUL2 = 1 / 73.377132366608819, LOG2 = 0.693147181;
 const SG_RET_MUL = SG_MUL2 * 18.6580932135 * LOG2, SG_V_OFFSET = 7.7825991679894591;
@@ -35,36 +36,65 @@ export interface QuantField {
     mask1x1: Float32Array;
 }
 
+/** The masking of bands of 64 rows (see initialQuantField) */
+export interface MaskingTask extends ChunkTask {
+    y: Float32Array;
+    w: number;
+    h: number;
+    // the results: the masking of the pixels (and blurred), the pre-erosion
+    mask1x1: Float32Array;
+    blurred: Float32Array;
+    preErosion: Float32Array;
+}
+
 /**
- * libjxl's initial quantization field of XYB planes (w x h, whole blocks) for a Butteraugli distance (scale: the
- * field's scale)
+ * The masking of a band of 64 rows (the chunk): of each pixel the inverse of the gamma corrected Laplacian (of Y);
+ * the pre-erosion: its square (limited), through maskingSqrt, averaged over 4x4 pixels
  */
-export function initialQuantField(distance: number, planes: Float32Array[], w: number, h: number, scale: number): QuantField {
-    const [px, py, pb] = planes, bw = w / 8, bh = h / 8;
-    const MATCH_GAMMA_OFFSET = 0.019, LIMIT = 0.2;
-    // the masking of each pixel: the inverse of the gamma corrected Laplacian; the pre-erosion: its square (limited),
-    // through maskingSqrt, averaged over 4x4 pixels
-    const mask1x1 = new Float32Array(w * h), preErosion = new Float32Array((w / 4) * (h / 4)), rowSum = new Float32Array(w);
-    for (let y = 0; y < h; y++) {
-        const r = y * w, r1 = (y > 0 ? y - 1 : y) * w, r2 = (y + 1 < h ? y + 1 : y) * w;
-        for (let x = 0; x < w; x++) {
-            const x1 = x > 0 ? x - 1 : x, x2 = x + 1 < w ? x + 1 : x;
-            const base = 0.25 * (py[r2 + x] + py[r1 + x] + py[r + x1] + py[r + x2]);
-            const gammac = gammaRatio(py[r + x] + MATCH_GAMMA_OFFSET);
-            const d = gammac * (py[r + x] - base);
-            mask1x1[r + x] = 1 / (Math.log1p(Math.abs(d)) + 0.01);
-            const v = maskingSqrt(Math.min(d * d, LIMIT));
-            rowSum[x] = y % 4 ? rowSum[x] + v : v;
-        }
-        if (y % 4 === 3) {
-            const out = (y >> 2) * (w / 4);
-            for (let x = 0; x < w / 4; x++) preErosion[out + x] = (rowSum[4 * x] + rowSum[4 * x + 1] + rowSum[4 * x + 2] + rowSum[4 * x + 3]) * 0.25;
+export function maskingChunks({ y: py, w, h, mask1x1, preErosion }: MaskingTask): (chunk: number) => void {
+    const MATCH_GAMMA_OFFSET = 0.019, LIMIT = 0.2, rowSum = new Float32Array(w);
+    return band => {
+        for (let y = band * 64; y < Math.min(h, band * 64 + 64); y++) {
+            const r = y * w, r1 = (y > 0 ? y - 1 : y) * w, r2 = (y + 1 < h ? y + 1 : y) * w;
+            for (let x = 0; x < w; x++) {
+                const x1 = x > 0 ? x - 1 : x, x2 = x + 1 < w ? x + 1 : x;
+                const base = 0.25 * (py[r2 + x] + py[r1 + x] + py[r + x1] + py[r + x2]);
+                const gammac = gammaRatio(py[r + x] + MATCH_GAMMA_OFFSET);
+                const d = gammac * (py[r + x] - base);
+                mask1x1[r + x] = 1 / (Math.log1p(Math.abs(d)) + 0.01);
+                const v = maskingSqrt(Math.min(d * d, LIMIT));
+                rowSum[x] = y % 4 ? rowSum[x] + v : v;
+            }
+            if (y % 4 === 3) {
+                const out = (y >> 2) * (w / 4);
+                for (let x = 0; x < w / 4; x++) preErosion[out + x] = (rowSum[4 * x] + rowSum[4 * x + 1] + rowSum[4 * x + 2] + rowSum[4 * x + 3]) * 0.25;
+            }
         }
-    }
-    const aq = fuzzyErosion(distance, preErosion, w / 4, h / 4);
+    };
+}
+
+/** The 5x5 blur of the pixels' masking of a band of 64 rows (the chunk) */
+export function maskBlurChunks({ w, h, mask1x1, blurred }: MaskingTask): (chunk: number) => void {
+    const k = [0.25647067633737225, 0.2050056912354399, 0.1540820486684973, 0.08149576591362004, 0.05127501048123085];
+    const n = 1 / Math.max(1 + 4 * (k[0] + k[1] + k[2] + k[4] + 2 * k[3]), 1e-5);
+    const weights = [n, n * k[0], n * k[2], n * k[1], n * k[4], n * k[3]];
+    return band => symmetric5(mask1x1, w, h, weights, blurred, band * 64, Math.min(h, band * 64 + 64));
+}
+
+/**
+ * libjxl's initial quantization field of XYB planes (w x h, whole blocks) for a Butteraugli distance (scale: the
+ * field's scale); the pixels' masking in threads, if given
+ */
+export async function initialQuantField(distance: number, planes: Float32Array[], w: number, h: number, scale: number, threads: Thread[] | null): Promise<QuantField> {
+    const [px, py, pb] = planes, bw = w / 8, bh = h / 8, bands = Math.ceil(h / 64);
+    const array = (n: number): Float32Array => new Float32Array(threads ? new SharedArrayBuffer(4 * n) : new ArrayBuffer(4 * n));
+    const task: MaskingTask = { task: 'jxl-masking', y: py, w, h, mask1x1: array(w * h), blurred: array(w * h), preErosion: array((w / 4) * (h / 4)) };
+    await runChunks(threads, task, bands, maskingChunks(task));
+    const aq = fuzzyErosion(distance, task.preErosion, w / 4, h / 4);
     const mask = Float32Array.from(aq, v => 1 / (v + 0.001));
     perBlockModulations(distance, px, py, pb, w, bw, bh, scale, aq);
-    return { quant: aq, mask, mask1x1: blur1x1(mask1x1, w, h) };
+    await runChunks(threads, { ...task, task: 'jxl-mask-blur' }, bands, maskBlurChunks(task));
+    return { quant: aq, mask, mask1x1: task.blurred };
 }
 
 // the weighted 4 smallest values of each 3x3 neighbourhood, summed over 2x2 (the output is half the size)
@@ -148,16 +178,15 @@ function perBlockModulations(distance: number, px: Float32Array, py: Float32Arra
     }
 }
 
-// a symmetric 5x5 convolution with mirrored borders (libjxl's Symmetric5): the weights of the center, the edges 1 and 2
-// away, the diagonals 1 and 2 away and the knight moves
-function symmetric5(m: Float32Array, w: number, h: number, [c, r, R, d, D, L]: number[]): Float32Array {
+// rows y0 to y1 of a symmetric 5x5 convolution with mirrored borders (libjxl's Symmetric5): the weights of the center,
+// the edges 1 and 2 away, the diagonals 1 and 2 away and the knight moves
+function symmetric5(m: Float32Array, w: number, h: number, [c, r, R, d, D, L]: number[], out: Float32Array, y0: number, y1: number): void {
     const weights = [D, L, R, L, D, L, d, r, d, L, R, r, c, r, R, L, d, r, d, L, D, L, R, L, D];
     const mirror = (v: number, n: number): number => {
         while (v < 0 || v >= n) v = v < 0 ? -v - 1 : 2 * n - 1 - v;
         return v;
     };
-    const out = new Float32Array(w * h);
-    for (let y = 0; y < h; y++) {
+    for (let y = y0; y < y1; y++) {
         const rows = [-2, -1, 0, 1, 2].map(dy => mirror(y + dy, h) * w);
         for (let x = 0; x < w; x++) {
             let sum = 0;
@@ -168,12 +197,4 @@ function symmetric5(m: Float32Array, w: number, h: number, [c, r, R, d, D, L]: n
             out[y * w + x] = sum;
         }
     }
-    return out;
-}
-
-// the 5x5 blur of the pixels' masking
-function blur1x1(m: Float32Array, w: number, h: number): Float32Array {
-    const k = [0.25647067633737225, 0.2050056912354399, 0.1540820486684973, 0.08149576591362004, 0.05127501048123085];
-    const n = 1 / Math.max(1 + 4 * (k[0] + k[1] + k[2] + k[4] + 2 * k[3]), 1e-5);
-    return symmetric5(m, w, h, [n, n * k[0], n * k[2], n * k[1], n * k[4], n * k[3]]);
 }

src/codecs/jxl/frame.ts

@@ -14,7 +14,7 @@ import { ModularImage, Channel, MAConfig, channel, readTree, decodeModular } fro
 import { upsample } from './upsample.ts';
 import { gaborish, gaborishRows, epf, epfStrip, epfScratch, epfSteps, epfStrips } from './filters.ts';
 import { VarDCT, VarDCTGlobal, readVarDCTGlobal, type VarDCTShared } from './vardct.ts';
-import { runChunks, shared, type ChunkTask, type Thread } from '../../util/threads.ts';
+import { runChunks, shared, sharedOf, type ChunkTask, type Thread } from '../../util/threads.ts';
 import { Patch, Splines, readPatches, applyPatches, readSplines, renderSplines, readNoise, renderNoise } from './features.ts';
 import { ColorOutput, xybToBytes } from './color.ts';
 import { Blending, performBlending, PB_REPLACE, PB_ADD, PB_BLEND_ABOVE, PB_ADD_ABOVE, PB_MUL } from './blend.ts';
@@ -372,7 +372,6 @@ async function threaded<T>(pool: FrameThreads, pixels: number, inThreads: (threa
 }
 
 // a typed array in shared memory (itself if it is)
-const sharedOf = <T extends Float32Array>(a: T): T => a.buffer instanceof SharedArrayBuffer ? a : shared(a);
 const sharedPlanes = (n: number, w: number, h: number): Float32Array[] => Array.from({ length: n }, () => new Float32Array(new SharedArrayBuffer(4 * w * h)));
 
 /** The rows of the steps after decoding for threads (see postChunks): Gaborish, an EPF pass, XYB to bytes */