AVIF encoder: unrolled forward DCTs, fewer trials

- the 4-, 8- and 16-point DCTs (also the leaves of the 32 and 64 point ones) as straight-line code, generated by
  scripts/av1-fdct.py: 2-4 times faster, the same coefficients
- the projections of the other transform types four coefficients at a time
- below effort 2, the second luma and chroma modes are only coded if their SATD is within 10% / 20% of the best
- the color conversion without Uint8Array.from's callbacks

Together with d5b9af2 about 40% less CPU time at effort 1 than 35be9bf, BD-rate -0.8% to +1.8% per image.

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

Author: matmen

scripts/av1-fdct.py

@@ -0,0 +1,41 @@
+# Writes src/codecs/av1/fdct-lee.ts: the unnormalized 4-, 8- and 16-point DCT-IIs of src/codecs/av1/fdct.ts's recursion
+# (Lee's) as straight-line code (the same operations, on locals instead of scratch arrays):
+#   python scripts/av1-fdct.py > src/codecs/av1/fdct-lee.ts
+import math
+counter = [0]
+lines = []
+def tmp(expr):
+    name = f't{counter[0]}'
+    counter[0] += 1
+    lines.append(f'    const {name} = {expr};')
+    return name
+
+# 1 / (2 cos(pi (2i + 1) / 2n))
+twiddle = lambda i, n: repr(1 / (2 * math.cos(math.pi * (2 * i + 1) / (2 * n))))
+
+def dct(x):
+    n = len(x)
+    if n == 2:
+        return [tmp(f'{x[0]} + {x[1]}'), tmp(f'({x[0]} - {x[1]}) * {twiddle(0, 2)}')]
+    half = n // 2
+    a = [tmp(f'{x[i]} + {x[n - 1 - i]}') for i in range(half)]
+    b = [tmp(f'({x[i]} - {x[n - 1 - i]}) * {twiddle(i, n)}') for i in range(half)]
+    even, odd = dct(a), dct(b)
+    out = [None] * n
+    for k in range(half):
+        out[2 * k] = even[k]
+        out[2 * k + 1] = tmp(f'{odd[k]} + {odd[k + 1]}') if k + 1 < half else odd[k]
+    return out
+
+print('''/*!
+ * Generated by scripts/av1-fdct.py: the 4-, 8- and 16-point DCTs of fdct.ts's recursion, unrolled
+ */''')
+for n in (4, 8, 16):
+    counter[0] = 0; lines.clear()
+    at = lambda i: f'x[o + {i}]' if i else 'x[o]'
+    out = dct([at(i) for i in range(n)])
+    print(f'\n// the unnormalized {n}-point DCT-II of x[o] in place')
+    print(f'export function dct{n}(x: Float64Array, o: number): void {{')
+    for l in lines: print(l)
+    for i in range(n): print(f'    {at(i)} = {out[i]};')
+    print('}')

src/codecs/av1/encode.ts

@@ -332,7 +332,9 @@ class TileEncoder {
         candidates.sort((x, y) => x.est - y.est);
         // (the best one's reconstruction kept)
         let best = Infinity, bestMode = candidates[0], dist = 0, bits = 0, coded = false, kept: Uint16Array | null = null, bestTx = DCT_DCT;
-        const tries = candidates.slice(0, [1, 2, 3][this.params.effort]);
+        // (below effort 2 only those about as good as the best)
+        const close = this.params.effort < 2 ? 1.1 : Infinity;
+        const tries = candidates.slice(0, [1, 2, 3][this.params.effort]).filter(c => c.est <= candidates[0].est * close);
         for (const [i, c] of tries.entries()) {
             b.yMode = c.mode;
             b.yAngle = c.angle;
@@ -396,7 +398,8 @@ class TileEncoder {
             }
             uvCandidates.sort((x, y) => x.est - y.est);
             let bestUv = Infinity, bestUvMode = uvCandidates[0], uvDist = 0, uvBits = 0, uvCoded = false, keptUv: Uint16Array[] | null = null;
-            const uvTries = uvCandidates.slice(0, [1, 2, 2][this.params.effort]);
+            const uvClose = this.params.effort < 2 ? 1.2 : Infinity;
+            const uvTries = uvCandidates.slice(0, [1, 2, 2][this.params.effort]).filter(c => c.est <= uvCandidates[0].est * uvClose);
             for (const [i, c] of uvTries.entries()) {
                 b.uvMode = c.mode;
                 b.uvAngle = c.angle;

src/codecs/av1/fdct-lee.ts

@@ -0,0 +1,162 @@
+/*!
+ * Generated by scripts/av1-fdct.py: the 4-, 8- and 16-point DCTs of fdct.ts's recursion, unrolled
+ */
+
+// the unnormalized 4-point DCT-II of x[o] in place
+export function dct4(x: Float64Array, o: number): void {
+    const t0 = x[o] + x[o + 3];
+    const t1 = x[o + 1] + x[o + 2];
+    const t2 = (x[o] - x[o + 3]) * 0.541196100146197;
+    const t3 = (x[o + 1] - x[o + 2]) * 1.3065629648763764;
+    const t4 = t0 + t1;
+    const t5 = (t0 - t1) * 0.7071067811865475;
+    const t6 = t2 + t3;
+    const t7 = (t2 - t3) * 0.7071067811865475;
+    const t8 = t6 + t7;
+    x[o] = t4;
+    x[o + 1] = t8;
+    x[o + 2] = t5;
+    x[o + 3] = t7;
+}
+
+// the unnormalized 8-point DCT-II of x[o] in place
+export function dct8(x: Float64Array, o: number): void {
+    const t0 = x[o] + x[o + 7];
+    const t1 = x[o + 1] + x[o + 6];
+    const t2 = x[o + 2] + x[o + 5];
+    const t3 = x[o + 3] + x[o + 4];
+    const t4 = (x[o] - x[o + 7]) * 0.5097955791041592;
+    const t5 = (x[o + 1] - x[o + 6]) * 0.6013448869350453;
+    const t6 = (x[o + 2] - x[o + 5]) * 0.8999762231364156;
+    const t7 = (x[o + 3] - x[o + 4]) * 2.5629154477415055;
+    const t8 = t0 + t3;
+    const t9 = t1 + t2;
+    const t10 = (t0 - t3) * 0.541196100146197;
+    const t11 = (t1 - t2) * 1.3065629648763764;
+    const t12 = t8 + t9;
+    const t13 = (t8 - t9) * 0.7071067811865475;
+    const t14 = t10 + t11;
+    const t15 = (t10 - t11) * 0.7071067811865475;
+    const t16 = t14 + t15;
+    const t17 = t4 + t7;
+    const t18 = t5 + t6;
+    const t19 = (t4 - t7) * 0.541196100146197;
+    const t20 = (t5 - t6) * 1.3065629648763764;
+    const t21 = t17 + t18;
+    const t22 = (t17 - t18) * 0.7071067811865475;
+    const t23 = t19 + t20;
+    const t24 = (t19 - t20) * 0.7071067811865475;
+    const t25 = t23 + t24;
+    const t26 = t21 + t25;
+    const t27 = t25 + t22;
+    const t28 = t22 + t24;
+    x[o] = t12;
+    x[o + 1] = t26;
+    x[o + 2] = t16;
+    x[o + 3] = t27;
+    x[o + 4] = t13;
+    x[o + 5] = t28;
+    x[o + 6] = t15;
+    x[o + 7] = t24;
+}
+
+// the unnormalized 16-point DCT-II of x[o] in place
+export function dct16(x: Float64Array, o: number): void {
+    const t0 = x[o] + x[o + 15];
+    const t1 = x[o + 1] + x[o + 14];
+    const t2 = x[o + 2] + x[o + 13];
+    const t3 = x[o + 3] + x[o + 12];
+    const t4 = x[o + 4] + x[o + 11];
+    const t5 = x[o + 5] + x[o + 10];
+    const t6 = x[o + 6] + x[o + 9];
+    const t7 = x[o + 7] + x[o + 8];
+    const t8 = (x[o] - x[o + 15]) * 0.5024192861881557;
+    const t9 = (x[o + 1] - x[o + 14]) * 0.5224986149396889;
+    const t10 = (x[o + 2] - x[o + 13]) * 0.5669440348163577;
+    const t11 = (x[o + 3] - x[o + 12]) * 0.6468217833599901;
+    const t12 = (x[o + 4] - x[o + 11]) * 0.7881546234512502;
+    const t13 = (x[o + 5] - x[o + 10]) * 1.060677685990347;
+    const t14 = (x[o + 6] - x[o + 9]) * 1.7224470982383342;
+    const t15 = (x[o + 7] - x[o + 8]) * 5.101148618689155;
+    const t16 = t0 + t7;
+    const t17 = t1 + t6;
+    const t18 = t2 + t5;
+    const t19 = t3 + t4;
+    const t20 = (t0 - t7) * 0.5097955791041592;
+    const t21 = (t1 - t6) * 0.6013448869350453;
+    const t22 = (t2 - t5) * 0.8999762231364156;
+    const t23 = (t3 - t4) * 2.5629154477415055;
+    const t24 = t16 + t19;
+    const t25 = t17 + t18;
+    const t26 = (t16 - t19) * 0.541196100146197;
+    const t27 = (t17 - t18) * 1.3065629648763764;
+    const t28 = t24 + t25;
+    const t29 = (t24 - t25) * 0.7071067811865475;
+    const t30 = t26 + t27;
+    const t31 = (t26 - t27) * 0.7071067811865475;
+    const t32 = t30 + t31;
+    const t33 = t20 + t23;
+    const t34 = t21 + t22;
+    const t35 = (t20 - t23) * 0.541196100146197;
+    const t36 = (t21 - t22) * 1.3065629648763764;
+    const t37 = t33 + t34;
+    const t38 = (t33 - t34) * 0.7071067811865475;
+    const t39 = t35 + t36;
+    const t40 = (t35 - t36) * 0.7071067811865475;
+    const t41 = t39 + t40;
+    const t42 = t37 + t41;
+    const t43 = t41 + t38;
+    const t44 = t38 + t40;
+    const t45 = t8 + t15;
+    const t46 = t9 + t14;
+    const t47 = t10 + t13;
+    const t48 = t11 + t12;
+    const t49 = (t8 - t15) * 0.5097955791041592;
+    const t50 = (t9 - t14) * 0.6013448869350453;
+    const t51 = (t10 - t13) * 0.8999762231364156;
+    const t52 = (t11 - t12) * 2.5629154477415055;
+    const t53 = t45 + t48;
+    const t54 = t46 + t47;
+    const t55 = (t45 - t48) * 0.541196100146197;
+    const t56 = (t46 - t47) * 1.3065629648763764;
+    const t57 = t53 + t54;
+    const t58 = (t53 - t54) * 0.7071067811865475;
+    const t59 = t55 + t56;
+    const t60 = (t55 - t56) * 0.7071067811865475;
+    const t61 = t59 + t60;
+    const t62 = t49 + t52;
+    const t63 = t50 + t51;
+    const t64 = (t49 - t52) * 0.541196100146197;
+    const t65 = (t50 - t51) * 1.3065629648763764;
+    const t66 = t62 + t63;
+    const t67 = (t62 - t63) * 0.7071067811865475;
+    const t68 = t64 + t65;
+    const t69 = (t64 - t65) * 0.7071067811865475;
+    const t70 = t68 + t69;
+    const t71 = t66 + t70;
+    const t72 = t70 + t67;
+    const t73 = t67 + t69;
+    const t74 = t57 + t71;
+    const t75 = t71 + t61;
+    const t76 = t61 + t72;
+    const t77 = t72 + t58;
+    const t78 = t58 + t73;
+    const t79 = t73 + t60;
+    const t80 = t60 + t69;
+    x[o] = t28;
+    x[o + 1] = t74;
+    x[o + 2] = t42;
+    x[o + 3] = t75;
+    x[o + 4] = t32;
+    x[o + 5] = t76;
+    x[o + 6] = t43;
+    x[o + 7] = t77;
+    x[o + 8] = t29;
+    x[o + 9] = t78;
+    x[o + 10] = t44;
+    x[o + 11] = t79;
+    x[o + 12] = t31;
+    x[o + 13] = t80;
+    x[o + 14] = t40;
+    x[o + 15] = t69;
+}

src/codecs/av1/fdct.ts

@@ -5,6 +5,7 @@
  */
 import { TXFM_DIMENSIONS } from './tables.ts';
 import { inverseBasis, transformTypes1D, passShift } from './itx.ts';
+import { dct4, dct8, dct16 } from './fdct-lee.ts';
 
 // the pseudo inverse (B^T B)^-1 B^T of a basis of n x k values: k x n
 function pseudoInverse(b: Float64Array, n: number, k: number): Float64Array {
@@ -62,19 +63,16 @@ const tmp = new Float64Array(64 * 32);
 
 // the unnormalized DCT-II (X_k = sum of x_i cos(pi (2i + 1) k / 2n)) of n values at x[o] (step 1) in place, by Lee's
 // recursion (scratch: n values)
-// (1 / (2 cos(pi (2i + 1) / 2n)) for n = 2-64, at n / 2 + i)
+// (1 / (2 cos(pi (2i + 1) / 2n)) for n = 32 and 64, at n / 2 + i)
 const TWIDDLES = new Float64Array(64);
-for (let n = 2; n <= 64; n *= 2) for (let i = 0; i < n / 2; i++) TWIDDLES[n / 2 + i] = 1 / (2 * Math.cos(Math.PI * (2 * i + 1) / (2 * n)));
+for (let n = 32; n <= 64; n *= 2) for (let i = 0; i < n / 2; i++) TWIDDLES[n / 2 + i] = 1 / (2 * Math.cos(Math.PI * (2 * i + 1) / (2 * n)));
 
 function dct(x: Float64Array, o: number, n: number, scratch: Float64Array, so: number): void {
-    if (n === 1) return;
+    // (up to 16: unrolled)
+    if (n === 4) return dct4(x, o);
+    if (n === 8) return dct8(x, o);
+    if (n === 16) return dct16(x, o);
     const half = n / 2;
-    if (n === 2) {
-        const a = x[o], b = x[o + 1];
-        x[o] = a + b;
-        x[o + 1] = (a - b) * TWIDDLES[1];
-        return;
-    }
     for (let i = 0; i < half; i++) {
         const a = x[o + i], b = x[o + n - 1 - i];
         scratch[so + i] = a + b;
@@ -113,11 +111,22 @@ function forward1D(log2: number, type: number, inp: Float64Array, io: number, st
         for (let j = 0; j < k; j++) out[oo + j * outStep] = line[j] * s[j];
         return;
     }
+    // (four coefficients at a time)
     const p = projection(log2, type);
-    for (let j = 0; j < k; j++) {
-        let sum = 0;
-        for (let i = 0; i < n; i++) sum += p[j * n + i] * inp[io + i * step];
-        out[oo + j * outStep] = sum;
+    for (let i = 0; i < n; i++) line[i] = inp[io + i * step];
+    for (let j = 0; j < k; j += 4) {
+        let s0 = 0, s1 = 0, s2 = 0, s3 = 0;
+        for (let i = 0, r = j * n; i < n; i++, r++) {
+            const v = line[i];
+            s0 += p[r] * v;
+            s1 += p[r + n] * v;
+            s2 += p[r + 2 * n] * v;
+            s3 += p[r + 3 * n] * v;
+        }
+        out[oo + j * outStep] = s0;
+        out[oo + (j + 1) * outStep] = s1;
+        out[oo + (j + 2) * outStep] = s2;
+        out[oo + (j + 3) * outStep] = s3;
     }
 }
 

src/codecs/avif-enc.ts

@@ -40,7 +40,14 @@ function toYUV(rgba: Uint8Array, width: number, height: number, subsampled: bool
         v[i] = (r - l) / (2 * (1 - KR)) + 128;
     }
     const byte = (x: number): number => x < 0 ? 0 : x > 255 ? 255 : Math.round(x);
-    if (!subsampled) return [y, Uint8Array.from(u, byte), Uint8Array.from(v, byte)];
+    if (!subsampled) {
+        const cu = new Uint8Array(n), cv = new Uint8Array(n);
+        for (let i = 0; i < n; i++) {
+            cu[i] = byte(u[i]);
+            cv[i] = byte(v[i]);
+        }
+        return [y, cu, cv];
+    }
     // (the average of each 2x2 block)
     const cw = width + 1 >> 1, ch = height + 1 >> 1, cu = new Uint8Array(cw * ch), cv = new Uint8Array(cw * ch);
     for (let cy = 0; cy < ch; cy++) {