JPEG XL: faster lossy encoding

Author: matmen

CHANGELOG.md

@@ -1,5 +1,17 @@
 # Changelog
 
+## 3.13.3
+
+### Performance
+
+- Lossy `encodeJXL()` is faster: a 1920x1080 photo at quality 90 takes 1.6 s instead of 2.3 s at effort 1 and 2.4 s
+  instead of 4.8 s at effort 2 on one thread; with 8 threads 0.72 s instead of 0.81 s and 0.95 s instead of 1.46 s
+  (3840x2160: 1.6 s instead of 2.0 s and 2.5 s instead of 4.1 s). The 4x8 and AFV transforms are computed directly
+  instead of by a 64x64 matrix, the varblock search rounds without `Math.round()` (slow in V8), patch search skips
+  large components after a quick labeling, and the search is split into smaller pieces for the threads. The files stay
+  the same, except for rare coefficients the new transforms round differently (2 of 108 test encodes, both the same
+  size).
+
 ## 3.13.2
 
 ### Fixes

src/codecs/jxl/acs-enc.ts

@@ -6,9 +6,10 @@
  */
 import { runChunks, shared, sharedOf, type ChunkTask, type Thread } from '../../util/threads.ts';
 import { forwardTransform } from './transforms-enc.ts';
-import { inverseTransform } from './dct.ts';
+import { idct8, inverseTransform } from './dct.ts';
+import { idct16, idct32, idct64 } from './idct.ts';
 import { defaultDequant, MATRIX_INDEX } from './quant.ts';
-import { SIZES, transposed } from './vardct.ts';
+import { SIZES, single, transposed } from './vardct.ts';
 import type { QuantField } from './aq-enc.ts';
 
 const DCT8 = 0, IDENTITY = 1, DCT2X2 = 2, DCT4X4 = 3, DCT16X16 = 4, DCT32X32 = 5, DCT16X8 = 6, DCT8X16 = 7;
@@ -41,6 +42,19 @@ function inverseDequant(type: number): Float32Array[] {
     return inv;
 }
 
+// the inverse DCT of a w x h block (8 to 64 points: like dct2d, rows then columns, by direct calls)
+function idct2(buf: Float32Array, w: number, h: number): void {
+    for (let y = 0; y < h; y++) idct(buf, y * w, 1, w);
+    for (let x = 0; x < w; x++) idct(buf, x, w, h);
+}
+
+function idct(buf: Float32Array, o: number, step: number, n: number): void {
+    if (n === 8) idct8(buf, o, step);
+    else if (n === 16) idct16(buf, o, step);
+    else if (n === 32) idct32(buf, o, step);
+    else idct64(buf, o, step);
+}
+
 // the estimated cost of a varblock of a type at pixels x, y: its bits (by the square roots of the quantized
 // coefficients and the non-zero counts) and the information lost (an 8-norm of the masked errors in pixels)
 function estimateEntropy(type: number, entropyMul: number, x: number, y: number, cfg: Config): number {
@@ -65,14 +79,17 @@ function estimateEntropy(type: number, entropyMul: number, x: number, y: number,
     for (let c = 0; c < 3; c++) {
         const src = scratch[c], ys = scratch[1], m = matrices[c], im = inverses[c], cmap = c === 2 ? 1 : 0;
         let sqrtSum = 0, nonZeros = 0;
+        // (rounded by Math.floor(v + 0.5): Math.round but for values within an ulp below halves, and faster in V8)
         for (let i = 0; i < size; i++) {
-            const v = (src[i] - ys[i] * cmap) * im[i] * quant, r = Math.round(v);
+            const v = (src[i] - ys[i] * cmap) * im[i] * quant, r = Math.floor(v + 0.5);
             errors[i] = m[i] * (v - r);
-            const q = Math.abs(r);
-            sqrtSum += Math.sqrt(q);
-            if (q) nonZeros++;
+            if (r) {
+                sqrtSum += Math.sqrt(Math.abs(r));
+                nonZeros++;
+            }
         }
-        inverseTransform({ buf: errors, offset: 0, stride: w, w, h }, type);
+        if (type === 0 || !single(type)) idct2(errors, w, h);
+        else inverseTransform({ buf: errors, offset: 0, stride: w, w, h }, type);
         let lossc = 0;
         for (let py = 0; py < h; py++) {
             const mo = (y + py) * stride + x;
@@ -293,8 +310,8 @@ export async function chooseVarblocks(planes: Float32Array[], stride: number, bw
         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));
+    // (up to 8 64x64 areas of a row each, independent of the others: small enough to keep all threads busy)
+    await runChunks(threads, task, Math.ceil(bh / 8) * Math.ceil(bw / 64), varblockChunks(task));
     return types;
 }
 
@@ -310,7 +327,7 @@ export interface VarblockTask extends ChunkTask {
     types: Int8Array;
 }
 
-/** The search of a row of 64x64 areas (the chunk), its types set */
+/** The search of up to 8 64x64 areas of a row (the chunk: of each row's segments of 8), its types set */
 export function varblockChunks({ planes, stride, bw, bh, field, distance, speed, types }: VarblockTask): (chunk: number) => void {
     const ratio = (distance + 0.13731742964354549) / (1 + 0.13731742964354549);
     const cfg: Config = {
@@ -320,8 +337,9 @@ export function varblockChunks({ planes, stride, bw, bh, field, distance, speed,
         costDelta: 10.833273317067883 * ratio ** 0.36702940662370243
     };
     const s = new Strategy(bw, bh, types);
-    return row => {
-        const by = row * 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);
+    const segments = Math.ceil(bw / 64);
+    return chunk => {
+        const by = Math.floor(chunk / segments) * 8, x0 = (chunk % segments) * 64;
+        for (let bx = x0; bx < Math.min(bw, x0 + 64); bx += 8) processRect(bx, by, Math.min(8, bw - bx), Math.min(8, bh - by), speed, cfg, s);
     };
 }

src/codecs/jxl/aq-enc.ts

@@ -138,62 +138,67 @@ function perBlockModulations(distance: number, px: Float32Array, py: Float32Arra
     if (distance >= 2) dampen = Math.max(0, 1 - (distance - 2) / (14 - 2));
     const mul = scale * dampen, add = (1 - dampen) * baseLevel;
     for (let by = 0; by < bh; by++) {
-        for (let bx = 0; bx < bw; bx++) {
-            const o = by * 8 * w + bx * 8;
-            let v = computeMask(aq[by * bw + bx]);
-            // high frequencies: the (limited) differences of Y to the right and below
-            let hf = 0;
-            for (let dy = 0; dy < 8; dy++) {
-                const r = o + dy * w, next = dy === 7 ? r : r + w;
-                for (let dx = 0; dx < 8; dx++) {
-                    const p = py[r + dx];
-                    if (dx < 7) hf += Math.min(0.0206, Math.abs(p - py[r + dx + 1]));
-                    hf += Math.min(0.0206, Math.abs(p - py[next + dx]));
-                }
-            }
-            v += hf * -0.38 + 0.42;
-            // gamma
-            let ratio = 0;
-            for (let dy = 0; dy < 8; dy++) {
-                for (let dx = 0; dx < 8; dx++) {
-                    const iny = py[o + dy * w + dx] + 0.16, inx = px[o + dy * w + dx];
-                    ratio += gammaRatioInverse(iny - inx) + gammaRatioInverse(iny + inx);
-                }
-            }
-            v += 0.1005613337192697 * Math.log2(ratio * (0.5 / 64));
-            // blue (but not all blue)
-            const LIMIT = 0.027121074570634722;
-            let blue = 0;
-            for (let dy = 0; dy < 8; dy++) {
-                for (let dx = 0; dx < 8; dx++) {
-                    const i = o + dy * w + dx, eff = py[i] + 0.084381641171960495 + Math.abs(px[i]);
-                    if (pb[i] > eff) blue += Math.min(pb[i] - eff, LIMIT);
-                }
-            }
-            if (blue >= 32 * LIMIT) blue = 64 * LIMIT - blue;
-            if (blue >= 15.398788439047934 * LIMIT) blue = 15.398788439047934 * LIMIT;
-            v += blue * 0.14207000358439159;
-            aq[by * bw + bx] = Math.exp(v) * mul + add;
+        for (let bx = 0; bx < bw; bx++) aq[by * bw + bx] = Math.exp(blockModulation(px, py, pb, by * 8 * w + bx * 8, w, aq[by * bw + bx])) * mul + add;
+    }
+}
+
+// the modulation of a block (at o; e: its erosion) before the exponentiation (a function of its own: V8 kept leaving
+// and re-entering perBlockModulations' inner loops compiled on stack)
+function blockModulation(px: Float32Array, py: Float32Array, pb: Float32Array, o: number, w: number, e: number): number {
+    let v = computeMask(e);
+    // high frequencies: the (limited) differences of Y to the right and below
+    let hf = 0;
+    for (let dy = 0; dy < 8; dy++) {
+        const r = o + dy * w, next = dy === 7 ? r : r + w;
+        for (let dx = 0; dx < 8; dx++) {
+            const p = py[r + dx];
+            if (dx < 7) hf += Math.min(0.0206, Math.abs(p - py[r + dx + 1]));
+            hf += Math.min(0.0206, Math.abs(p - py[next + dx]));
         }
     }
+    v += hf * -0.38 + 0.42;
+    // gamma
+    let ratio = 0;
+    for (let dy = 0; dy < 8; dy++) {
+        for (let dx = 0; dx < 8; dx++) {
+            const iny = py[o + dy * w + dx] + 0.16, inx = px[o + dy * w + dx];
+            ratio += gammaRatioInverse(iny - inx) + gammaRatioInverse(iny + inx);
+        }
+    }
+    v += 0.1005613337192697 * Math.log2(ratio * (0.5 / 64));
+    // blue (but not all blue)
+    const LIMIT = 0.027121074570634722;
+    let blue = 0;
+    for (let dy = 0; dy < 8; dy++) {
+        for (let dx = 0; dx < 8; dx++) {
+            const i = o + dy * w + dx, eff = py[i] + 0.084381641171960495 + Math.abs(px[i]);
+            if (pb[i] > eff) blue += Math.min(pb[i] - eff, LIMIT);
+        }
+    }
+    if (blue >= 32 * LIMIT) blue = 64 * LIMIT - blue;
+    if (blue >= 15.398788439047934 * LIMIT) blue = 15.398788439047934 * LIMIT;
+    return v + blue * 0.14207000358439159;
 }
 
 // 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;
     };
+    // (the weights by column: of rows -2 to 2; summed by column, then row)
     for (let y = y0; y < y1; y++) {
-        const rows = [-2, -1, 0, 1, 2].map(dy => mirror(y + dy, h) * w);
+        const r0 = mirror(y - 2, h) * w, r1 = mirror(y - 1, h) * w, r2 = y * w, r3 = mirror(y + 1, h) * w, r4 = mirror(y + 2, h) * w;
         for (let x = 0; x < w; x++) {
+            const inside = x >= 2 && x + 2 < w;
+            const c0 = inside ? x - 2 : mirror(x - 2, w), c1 = inside ? x - 1 : mirror(x - 1, w), c3 = inside ? x + 1 : mirror(x + 1, w), c4 = inside ? x + 2 : mirror(x + 2, w);
             let sum = 0;
-            for (let dx = -2; dx <= 2; dx++) {
-                const col = x + dx >= 0 && x + dx < w ? x + dx : mirror(x + dx, w);
-                for (let k = 0; k < 5; k++) sum += weights[k * 5 + dx + 2] * m[rows[k] + col];
-            }
+            sum += D * m[r0 + c0]; sum += L * m[r1 + c0]; sum += R * m[r2 + c0]; sum += L * m[r3 + c0]; sum += D * m[r4 + c0];
+            sum += L * m[r0 + c1]; sum += d * m[r1 + c1]; sum += r * m[r2 + c1]; sum += d * m[r3 + c1]; sum += L * m[r4 + c1];
+            sum += R * m[r0 + x]; sum += r * m[r1 + x]; sum += c * m[r2 + x]; sum += r * m[r3 + x]; sum += R * m[r4 + x];
+            sum += L * m[r0 + c3]; sum += d * m[r1 + c3]; sum += r * m[r2 + c3]; sum += d * m[r3 + c3]; sum += L * m[r4 + c3];
+            sum += D * m[r0 + c4]; sum += L * m[r1 + c4]; sum += R * m[r2 + c4]; sum += L * m[r3 + c4]; sum += D * m[r4 + c4];
             out[y * w + x] = sum;
         }
     }

src/codecs/jxl/dct.ts

@@ -102,7 +102,7 @@ const work = { row: new Float64Array(256), scratch: new Float64Array(256) };
 const SEC8 = secHalf(8);
 
 // the inverse 8-point DCT of buf[o + k * step] in place (like dct's recursion, unrolled)
-function idct8(buf: Float32Array, o: number, step: number): void {
+export function idct8(buf: Float32Array, o: number, step: number): void {
     const v0 = buf[o], v1 = buf[o + step], v2 = buf[o + 2 * step], v3 = buf[o + 3 * step];
     const v4 = buf[o + 4 * step], v5 = buf[o + 5 * step], v6 = buf[o + 6 * step], v7 = buf[o + 7 * step];
     // the even half: v0, v2, v4, v6
@@ -229,7 +229,28 @@ const set = (b: Block, x: number, y: number, v: number): void => {
     b.buf[b.offset + y * b.stride + x] = v;
 };
 const tmp = { a: new Float32Array(64), b: new Float32Array(64) };
-const blockOf = (buf: Float32Array, offset: number, w: number, h: number, stride = w): Block => ({ buf, offset, stride, w, h });
+
+// the inverse 4-point DCT of buf[o + k * step] in place (dct4's)
+function idct4(buf: Float32Array, o: number, step: number): void {
+    const i0 = buf[o], i1 = buf[o + step], i2 = buf[o + 2 * step], i3 = buf[o + 3 * step];
+    const t0 = i1 * Math.SQRT2, t1 = i1 + i3, out0 = (t0 + t1) * SEC0, out1 = (t0 - t1) * SEC1;
+    const sum02 = i0 + i2, sub02 = i0 - i2;
+    buf[o] = sum02 + out0;
+    buf[o + step] = sub02 + out1;
+    buf[o + 2 * step] = sub02 - out1;
+    buf[o + 3 * step] = sum02 - out0;
+}
+
+// the inverse 2D DCTs of 8x4 and 4x4 blocks at buf[o] (rows, then columns: dct2d's)
+function idct8x4(buf: Float32Array, o: number): void {
+    for (let y = 0; y < 4; y++) idct8(buf, o + y * 8, 1);
+    for (let x = 0; x < 8; x++) idct4(buf, o + x, 8);
+}
+
+function idct4x4(buf: Float32Array, o: number): void {
+    for (let y = 0; y < 4; y++) idct4(buf, o + y * 4, 1);
+    for (let x = 0; x < 4; x++) idct4(buf, o + x, 4);
+}
 
 function idct2InPlace(b: Block, size: number): void {
     const half = size / 2, s = tmp.a;
@@ -259,7 +280,7 @@ function dct4x4(b: Block): void {
             const o = (y * 2 + x) * 16;
             // (transposed)
             for (let iy = 0; iy < 4; iy++) for (let ix = 0; ix < 4; ix++) s[o + ix * 4 + iy] = get(b, x + ix * 2, y + iy * 2);
-            dct2d(blockOf(s, o, 4, 4), true);
+            idct4x4(s, o);
         }
     }
     for (let y = 0; y < 2; y++) for (let x = 0; x < 2; x++) for (let iy = 0; iy < 4; iy++) for (let ix = 0; ix < 4; ix++) set(b, x * 4 + ix, y * 4 + iy, s[(y * 2 + x) * 16 + iy * 4 + ix]);
@@ -290,27 +311,28 @@ function dct4x8(b: Block, transposed: boolean): void {
     const s = tmp.b;
     for (let i = 0; i < 2; i++) {
         for (let iy = 0; iy < 4; iy++) for (let ix = 0; ix < 8; ix++) s[i * 32 + iy * 8 + ix] = get(b, ix, iy * 2 + i);
-        dct2d(blockOf(s, i * 32, 8, 4), true);
+        idct8x4(s, i * 32);
     }
     for (let y = 0; y < 8; y++) for (let x = 0; x < 8; x++) {
         if (transposed) set(b, y, x, s[y * 8 + x]);
         else set(b, x, y, s[y * 8 + x]);
     }
 }
 
+const afvScratch = { coeffs: new Float64Array(16), samples: new Float64Array(16), s4: new Float32Array(16), s8: new Float32Array(32) };
+
 function afv(b: Block, n: number): void {
     const flipX = n % 2, flipY = n >> 1;
-    const coeffs = new Float64Array(16), samples = new Float64Array(16);
+    const coeffs = afvScratch.coeffs, samples = afvScratch.samples.fill(0), s4 = afvScratch.s4, s8 = afvScratch.s8;
     coeffs[0] = (get(b, 0, 0) + get(b, 1, 0) + get(b, 0, 1)) * 4;
     for (let i = 1; i < 16; i++) coeffs[i] = get(b, 2 * (i % 4), 2 * (i >> 2));
     for (let i = 0; i < 16; i++) for (let k = 0; k < 16; k++) samples[k] += coeffs[i] * AFV_BASIS[i * 16 + k];
-    const s4 = new Float32Array(16), s8 = new Float32Array(32);
     s4[0] = get(b, 0, 0) - get(b, 1, 0) + get(b, 0, 1);
     for (let iy = 0; iy < 4; iy++) for (let ix = 0; ix < 4; ix++) if (ix | iy) s4[ix * 4 + iy] = get(b, 2 * ix + 1, 2 * iy);
-    dct2d(blockOf(s4, 0, 4, 4), true);
+    idct4x4(s4, 0);
     s8[0] = get(b, 0, 0) - get(b, 0, 1);
     for (let iy = 0; iy < 4; iy++) for (let ix = 0; ix < 8; ix++) if (ix | iy) s8[iy * 8 + ix] = get(b, ix, 2 * iy + 1);
-    dct2d(blockOf(s8, 0, 8, 4), true);
+    idct8x4(s8, 0);
     for (let iy = 0; iy < 4; iy++) {
         const ay = flipY ? 3 - iy : iy;
         for (let ix = 0; ix < 4; ix++) set(b, flipX * 4 + ix, flipY * 4 + iy, samples[ay * 4 + (flipX ? 3 - ix : ix)]);
@@ -337,7 +359,7 @@ export function inverseTransform(b: Block, type: number): void {
 export { sub as subBlock };
 
 // the AFV basis (16 x 16)
-const AFV_BASIS = [
+export const AFV_BASIS = [
     0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25,
     0.876902929799142, 0.2206518106944235, -0.10140050393753763, -0.1014005039375375, 0.2206518106944236, -0.10140050393753777, -0.10140050393753772, -0.10140050393753763, -0.10140050393753758, -0.10140050393753769, -0.1014005039375375, -0.10140050393753768, -0.10140050393753768, -0.10140050393753759, -0.10140050393753763, -0.10140050393753741,
     0.0, 0.0, 0.40670075830260755, 0.44444816619734445, 0.0, 0.0, 0.19574399372042936, 0.2929100136981264, -0.40670075830260716, -0.19574399372042872, 0.0, 0.11379074460448091, -0.44444816619734384, -0.29291001369812636, -0.1137907446044814, 0.0,