More experiments for different numerical kernels

Author: FrancescAlted

bench/js-transpiler/README.md

@@ -17,20 +17,37 @@ then benches a 24-frame `relax` sweep.
 
 ```sh
 npm i                                         # pulls pyodide@314 (see package.json)
-node bench/js-transpiler/dsl-js-node.mjs       # correctness + 24-frame bench
-node bench/js-transpiler/dsl-js-node.mjs 48    # N frames
+node bench/js-transpiler/dsl-js-node.mjs       # correctness + kernel sweep, 12 reps
+node bench/js-transpiler/dsl-js-node.mjs 24    # N reps
 ```
 
 Needs network on first run (PyPI wheel via micropip). Exits non-zero on a correctness
-mismatch, so it works as a smoke test. Typical output (Apple M-series, blosc2 4.6.0):
+mismatch, so it works as a smoke test. It benches four kernel shapes so the js-vs-JIT ratio
+can be read against the kernel, not generalized from one. Representative (Apple M2, 4.6.0):
 
 ```
-correctness vs numpy:  js maxdiff=0.00e+0  jit maxdiff=0.00e+0  OK
-  jit_backend="js" : ~16 ms/frame
-  jit (miniexpr)   : ~31 ms/frame      -> js ~2x faster
-  no-JIT           : ~130 ms/frame     -> js ~8x faster
+  kernel    js     jit    nojit   js/jit  js/nojit
+  newton    12.1   26.1   114.6   2.15x    9.44x     arithmetic + branches + early-exit
+  deepar    22.5   50.2    53.1   2.23x    2.36x     deep pure-arithmetic loop
+  deep     120.1  143.1   104.4   1.19x    0.87x     deep loop, libm sin every iter
+  trans      5.0    5.0     6.8   1.00x    1.34x     transcendental-heavy
+  poly       3.4    3.0     3.6   0.87x    1.05x     light, branch-free
 ```
 
+**The takeaway: there is no single "js is N× the JIT" number — it depends on what the kernel
+is bottlenecked on.**
+
+- **Arithmetic / control-flow bound** (newton, deepar) → V8's optimizing JIT beats blosc2's
+  miniexpr WASM codegen by **~2×**. This is the sweet spot.
+- **Transcendental bound** (trans, deep) → **~1×**: time is spent in `sin`/`exp`/`log` (libm),
+  which costs about the same whoever runs the loop — `nojit` even edges `js` on `deep`.
+- **Light / trivial** (poly) → **<1×**: the kernel does almost no compute, so the blosc2
+  pipeline + per-call JS marshaling dominate, and `js` can be *slightly slower* than the JIT.
+
+So the honest generalization is qualitative: transpiling to JS wins (~2×, single-threaded)
+for **compute-bound float kernels dominated by arithmetic and control flow**, and is roughly
+a wash for transcendental-bound or trivial kernels.
+
 > The overlay pins `blosc2==4.6.0` to keep the compiled `blosc2_ext` ABI in step with the
 > pure-Python we drop on top. Once these changes ship in a Pyodide-installable wheel, the
 > overlay can go away. If the overlay import ever breaks on version skew, overlay all of

bench/js-transpiler/dsl-js-node.mjs

@@ -3,18 +3,22 @@
 // working tree's pure-Python (src/blosc2/dsl_js.py + lazyexpr.py) on top of it before
 // importing blosc2 -- so the wired path runs without waiting for a new wheel.
 //
+// Benches a spread of kernel shapes to show how the js-vs-JIT ratio depends on the kernel:
+// branchy + early-exit (newton), branch-free light (poly), transcendental-heavy (trans),
+// deep no-exit loop (deep). Reports js / jit / no-jit per kernel.
+//
 //   npm i                                         # pulls pyodide@314 (see package.json)
-//   node bench/js-transpiler/dsl-js-node.mjs      # correctness + 24-frame bench
-//   node bench/js-transpiler/dsl-js-node.mjs 48   # bench with N frames
+//   node bench/js-transpiler/dsl-js-node.mjs       # correctness + bench, 12 reps
+//   node bench/js-transpiler/dsl-js-node.mjs 24    # N reps
 import { loadPyodide } from "pyodide";
 import { readFileSync } from "node:fs";
 import { fileURLToPath } from "node:url";
 
 // Resolve paths from this file, so the harness runs from any CWD (repo root is ../../).
 const ROOT = fileURLToPath(new URL("../../", import.meta.url));
-const NFRAMES = Number(process.argv[2]) || 24;
+const NFRAMES = Number(process.argv[2]) || 12;
 
-// Kernel + bench live in a real module file: @blosc2.dsl_kernel runs inspect.getsource(),
+// Kernels + bench live in a real module file: @blosc2.dsl_kernel runs inspect.getsource(),
 // which needs the function to be file-backed (not exec'd from a string).
 const PYSRC = String.raw`
 import json, time
@@ -26,7 +30,8 @@ SPANX = 3.4
 ASPECT = HEIGHT / WIDTH
 DTYPE = np.float64
 
-@blosc2.dsl_kernel
+# --- kernels spanning the cost/control-flow spectrum -------------------------------------
+@blosc2.dsl_kernel  # branchy, deep, per-pixel early exit
 def newton_dsl(a, b, max_iter, relax):
     za = a
     zb = b
@@ -59,27 +64,34 @@ def newton_dsl(a, b, max_iter, relax):
         root = 2.0
     return root + 0.9 * (it / mif)
 
-def newton_numpy(a, b, max_iter, relax):
-    za = a.copy(); zb = b.copy()
-    mif = float(max_iter)
-    it = np.full(a.shape, mif)
-    alive = np.ones(a.shape, dtype=bool)
-    for k in range(max_iter):
-        a2 = za * za; b2 = zb * zb
-        fr = za * a2 - 3.0 * za * b2 - 1.0; fi = 3.0 * a2 * zb - zb * b2
-        dr = 3.0 * (a2 - b2); di = 6.0 * za * zb; den = dr * dr + di * di + 1e-12
-        qr = relax * (fr * dr + fi * di) / den; qi = relax * (fi * dr - fr * di) / den
-        za = za - qr; zb = zb - qi
-        c = alive & ((qr * qr + qi * qi) < 1e-6); it[c] = k; alive &= ~c
-        if not alive.any():
-            break
-    d0 = (za - 1.0) ** 2 + zb ** 2
-    d1 = (za + 0.5) ** 2 + (zb - 0.8660254) ** 2
-    d2 = (za + 0.5) ** 2 + (zb + 0.8660254) ** 2
-    root = np.zeros(a.shape); md = d0.copy()
-    m = d1 < md; md = np.where(m, d1, md); root = np.where(m, 1.0, root)
-    m = d2 < md; root = np.where(m, 2.0, root)
-    return root + 0.9 * (it / mif)
+@blosc2.dsl_kernel  # light, branch-free, vectorizable arithmetic
+def poly_dsl(a, b):
+    a2 = a * a
+    b2 = b * b
+    return a2 * a - 3.0 * a * b2 + 2.0 * b2 * b - a + 0.5 * b
+
+@blosc2.dsl_kernel  # transcendental-heavy (exercises each engine's libm). miniexpr wants
+def trans_dsl(a, b):  # bare sin/cos/... (not np.sin); the transpiler maps both to Math.*
+    msq = a * a + b * b
+    sc = sin(a * 3.0) * cos(b * 2.0)
+    ex = exp(msq * -0.5)
+    return sc + ex + sqrt(msq + 1.0) + log(msq + 2.0)
+
+@blosc2.dsl_kernel  # deep fixed loop, transcendental-bound (libm sin every iter)
+def deep_dsl(a, b):
+    acc = a
+    for k in range(64):
+        acc = acc * 0.99 + sin(acc + b)
+    return acc
+
+@blosc2.dsl_kernel  # deep fixed loop, pure arithmetic (no libm, no branches); contractive
+def deepar_dsl(a, b):
+    acc = a * 0.1
+    t = b * 0.1
+    for k in range(64):
+        t = t * 0.5 - acc * 0.25 + 0.1
+        acc = acc * 0.5 + t * 0.25 + 0.1
+    return acc + t
 
 _x = np.linspace(-SPANX / 2, SPANX / 2, WIDTH, dtype=DTYPE)
 _y = np.linspace(-SPANX * ASPECT / 2, SPANX * ASPECT / 2, HEIGHT, dtype=DTYPE)
@@ -90,40 +102,36 @@ _cp = blosc2.CParams(codec=blosc2.Codec.LZ4, clevel=1)
 A_B2 = blosc2.asarray(A_NP, chunks=_chunks, blocks=_blocks, cparams=_cp)
 B_B2 = blosc2.asarray(B_NP, chunks=_chunks, blocks=_blocks, cparams=_cp)
 
-def run(relax, backend):
+# (name, kernel, operand tuple). Fixed inputs -> each rep does identical work.
+KERNELS = [
+    ("newton", newton_dsl, (A_B2, B_B2, MAXITER, 1.37)),
+    ("poly",   poly_dsl,   (A_B2, B_B2)),
+    ("trans",  trans_dsl,  (A_B2, B_B2)),
+    ("deep",   deep_dsl,   (A_B2, B_B2)),
+    ("deepar", deepar_dsl, (A_B2, B_B2)),
+]
+
+def run(func, ops, backend):
     kw = {"dtype": DTYPE, "cparams": _cp}
     if backend == "js":
         kw["jit_backend"] = "js"
     elif backend == "jit":
         kw["jit"] = True
     else:
         kw["jit"] = False
-    return blosc2.lazyudf(newton_dsl, (A_B2, B_B2, MAXITER, relax), **kw)[:]
-
-def correctness():
-    rx = 1.37
-    ref = newton_numpy(A_NP, B_NP, MAXITER, rx)
-    return {
-        "diff_js": float(np.max(np.abs(run(rx, "js") - ref))),
-        "diff_jit": float(np.max(np.abs(run(rx, "jit") - ref))),
-    }
-
-def bench(nframes):
-    relaxes = [float(v) for v in np.linspace(1.0, 1.85, nframes)]
-    def sweep_ms(backend):
-        for r in relaxes[:2]:
-            run(r, backend)            # warm
-        best = float("inf")
-        for _ in range(3):
-            t = time.perf_counter()
-            for r in relaxes:
-                run(r, backend)
-            best = min(best, (time.perf_counter() - t) * 1000)
-        return best
-    return {b: sweep_ms(b) for b in ("js", "jit", "nojit")}
+    return blosc2.lazyudf(func, ops, **kw)[:]
+
+def bench(func, ops, backend, reps):
+    run(func, ops, backend)            # warm
+    best = float("inf")
+    for _ in range(3):
+        t = time.perf_counter()
+        for _ in range(reps):
+            run(func, ops, backend)
+        best = min(best, (time.perf_counter() - t) * 1000 / reps)
+    return best
 
 def debug_bridge():
-    # Call the JS bridge directly (no lazyudf/blosc2 machinery) to isolate marshal+compile+compute.
     import blosc2.dsl_js as dj
     bridge = dj.js_kernel(newton_dsl)
     out = np.empty((HEIGHT, WIDTH), dtype=DTYPE)
@@ -134,12 +142,17 @@ def debug_bridge():
         t = time.perf_counter(); bridge(inp, out, 0); best = min(best, (time.perf_counter() - t) * 1000)
     return {"first_ms": first, "warm_ms": best}
 
-def result(nframes):
-    out = correctness()
-    out["ms"] = bench(nframes)
-    out["bridge"] = debug_bridge()
-    out["nframes"] = nframes
-    return json.dumps(out)
+def result(reps):
+    import math
+    kernels = []
+    for name, func, ops in KERNELS:
+        rj = run(func, ops, "js")
+        rjit = run(func, ops, "jit")
+        diff = float(np.max(np.abs(rj - rjit)))
+        diff = diff if math.isfinite(diff) else 1e30   # keep JSON valid; flags as mismatch
+        ms = {b: bench(func, ops, b, reps) for b in ("js", "jit", "nojit")}
+        kernels.append({"name": name, "ms": ms, "diff": diff})
+    return json.dumps({"kernels": kernels, "bridge": debug_bridge(), "reps": reps})
 `;
 
 const py = await loadPyodide();
@@ -160,27 +173,28 @@ import sys, blosc2
 assert hasattr(sys.modules["blosc2.lazyexpr"], "_as_js_udf"), "overlay did not take"
 `);
 console.log("blosc2", await py.runPythonAsync("blosc2.__version__"),
-            "| Pyodide", py.version, "| frames", NFRAMES);
+            "| Pyodide", py.version, "| reps", NFRAMES);
 
-py.FS.writeFile("/newton_bench.py", new TextEncoder().encode(PYSRC));
+py.FS.writeFile("/kernel_bench.py", new TextEncoder().encode(PYSRC));
 const out = await py.runPythonAsync(`
 import sys
 if "/" not in sys.path: sys.path.insert(0, "/")
-import newton_bench
-newton_bench.result(${NFRAMES})
+import kernel_bench
+kernel_bench.result(${NFRAMES})
 `);
 
 const r = JSON.parse(out);
-const ok = r.diff_js < 1e-9 && r.diff_jit < 1e-9;
-console.log(`\ncorrectness vs numpy:  js maxdiff=${r.diff_js.toExponential(2)}  ` +
-            `jit maxdiff=${r.diff_jit.toExponential(2)}  ${ok ? "OK" : "MISMATCH"}`);
-const per = ms => `${ms.toFixed(0)} ms total (${(ms / r.nframes).toFixed(2)} ms/frame)`;
-console.log(`\nperf (${r.nframes}-frame relax sweep, best of 3):`);
-console.log(`  jit_backend="js" : ${per(r.ms.js)}`);
-console.log(`  jit (miniexpr)   : ${per(r.ms.jit)}`);
-console.log(`  no-JIT           : ${per(r.ms.nojit)}`);
-console.log(`  js vs jit  : ${(r.ms.jit / r.ms.js).toFixed(2)}x   ` +
-            `js vs no-jit : ${(r.ms.nojit / r.ms.js).toFixed(2)}x`);
-console.log(`\nbridge probe (direct call, no blosc2 machinery):  ` +
-            `first=${r.bridge.first_ms.toFixed(1)} ms  warm=${r.bridge.warm_ms.toFixed(1)} ms`);
-if (!ok) process.exit(1);
+const fmt = (x, w) => String(x).padStart(w);
+const bad = r.kernels.filter((k) => k.diff > 1e-5);
+console.log(`\ncorrectness (js vs JIT maxdiff): ${bad.length ? "MISMATCH " + bad.map((k) => k.name) : "OK"}`);
+console.log("\nper-kernel bench (ms/frame, lower is better):");
+console.log("  kernel    js     jit    nojit   js/jit  js/nojit  diff");
+for (const k of r.kernels) {
+  const { js, jit, nojit } = k.ms;
+  console.log(
+    `  ${k.name.padEnd(7)} ${fmt(js.toFixed(1), 6)} ${fmt(jit.toFixed(1), 6)} ${fmt(nojit.toFixed(1), 7)}` +
+    `  ${fmt((jit / js).toFixed(2) + "x", 6)} ${fmt((nojit / js).toFixed(2) + "x", 8)}  ${k.diff.toExponential(1)}`,
+  );
+}
+console.log(`\nnewton bridge probe (no blosc2 machinery): first=${r.bridge.first_ms.toFixed(1)} ms  warm=${r.bridge.warm_ms.toFixed(1)} ms`);
+if (bad.length) process.exit(1);

plans/dsl-js.md

@@ -142,6 +142,13 @@ Measured on Apple M-series, blosc2 4.6.0 under Pyodide 314, Newton 320×213 / ma
 
 Correctness exact: `js` and JIT both `maxdiff=0.00` vs numpy.
 
+**The ~2× is kernel-dependent, not a flat rule** (kernel sweep in `dsl-js-node.mjs`, see
+`bench/js-transpiler/README.md`). The js-vs-JIT win tracks what the kernel is bottlenecked on:
+**arithmetic / control-flow** bound → ~2× (newton 2.15×, a deep pure-arithmetic loop 2.23×);
+**transcendental** bound (`sin`/`exp`/`log`) → ~1× (libm cost is engine-independent; `nojit`
+can even edge `js`); **light / trivial** → <1× (blosc2 pipeline + marshaling dominate, `js`
+slightly loses). So JS helps for compute-bound float kernels heavy on arithmetic and branches.
+
 **The PyProxy gotcha (the one real bug the headless harness caught).** The bridge must pass
 the per-call operands to the JS driver as real **JS `Array`s**, not Python lists. A Python
 list arrives in JS as a `PyProxy`, so every `ops[k][i]` in the hot inner loop crosses the