Skip to content

Add a DSL → JavaScript transpiler backend (jit_backend="js")#670

Merged
FrancescAlted merged 13 commits into
mainfrom
js-transpiler
Jun 29, 2026
Merged

Add a DSL → JavaScript transpiler backend (jit_backend="js")#670
FrancescAlted merged 13 commits into
mainfrom
js-transpiler

Conversation

@FrancescAlted

Copy link
Copy Markdown
Member

Adds a new execution backend for @blosc2.dsl_kernel kernels that transpiles them to JavaScript and runs them in-browser (Pyodide/WASM), where V8 JIT-compiles the scalar loop to native code. On compute-heavy kernels it beats blosc2's WASM TinyCC JIT substantially.

What's new

JS transpiler (src/blosc2/dsl_js.py)

  • Transpiles a DSL kernel's AST to a self-contained JS module (arithmetic, comparisons, where, if/elif/else, range loops, **////%, and a whitelist of math functions).
  • Index/shape symbols (_i0/_n0/_flat_idx, …): emitted as trailing kernel params; the runtime driver reconstructs each element's global coordinate per block from (offset, gshape, cshape).
  • Transpile result and the V8-compiled __run are memoized so repeated evaluations (e.g. animation loops) don't re-parse/re-eval.

Backend routing (src/blosc2/lazyexpr.py)

  • Under WebAssembly, the default prefers JS for transpilable float DSL kernels and silently falls back to miniexpr+jit-wasm for anything it can't handle (non-float output, reductions, unsupported constructs) — no regression.
  • Explicit jit_backend="js" transpiles or raises; jit=False / strict_miniexpr=True / an explicit jit_backend opt out.
  • Integer inputs with a floating output now use JS too (the bridge float64-converts operands, exactly matching miniexpr's promotion). Integer/complex output still routes to miniexpr.

Coverage (what runs on JS vs miniexpr)

Routed to JS Falls back to miniexpr+jit-wasm
float64/float32 element-wise kernels integer/complex output
index/shape symbols (_i0/_n0/_flat_idx) reductions
integer inputs + float output zero-input DSL kernels
arithmetic, where, if/elif/else, range, math whitelist unsupported constructs (ternary, chained comparisons, tuple/subscript assignment, non-whitelisted calls)

Full breakdown and remaining work (integer output, zero-copy block I/O) in plans/dsl-js-coverage.md.

Performance (Pyodide, ms/frame; bench/js-transpiler/dsl-js-node.mjs)

JS wins on compute-heavy kernels and is at parity/slightly behind on light, vectorizable ones (where per-block marshaling dominates):

newton 2.80×   deepar 2.78×   idxgrad 2.00×   deep 1.30×
trans  0.99×   intmix 0.87×   poly   0.86×

The remaining sub-1.0 cases are marshaling cost (the bridge copies blocks in/out; miniexpr computes in place), documented as a future zero-copy HEAPF64 optimization.

Tests & CI

  • New transpiler tests (tests/ndarray/test_dsl_js.py) and WASM integration tests (tests/ndarray/test_wasm_dsl_jit.py) covering index symbols, integer inputs, prefer-js selection, and silent fallback.
  • Node-subprocess equivalence tests skip on emscripten (can't spawn processes); JIT-assertion tests skip on Windows wheels (no bundled miniexpr JIT backend).
  • Hardened a flaky b2view TUI test (wait_until poll instead of a single pilot.pause()).
  • Bench harness streams per-kernel rows as computed and verifies all backend paths agree with miniexpr.

Transpile a @blosc2.dsl_kernel (the bounded Python-ast subset DSLValidator
accepts) to JavaScript, so kernels run at V8-optimized native speed in the
browser/Pyodide. Wired into chunked_eval as a new backend alongside no-JIT and
miniexpr-JIT: the kernel becomes a plain per-block UDF callable, so there are no
compiled-code changes and the default path is untouched.

Verified end-to-end under Pyodide (Newton 320x213, 24-frame sweep): correctness
exact vs numpy, ~2x faster than the miniexpr JIT, ~8x over no-JIT. The bridge
must hand the JS driver real JS Arrays (not Python lists), else the hot loop
crosses the Python<->JS boundary per element (~10x slower).

- src/blosc2/dsl_js.py: dsl_to_js(), build_js_module(), js_kernel() bridge
- src/blosc2/lazyexpr.py: _as_js_udf() + jit_backend="js" swap in chunked_eval
- tests/ndarray/test_dsl_js.py: transpiler + node numeric-equivalence tests
- bench/js-transpiler/: headless Node+Pyodide bench, browser demo, README
- plans/dsl-js.md: design and verified bench findings
…ack)

Under WASM, a float/transpilable/non-reduction DSL kernel now auto-routes to
jit_backend="js" unless the user opts out; anything JS can't do (non-float
dtypes, reductions, unsupported constructs) silently falls back to miniexpr,
so there's no regression. Since JS is itself a JIT, jit=True prefers it too —
only jit=False, strict_miniexpr=True, or an explicit jit_backend opts out
(force miniexpr with jit_backend="tcc"/"cc").

- lazyexpr.py: _maybe_js_backend prefer-js logic + _js_dtypes_ok gating; "js"
  documented and listed for the jit_backend param.
- test_dsl_kernels.py: autouse fixture keeps this miniexpr-semantics module on
  miniexpr (the prefer-js default would bypass its _set_pref_expr assertions).
- test_wasm_dsl_jit.py: native-CI coverage for explicit js, the prefer-js
  default, and the int fallback (counterpart of the node overlay harness).
- bench/js-transpiler: kernel sweep (newton/poly/trans/deep/deepar) showing the
  js-vs-tcc win is kernel-dependent (~2x arithmetic, ~1x transcendental/light).
Add P1 (index/shape symbols _i0/_n0/_flat_idx) by emitting them as trailing
kernel params and reconstructing per-block global coords in the driver, and
P2 (integer inputs with floating output, matching miniexpr's float promotion).
Memoize transpile + js.eval so repeated evaluations don't re-parse/re-compile,
which closes most of the gap on light kernels. Update tests, the node bench
(new P1/P2 kernels, streamed rows), and plans/dsl-js-coverage.md.
The explicit JS path bypassed the output-dtype check, so an integer/complex
output kernel would silently compute in float64 instead of failing. Raise a
clear ValueError up front (matching the documented "explicit js raises on
gaps" contract); integer inputs with a floating output still use JS. Keeps
integer output entirely on miniexpr. Add a test and update the coverage plan.
@FrancescAlted FrancescAlted merged commit e23cf67 into main Jun 29, 2026
21 checks passed
@FrancescAlted FrancescAlted deleted the js-transpiler branch June 29, 2026 12:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant