MATLAB-syntax scripts on WebGPU: fused kernels, tic/toc timing, CPU comparison
37 changed files+8644−0
.github/workflows/ci.ymladded+42−0View file
@@ -0,0 +1,42 @@
1+name: ci
2+on:
3+ push:
4+ branches: [main]
5+ pull_request:
6+
7+jobs:
8+ test:
9+ runs-on: ubuntu-latest
10+ steps:
11+ - uses: actions/checkout@v4
12+ - uses: actions/setup-node@v4
13+ with:
14+ node-version: 24
15+ cache: npm
16+ # numbl is a `file:../../numbl` dependency: the GPU compiler reaches its
17+ # JIT internals (parser, lowerer, IR, inline pass, builtin registry)
18+ # through the `numbl-src` vite alias — those need no build — and the CPU
19+ # comparison imports `numbl/browser`, whose dist-browser/ is NOT
20+ # committed, so it is built here. Pinned so a change to the internals
21+ # cannot silently break this repo — the surface we rely on is written
22+ # down in src/mgpu/numbl.d.ts.
23+ - name: Check out numbl (sibling dependency)
24+ env:
25+ NUMBL_REF: 38ce14046d64d03ecf05cb57def53057a6bc64ab
26+ run: |
27+ git clone --filter=blob:none --no-checkout \
28+ https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
29+ git -C "$GITHUB_WORKSPACE/../../numbl" checkout --quiet "$NUMBL_REF"
30+ # --ignore-scripts both times: npm would run numbl's `prepare` (husky).
31+ - name: Build numbl/browser (dist-browser)
32+ run: |
33+ cd "$GITHUB_WORKSPACE/../../numbl"
34+ npm ci --ignore-scripts
35+ npm run build:browser
36+ - run: npm ci --ignore-scripts
37+ # The suite compiles MATLAB to compute shaders, so it needs a GPU; the
38+ # browser run below falls back to SwiftShader when there is none.
39+ - run: npm run test:node -- --skip-without-gpu
40+ - run: npm run test:gpu
41+ env:
42+ CHROME_PATH: /usr/bin/google-chrome
.github/workflows/deploy.ymladded+52−0View file
@@ -0,0 +1,52 @@
1+name: deploy
2+on:
3+ push:
4+ branches: [main]
5+ workflow_dispatch:
6+
7+permissions:
8+ contents: read
9+ pages: write
10+ id-token: write
11+
12+concurrency:
13+ group: pages
14+ cancel-in-progress: true
15+
16+jobs:
17+ build-deploy:
18+ runs-on: ubuntu-latest
19+ environment:
20+ name: github-pages
21+ url: ${{ steps.deployment.outputs.page_url }}
22+ steps:
23+ - uses: actions/checkout@v4
24+ - uses: actions/setup-node@v4
25+ with:
26+ node-version: 24
27+ cache: npm
28+ # See ci.yml for why numbl is cloned and dist-browser built.
29+ - name: Check out numbl (sibling dependency)
30+ env:
31+ NUMBL_REF: 38ce14046d64d03ecf05cb57def53057a6bc64ab
32+ run: |
33+ git clone --filter=blob:none --no-checkout \
34+ https://github.com/flatironinstitute/numbl.git "$GITHUB_WORKSPACE/../../numbl"
35+ git -C "$GITHUB_WORKSPACE/../../numbl" checkout --quiet "$NUMBL_REF"
36+ - name: Build numbl/browser (dist-browser)
37+ run: |
38+ cd "$GITHUB_WORKSPACE/../../numbl"
39+ npm ci --ignore-scripts
40+ npm run build:browser
41+ - run: npm ci --ignore-scripts
42+ - run: npm run test:node -- --skip-without-gpu
43+ - run: npm run build
44+ # Pages must already be enabled with "GitHub Actions" as the source; the
45+ # workflow token cannot create the site itself (`enablement: true` fails
46+ # with "Resource not accessible by integration").
47+ - uses: actions/configure-pages@v5
48+ - uses: actions/upload-pages-artifact@v3
49+ with:
50+ path: dist
51+ - id: deployment
52+ uses: actions/deploy-pages@v4
.gitignoreadded+3−0View file
@@ -0,0 +1,3 @@
1+node_modules/
2+dist/
3+*.log
CLAUDE.mdadded+81−0View file
@@ -0,0 +1,81 @@
1+# CLAUDE.md
2+
3+Tips for future agents working in this repo.
4+
5+## Architecture
6+
7+```
8+src/mgpu/compile.ts MATLAB source -> numbl IR: parseMFile + lowerProgram
9+ (whole-script lowering, NOT turing-surface's per-function
10+ specialization), numbl's inlinePass, then fuse.ts.
11+src/mgpu/patches.ts Type-rule patches applied to numbl's JIT builtin
12+ registry (registerBuiltin overwrites by name): precise
13+ shapes for comparisons, tensor &/|/~, two-arg max/min,
14+ matrix rand, and randn registered from scratch. Only
15+ types — numbl's emitters never run here.
16+src/mgpu/fuse.ts Folds single-use _mtoc2_* temps numbl's inline pass
17+ declines (tensor-producing Calls like sin(x), logicals),
18+ so one source line = one kernel.
19+src/mgpu/wgsl.ts Fused elementwise emitter -> one WGSL kernel per Assign.
20+ Also: inline generators (rand/randn/linspace/ranges/eye
21+ computed from the linear index), loop-var uniforms,
22+ runtime scalars as 1-element buffers, exact-value
23+ folding at every node.
24+src/mgpu/kernels.ts Non-elementwise WGSL: tiled column-major GEMM, tiled
25+ transpose, 2-pass full reduction + per-column reduction,
26+ all with shapes baked in as constants.
27+src/mgpu/plan.ts IR statements -> static op sequence: buffers, pipelines
28+ (cached by WGSL text), bind groups, aliasing (X(:) etc.
29+ as views when the source has exactly one assignment),
30+ scratch+copy-back for in-place updates, `for` bodies
31+ planned once with per-iteration dynamic-offset uniform
32+ slots (256 B each), host ops (tic/toc/disp/fprintf/echo).
33+src/mgpu/run.ts Executor: streams ops into command encoders; tic/toc
34+ flush + await onSubmittedWorkDone (that's what makes toc
35+ MATLAB-comparable); readbacks for printing.
36+src/cpu/cpuRunner.ts Optional CPU column via createNumblSession
37+ (numbl/browser) — numbl's own worker, f64, lazy-loaded
38+ because it's a ~3 MB chunk.
39+test/cases.ts One suite, two harnesses: scripts/test-node.ts (Dawn via
40+ the `webgpu` npm package) and test/test-page.ts +
41+ scripts/test-gpu.mjs (headless Chrome, SwiftShader
42+ fallback).
43+```
44+
45+## Key gotchas
46+
47+- **`resolve.preserveSymlinks: true` in vite.config.ts is load-bearing.**
48+ numbl is a `file:` symlink; without it the dev server canonicalizes the
49+ symlink for some import chains but not others, loading numbl's builtin
50+ registry TWICE — patches.ts then patches one instance while the lowerer
51+ consults the other ("JS-JIT 'rand' supports only the scalar form"), and
52+ raw-realpath requests trip the fs allow list.
53+- **The registry patch is per-module-instance.** The CPU runner is safe from
54+ it because numbl/browser runs in its own worker; don't move patching
55+ somewhere the CPU path could share.
56+- **Everything is f32 and column-major.** `A(:)`, `reshape` and vector
57+ transpose are buffer views (or plain copies when the source is reassigned);
58+ matrix transpose is a real kernel.
59+- **Aliased writes need scratch+copy-back** (`u = u + 1`, GEMM with output
60+ aliasing an input): WebGPU forbids binding one buffer as both read-only and
61+ read-write in a bind group.
62+- **`toc` as a bare statement lowers to a Call named `toc_print`**, not `toc`.
63+ String literals carry a `Char` type kind, not Numeric.
64+- **Loop bodies may not contain host ops** (tic/toc/disp/fprintf/echoes) —
65+ they're compiled once and replayed; the planner declines with a message.
66+- **rand inside loops**: generator kernels bind every enclosing loop's
67+ iteration counter and mix it into the hash, or every iteration would draw
68+ identical values. Don't "simplify" that away.
69+- numbl HEAD is pinned in both workflows (`NUMBL_REF`); the compiler surface
70+ this repo relies on is declared in src/mgpu/numbl.d.ts, so a numbl change
71+ breaks the build here with a type diff. CI builds numbl's dist-browser
72+ (`npm run build:browser`) because it is not committed.
73+
74+## Testing
75+
76+- `npm run test:node` — 17-case correctness suite on desktop Dawn (real GPU
77+ on this machine). References computed in f64 in the test file; tolerances
78+ are f32-scale, loosened to 3e-4 where SwiftShader's transcendentals lag.
79+- `npm run test:gpu` — same suite, headless Chrome.
80+- `npx vite-node scripts/bench-probe.ts` — timing sanity (fused loop ~1
81+ kernel/iteration, GEMM GFLOP/s printout).
README.mdadded+100−0View file
@@ -0,0 +1,100 @@
1+# math-webgpu-sandbox
2+
3+Write a MATLAB script, run it on your GPU. The sandbox compiles MATLAB-syntax
4+scripts to fused WebGPU compute kernels, times them with the script's own
5+`tic`/`toc`, and — because everything you write is plain MATLAB — you can
6+paste the same script into real MATLAB and compare the numbers. An optional
7+in-browser CPU run through [numbl](https://numbl.org)'s normal engine gives a
8+third column without leaving the page.
9+
10+**Live page:** https://concept-collection.github.io/math-webgpu-sandbox/
11+
12+```matlab
13+n = 4000000;
14+x = rand(n, 1);
15+y = zeros(n, 1);
16+tic;
17+for k = 1:200
18+ y = y + 0.1*sin(x + k) .* exp(-x) + x.^2;
19+end
20+toc
21+fprintf('checksum %.4f\n', mean(y));
22+```
23+
24+That loop body is **one** GPU kernel, compiled once and replayed 200 times.
25+
26+## How it works
27+
28+The compiler front end is numbl's own: the script is parsed and lowered by
29+numbl's JIT pipeline (reached through a `numbl-src` vite alias, the same
30+arrangement [turing-surface](https://github.com/concept-collection/turing-surface)
31+uses), which fixes every type and shape at compile time — `n = 2048;
32+A = rand(n);` pins static shapes via exact-value propagation. The back end is
33+this repo's: a planner maps the typed IR onto WebGPU, and an executor replays
34+the resulting op sequence.
35+
36+- **Fusion.** numbl's inline pass folds single-use temps back into their
37+ consumer, and a sandbox-side pass (`src/mgpu/fuse.ts`) folds the rest —
38+ transcendentals, comparisons, logicals — which numbl's C backend declines
39+ but WGSL handles. One source line becomes one kernel. Generators fuse too:
40+ `x = 2*rand(n,1) - 1` is a single kernel that hashes its way to uniform
41+ variates per element, touching no other buffer.
42+- **Reductions** (`sum`/`mean`/`prod`/`max`/`min`/`norm`/`dot`) take a fused
43+ *loader*, so `sum(a.*b + c)` reads its operands exactly once. Vectors (and
44+ `X(:)`) reduce fully in two passes; matrices reduce per column, which in
45+ MATLAB's column-major layout is the contiguous direction.
46+- **`A * B`** is a 16×16-tiled shared-memory GEMM (column-major, adapted from
47+ [matmul-bench](https://github.com/concept-collection/matmul-bench));
48+ matrix transpose is a tiled relayout; vector transpose, `X(:)` and
49+ `reshape` are views of the same buffer (free) when the source is never
50+ reassigned.
51+- **`for` loops replay, they don't unroll.** The body is planned once; the
52+ loop variable lives in a dynamic-offset uniform with one slot per
53+ iteration, and all iterations are encoded into one submit. `rand` inside a
54+ loop mixes the iteration counter into its stream, so every pass draws
55+ fresh values.
56+- **`tic`/`toc` are synchronization points**: pending GPU work is submitted
57+ and awaited, then wall-clock time is taken on the host. That is the same
58+ thing MATLAB's synchronous tic/toc measures, which is what makes the
59+ number comparable when you paste the script there.
60+
61+## What is (and isn't) supported
62+
63+Supported: elementwise math on real arrays (including comparisons, `&`/`|`/`~`,
64+two-arg `max`/`min`, `mod`/`rem`, `.^`), `zeros`/`ones`/`eye`/`rand`/`randn`/
65+`linspace`/ranges, `A*B`, transpose, the reductions above, `A(:)`/`reshape`,
66+counted `for` loops, `tic`/`toc`, `disp`/`fprintf`, semicolon-suppressed or
67+echoed assignments, and literal row/matrix constants.
68+
69+Deliberately rejected, with a source-located error rather than a wrong
70+answer: indexing/slicing beyond `(:)`, `if`/`while`, complex numbers,
71+user-defined functions, variables that change size, and printing inside `for`
72+loops (the body has nowhere to run host I/O — it replays on the GPU).
73+
74+## The comparison is honest, with two caveats
75+
76+- The GPU computes in **f32**; WebGPU has no f64. MATLAB defaults to double.
77+ Timing comparisons are still meaningful; for the closest apples-to-apples,
78+ use `single` arrays in MATLAB. Values agree to single precision, and the
79+ scripts print checksums so you can see that they do.
80+- `rand`/`randn` here are deterministic counter-based generators
81+ (PCG-flavored hash of element index and call site). Statistics match
82+ MATLAB's; individual draws do not, so checksums of random data are close
83+ but not equal.
84+
85+## Development
86+
87+```
88+npm install # numbl must be checked out as a sibling: ../../numbl
89+npm run dev
90+npm run test:node # correctness suite on desktop WebGPU (Dawn)
91+npm run test:gpu # same suite in headless Chrome (SwiftShader fallback)
92+```
93+
94+The CPU-comparison button needs numbl's `dist-browser` build
95+(`npm run build:browser` in the numbl checkout); the GPU path runs from
96+numbl's TypeScript sources directly and needs no numbl build.
97+
98+## License
99+
100+Apache-2.0
examples/fused_loop.madded+13−0View file
@@ -0,0 +1,13 @@
1+% A long chain of elementwise math on 4 million elements, fused into ONE
2+% GPU kernel per source line and replayed 200 times.
3+% Paste this whole script into MATLAB to compare timings.
4+n = 4000000;
5+x = rand(n, 1);
6+y = zeros(n, 1);
7+tic;
8+for k = 1:200
9+ y = y + 0.1*sin(x + k) .* exp(-x) + x.^2;
10+end
11+t = toc;
12+fprintf('%.0f million element-updates/sec\n', 200*n/1e6/t);
13+fprintf('checksum %.4f (rand differs from MATLAB; expect ~equal, not equal)\n', mean(y));
examples/logistic_ensemble.madded+11−0View file
@@ -0,0 +1,11 @@
1+% One million logistic maps x <- r x (1 - x) iterated in lockstep.
2+% The loop body compiles once; 500 iterations replay on the GPU.
3+n = 1000000;
4+r = 3.5 + 0.5*rand(n, 1);
5+x = rand(n, 1);
6+tic;
7+for k = 1:500
8+ x = r .* x .* (1 - x);
9+end
10+toc
11+fprintf('mean x after 500 iterations: %.4f\n', mean(x));
examples/matmul.madded+14−0View file
@@ -0,0 +1,14 @@
1+% Tiled f32 matrix multiply (the GPU has no f64).
2+% For an apples-to-apples MATLAB run, also try single precision there:
3+% A = rand(n, n, 'single'); B = rand(n, n, 'single');
4+n = 1024;
5+A = rand(n, n);
6+B = rand(n, n);
7+C = A * B; % warm-up
8+tic;
9+for k = 1:20
10+ C = A * B;
11+end
12+t = toc;
13+fprintf('%.1f GFLOP/s over 20 multiplies of %dx%d\n', 2*n^3*20/1e9/t, n, n);
14+fprintf('checksum %.2f\n', sum(C(:))/n^2);
examples/monte_carlo_pi.madded+9−0View file
@@ -0,0 +1,9 @@
1+% Monte Carlo estimate of pi: comparisons and logicals fuse too, so the
2+% mask never touches memory -- one fused pass plus the reduction.
3+n = 10000000;
4+tic;
5+x = 2*rand(n, 1) - 1;
6+y = 2*rand(n, 1) - 1;
7+pi_est = 4*mean((x.^2 + y.^2) <= 1);
8+toc
9+fprintf('pi ~ %.6f (n = %d)\n', pi_est, n);
examples/reductions.madded+11−0View file
@@ -0,0 +1,11 @@
1+% Reductions take a fused loader: dot, norm and max each read their
2+% operands exactly once -- sum(a.*b) is one pass, not two.
3+n = 10000000;
4+a = rand(n, 1);
5+b = rand(n, 1);
6+tic;
7+d = dot(a, b);
8+nrm = norm(a - b);
9+mx = max(abs(a - b));
10+toc
11+fprintf('dot %.0f norm %.2f max %.6f\n', d, nrm, mx);
index.htmladded+198−0View file
@@ -0,0 +1,198 @@
1+<!doctype html>
2+<html lang="en">
3+ <head>
4+ <meta charset="utf-8" />
5+ <meta name="viewport" content="width=device-width, initial-scale=1" />
6+ <link rel="icon" href="data:image/svg+xml,<svg xmlns=%22http://www.w3.org/2000/svg%22 viewBox=%220 0 100 100%22><rect x=%225%22 y=%225%22 width=%2290%22 height=%2290%22 rx=%2216%22 fill=%22%230969da%22/><text x=%2250%22 y=%2268%22 font-size=%2252%22 text-anchor=%22middle%22 fill=%22white%22 font-family=%22monospace%22 font-weight=%22bold%22>⚡</text></svg>" />
7+ <title>math-webgpu-sandbox — MATLAB-syntax scripts on WebGPU</title>
8+ <style>
9+ :root {
10+ --bg: #ffffff;
11+ --ink: #1f2328;
12+ --ink-2: #57606a;
13+ --line: #d0d7de;
14+ --accent: #0969da;
15+ --pane-bg: #f4f6f8;
16+ --tok-com: #6e7781;
17+ --tok-str: #0a3069;
18+ --tok-num: #0550ae;
19+ --tok-kw: #cf222e;
20+ --tok-ext: #8250df;
21+ --warn-bg: #fff8e5;
22+ --warn-line: #e3c37a;
23+ --warn-edge: #bf8700;
24+ --err: #b35900;
25+ color-scheme: light dark;
26+ }
27+ @media (prefers-color-scheme: dark) {
28+ :root {
29+ --bg: #14171a;
30+ --ink: #e6e9ec;
31+ --ink-2: #9aa4af;
32+ --line: #333b44;
33+ --accent: #58a6ff;
34+ --pane-bg: #191d22;
35+ --tok-com: #8b949e;
36+ --tok-str: #a5d6ff;
37+ --tok-num: #79c0ff;
38+ --tok-kw: #ff7b72;
39+ --tok-ext: #d2a8ff;
40+ --warn-bg: #2b2410;
41+ --warn-line: #6b5518;
42+ --warn-edge: #e3b341;
43+ --err: #f0883e;
44+ }
45+ }
46+ body {
47+ margin: 0;
48+ background: var(--bg);
49+ color: var(--ink);
50+ font: 15px/1.5 system-ui, -apple-system, sans-serif;
51+ }
52+ main { max-width: 1200px; margin: 0 auto; padding: 20px 16px 48px; }
53+ h1 { font-size: 20px; margin: 0 0 2px; }
54+ .sub { color: var(--ink-2); margin: 0 0 10px; font-size: 13px; }
55+ .sub a { color: var(--accent); }
56+ .warn {
57+ margin: 0 0 12px; padding: 10px 14px;
58+ border: 1px solid var(--warn-line); border-left: 5px solid var(--warn-edge);
59+ border-radius: 6px; background: var(--warn-bg);
60+ font-size: 13.5px; line-height: 1.5;
61+ }
62+ .controls {
63+ display: flex; flex-wrap: wrap; gap: 8px 14px; align-items: center;
64+ padding: 6px 0;
65+ }
66+ .controls label { color: var(--ink-2); font-size: 13px; white-space: nowrap; }
67+ select, button {
68+ font: inherit; font-size: 13px;
69+ color: var(--ink); background: var(--bg);
70+ border: 1px solid var(--line); border-radius: 6px;
71+ padding: 4px 10px;
72+ }
73+ button { cursor: pointer; }
74+ button:hover:not(:disabled) { border-color: var(--accent); }
75+ button:disabled { opacity: 0.5; cursor: default; }
76+ button.primary { border-color: var(--accent); color: var(--accent); font-weight: 600; }
77+ #device { font-size: 12.5px; color: var(--ink-2); margin-left: auto; }
78+ .split { display: flex; gap: 14px; margin-top: 10px; align-items: stretch; }
79+ .box {
80+ border: 1px solid var(--line); border-radius: 8px; overflow: hidden;
81+ display: flex; flex-direction: column; min-width: 0;
82+ }
83+ #editorbox { flex: 1 1 56%; }
84+ #outbox { flex: 1 1 44%; }
85+ .box-head {
86+ display: flex; gap: 10px; align-items: center; justify-content: space-between;
87+ padding: 6px 10px; font-size: 12px; color: var(--ink-2);
88+ background: var(--pane-bg); border-bottom: 1px solid var(--line);
89+ }
90+ .editor-code { position: relative; flex: 1; min-height: 30em; }
91+ .editor-code > pre,
92+ .editor-code > textarea {
93+ margin: 0; padding: 10px 12px; border: 0;
94+ box-sizing: border-box; width: 100%; height: 100%;
95+ font: 12.5px/1.6 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
96+ tab-size: 2;
97+ white-space: pre; overflow-wrap: normal;
98+ }
99+ #highlight {
100+ position: absolute; inset: 0; overflow: hidden;
101+ pointer-events: none; background: var(--bg); color: var(--ink);
102+ }
103+ #source {
104+ position: relative; z-index: 1; display: block;
105+ resize: none; overflow: auto;
106+ background: transparent; color: transparent; caret-color: var(--ink);
107+ }
108+ #source:focus { outline: none; }
109+ #source::selection { background: color-mix(in srgb, var(--accent) 28%, transparent); }
110+ .console {
111+ flex: 1; margin: 0; padding: 10px 12px; overflow: auto;
112+ font: 12.5px/1.6 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
113+ white-space: pre-wrap; background: var(--pane-bg);
114+ min-height: 8em;
115+ }
116+ .console .err { color: var(--err); }
117+ .console .meta { color: var(--ink-2); }
118+ #cpuconsole { border-top: 1px solid var(--line); max-height: 14em; display: none; }
119+ #timings {
120+ margin-top: 12px; font-size: 13.5px; font-variant-numeric: tabular-nums;
121+ border-collapse: collapse; display: none;
122+ }
123+ #timings td, #timings th {
124+ border: 1px solid var(--line); padding: 4px 12px; text-align: right;
125+ }
126+ #timings th { background: var(--pane-bg); font-weight: 600; }
127+ #timings td:first-child, #timings th:first-child { text-align: left; }
128+ details { margin-top: 12px; font-size: 13px; }
129+ details pre {
130+ margin: 6px 0 0; padding: 8px 10px; overflow: auto;
131+ background: var(--pane-bg); border: 1px solid var(--line); border-radius: 6px;
132+ font: 12px/1.6 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
133+ }
134+ #blurb { margin-top: 14px; font-size: 13px; color: var(--ink-2); max-width: 72em; }
135+ #blurb code { font-size: 12px; }
136+ .tok-com { color: var(--tok-com); }
137+ .tok-str { color: var(--tok-str); }
138+ .tok-num { color: var(--tok-num); }
139+ .tok-kw { color: var(--tok-kw); font-weight: 600; }
140+ .tok-ext { color: var(--tok-ext); }
141+ @media (max-width: 900px) {
142+ .split { flex-direction: column; }
143+ .editor-code { min-height: 22em; }
144+ }
145+ </style>
146+ </head>
147+ <body>
148+ <main>
149+ <h1>math-webgpu-sandbox</h1>
150+ <p class="sub">
151+ MATLAB-syntax scripts compiled to fused WebGPU compute kernels — time them with
152+ <code>tic</code>/<code>toc</code>, then paste the same script into MATLAB and compare.
153+ <a href="https://github.com/concept-collection/math-webgpu-sandbox">source</a>
154+ </p>
155+ <div id="gpuwarn" class="warn" hidden></div>
156+ <div class="controls">
157+ <label>example
158+ <select id="example"></select>
159+ </label>
160+ <button id="run" class="primary">Run on GPU</button>
161+ <button id="runcpu" title="Run the same script through numbl's CPU engine in a worker">Run on CPU (numbl)</button>
162+ <button id="copy" title="Copy the script for pasting into MATLAB">Copy script</button>
163+ <span id="device"></span>
164+ </div>
165+ <div class="split">
166+ <div class="box" id="editorbox">
167+ <div class="box-head"><span>script.m — valid MATLAB</span><span id="editstate"></span></div>
168+ <div class="editor-code">
169+ <pre id="highlight" aria-hidden="true"></pre>
170+ <textarea id="source" spellcheck="false" autocomplete="off"></textarea>
171+ </div>
172+ </div>
173+ <div class="box" id="outbox">
174+ <div class="box-head"><span>output</span><span id="runstate"></span></div>
175+ <pre id="console" class="console"></pre>
176+ <pre id="cpuconsole" class="console"></pre>
177+ </div>
178+ </div>
179+ <table id="timings"></table>
180+ <details id="plandetails" hidden>
181+ <summary>compiled GPU plan</summary>
182+ <pre id="plan"></pre>
183+ </details>
184+ <p id="blurb">
185+ The GPU computes in <b>f32</b> (WebGPU has no f64) while MATLAB defaults to double —
186+ timings compare fairly, values agree to single precision. For the closest MATLAB
187+ comparison use <code>single</code> arrays. <code>rand</code> here is a deterministic
188+ counter-based generator: statistics match MATLAB's, individual draws do not.
189+ Supported: elementwise math (fused per source line, including comparisons and
190+ logicals), <code>A*B</code>, transpose, <code>sum/mean/prod/max/min/norm/dot</code>,
191+ <code>A(:)</code>/<code>reshape</code> (free), <code>for</code> loops (replayed, not
192+ unrolled), <code>tic/toc/disp/fprintf</code>. Not (yet) supported: indexing/slicing
193+ beyond <code>(:)</code>, <code>if/while</code>, complex numbers, user functions.
194+ </p>
195+ </main>
196+ <script type="module" src="/src/main.ts"></script>
197+ </body>
198+</html>
package-lock.jsonadded+3475−0View file
This diff is 3,480 lines long and is not shown.
package.jsonadded+31−0View file
@@ -0,0 +1,31 @@
1+{
2+ "name": "math-webgpu-sandbox",
3+ "version": "0.1.0",
4+ "description": "MATLAB-syntax playground that compiles scripts to fused WebGPU compute kernels, with tic/toc timing you can compare against real MATLAB",
5+ "type": "module",
6+ "engines": {
7+ "node": ">=22.6"
8+ },
9+ "license": "Apache-2.0",
10+ "scripts": {
11+ "dev": "vite",
12+ "build": "tsc --noEmit && vite build",
13+ "test:node": "vite-node scripts/test-node.ts",
14+ "test:gpu": "vite build && node scripts/test-gpu.mjs",
15+ "test": "npm run test:node && npm run test:gpu"
16+ },
17+ "dependencies": {
18+ "numbl": "file:../../numbl"
19+ },
20+ "optionalDependencies": {
21+ "webgpu": "^0.4.0"
22+ },
23+ "devDependencies": {
24+ "@types/node": "^26.1.1",
25+ "@webgpu/types": "^0.1.44",
26+ "puppeteer-core": "^23.0.0",
27+ "typescript": "^5.5.0",
28+ "vite": "^5.4.0",
29+ "vite-node": "^6.0.0"
30+ }
31+}
scripts/bench-probe.tsadded+58−0View file
@@ -0,0 +1,58 @@
1+/* Quick timing sanity: does the sandbox produce sane benchmark numbers? */
2+import { installWebGpu } from './nodeWebGpu.ts';
3+import { runScript } from '../src/mgpu/session.ts';
4+
5+const script = `
6+n = 4000000;
7+x = rand(n, 1);
8+y = zeros(n, 1);
9+tic;
10+for k = 1:100
11+ y = y + 0.1*sin(x + k) .* exp(-x) + x.^2;
12+end
13+toc
14+fprintf('checksum %.4f\\n', mean(y));
15+
16+m = 1024;
17+A = rand(m, m);
18+B = rand(m, m);
19+C = A * B;
20+tic;
21+for k = 1:10
22+ C = A * B;
23+end
24+toc
25+fprintf('gemm checksum %.2f\\n', sum(C(:)) / m);
26+
27+tic;
28+s = 0;
29+for k = 1:100
30+ s = s + sum(x);
31+end
32+toc
33+fprintf('reduce checksum %.1f\\n', s / 100);
34+`;
35+
36+async function main(): Promise<void> {
37+ console.log(await installWebGpu());
38+ const adapter = await navigator.gpu.requestAdapter();
39+ const device = await adapter!.requestDevice({
40+ requiredLimits: {
41+ maxStorageBufferBindingSize: adapter!.limits.maxStorageBufferBindingSize,
42+ maxBufferSize: adapter!.limits.maxBufferSize,
43+ },
44+ });
45+ const run = await runScript(device, script);
46+ console.log('--- output ---');
47+ console.log(run.result.output);
48+ console.log('--- plan ---');
49+ console.log(run.planDescription.join('\n'));
50+ console.log(`compile: ${(run.compileSeconds * 1000).toFixed(1)} ms`);
51+ console.log(`total: ${(run.result.totalSeconds * 1000).toFixed(1)} ms`);
52+ if (run.result.error) console.log('ERROR:', run.result.error);
53+ const flops = (2 * 1024 ** 3 * 10) / 1e9;
54+ const gemmSeg = run.result.segments[1];
55+ console.log(`gemm: ${(flops / gemmSeg.seconds).toFixed(1)} GFLOP/s over 10 reps`);
56+ process.exit(0);
57+}
58+main().catch((e) => { console.error(e); process.exit(1); });
scripts/nodeWebGpu.tsadded+71−0View file
@@ -0,0 +1,71 @@
1+/**
2+ * Desktop WebGPU for the command-line scripts, via the optional `webgpu`
3+ * package (prebuilt Google Dawn).
4+ *
5+ * Installs Dawn under the globals the transform code expects (navigator.gpu,
6+ * GPUBufferUsage, ...) so everything under src/ runs here unchanged —
7+ * including requestShtDevice(), which makes the same device request the
8+ * browser makes.
9+ */
10+
11+export const errMsg = (e: unknown): string =>
12+ e instanceof Error ? e.message : String(e);
13+
14+/**
15+ * Returns a human-readable runtime description. The import specifier is
16+ * indirect so typechecking does not require the optional package.
17+ */
18+export async function installWebGpu(): Promise<string> {
19+ const specifier = 'webgpu';
20+ let mod: {
21+ create: (flags: string[]) => GPU;
22+ globals: Record<string, unknown>;
23+ };
24+ try {
25+ mod = await import(specifier);
26+ } catch (e) {
27+ // Distinguish "not installed" from "installed but the prebuilt Dawn binary
28+ // will not load" — the second is what a machine missing a system library
29+ // looks like, and reporting it as the first sends people in circles.
30+ const detail = errMsg(e);
31+ if (/Cannot find (package|module) '?webgpu'?/.test(detail)) {
32+ throw new Error(
33+ 'desktop WebGPU needs the optional `webgpu` package (prebuilt Google Dawn):\n' +
34+ ' npm install webgpu\n' +
35+ 'It is an optionalDependency, so npm can skip it silently — `npm ls webgpu`\n' +
36+ 'says whether it is there.',
37+ );
38+ }
39+ const glibc = /GLIBC_([0-9.]+)/.exec(detail);
40+ throw new Error(
41+ `the \`webgpu\` package is installed but did not load:\n ${detail}\n` +
42+ (glibc
43+ ? `Dawn's prebuilt binary wants glibc ${glibc[1]} or newer and this host is older\n` +
44+ '(`ldd --version` says how old). No flag bridges that — use a container with a\n' +
45+ 'newer base image, or a newer host.\n'
46+ : 'That is usually the prebuilt Dawn binary missing a system library.\n'),
47+ );
48+ }
49+ Object.assign(globalThis, mod.globals);
50+ // DAWN_FLAGS is ';'-separated because individual Dawn options take
51+ // comma-separated lists, e.g. 'enable-dawn-features=allow_unsafe_apis,...'
52+ const dawnFlags = process.env.DAWN_FLAGS?.split(';').filter(Boolean) ?? [];
53+ Object.defineProperty(globalThis, 'navigator', {
54+ value: { gpu: mod.create(dawnFlags) },
55+ configurable: true,
56+ writable: true,
57+ });
58+ const { version } = await import(`${specifier}/package.json`, {
59+ with: { type: 'json' },
60+ }).then(
61+ (m) => m.default as { version: string },
62+ () => ({ version: '?' }),
63+ );
64+ return `node-webgpu ${version} (Google Dawn)`;
65+}
66+
67+/** The hint to print when Dawn loads but finds no adapter. */
68+export const NO_ADAPTER_HINT =
69+ ' Dawn reaches the GPU through Vulkan on Linux and Windows, Metal on macOS,\n' +
70+ " so a headless box may have no adapter at all. DAWN_FLAGS='backend=vulkan'\n" +
71+ ' makes it explain itself.';
scripts/test-gpu.mjsadded+69−0View file
@@ -0,0 +1,69 @@
1+/**
2+ * Headless GPU test runner: serves dist/, opens test.html in headless
3+ * Chrome (falling back to the SwiftShader software WebGPU adapter when no
4+ * hardware GPU is available), and reports the suite results.
5+ *
6+ * Run after `vite build`: node scripts/test-gpu.mjs
7+ */
8+import { createServer } from 'node:http';
9+import { readFile } from 'node:fs/promises';
10+import { extname, join } from 'node:path';
11+import puppeteer from 'puppeteer-core';
12+
13+const DIST = new URL('../dist/', import.meta.url).pathname;
14+const CHROME = process.env.CHROME_PATH ?? '/usr/bin/google-chrome';
15+const MIME = {
16+ '.html': 'text/html',
17+ '.js': 'text/javascript',
18+ '.css': 'text/css',
19+ '.json': 'application/json',
20+ '.wasm': 'application/wasm',
21+};
22+
23+const server = createServer(async (req, res) => {
24+ try {
25+ const path = req.url === '/' ? '/index.html' : req.url.split('?')[0];
26+ const data = await readFile(join(DIST, path));
27+ res.writeHead(200, { 'content-type': MIME[extname(path)] ?? 'application/octet-stream' });
28+ res.end(data);
29+ } catch {
30+ res.writeHead(404);
31+ res.end('not found');
32+ }
33+});
34+await new Promise((r) => server.listen(0, '127.0.0.1', r));
35+const port = server.address().port;
36+
37+const flagSets = [
38+ // hardware first, then SwiftShader (software) WebGPU
39+ ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu', '--enable-features=Vulkan'],
40+ ['--headless=new', '--no-sandbox', '--enable-unsafe-webgpu', '--use-webgpu-adapter=swiftshader', '--enable-unsafe-swiftshader'],
41+];
42+
43+let final = null;
44+for (const flags of flagSets) {
45+ const browser = await puppeteer.launch({ executablePath: CHROME, args: flags });
46+ try {
47+ const page = await browser.newPage();
48+ page.on('console', (msg) => console.log(` [page] ${msg.text()}`));
49+ page.on('pageerror', (err) => console.log(` [pageerror] ${err.message}`));
50+ await page.goto(`http://127.0.0.1:${port}/test.html`, { waitUntil: 'load' });
51+ const results = await page.waitForFunction(() => window.__RESULTS__, { timeout: 600_000 });
52+ final = await results.jsonValue();
53+ } catch (e) {
54+ console.error(`run with flags [${flags.join(' ')}] failed: ${e.message}`);
55+ } finally {
56+ await browser.close();
57+ }
58+ if (final && !final.fatal) break;
59+ console.log('retrying with next flag set…');
60+}
61+server.close();
62+
63+if (!final || final.fatal) {
64+ console.error(`GPU tests could not run: ${final?.fatal ?? 'no results'}`);
65+ process.exit(2);
66+}
67+if (!final.ok) for (const line of final.lines ?? []) console.log(` ${line}`);
68+console.log(final.ok ? 'GPU SUITE: PASS' : 'GPU SUITE: FAIL');
69+process.exit(final.ok ? 0 : 1);
scripts/test-node.tsadded+39−0View file
@@ -0,0 +1,39 @@
1+/**
2+ * The shared correctness suite on desktop WebGPU (Dawn). Run: npm run test:node
3+ *
4+ * `--skip-without-gpu` exits 0 when no adapter exists (CI runners without a
5+ * GPU for Node; the browser suite still covers everything on SwiftShader).
6+ */
7+import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
8+import { runCases } from '../test/cases.ts';
9+
10+async function main(): Promise<void> {
11+ const skipWithoutGpu = process.argv.includes('--skip-without-gpu');
12+ try {
13+ console.log(await installWebGpu());
14+ } catch (e) {
15+ console.error(errMsg(e));
16+ process.exit(skipWithoutGpu ? 0 : 1);
17+ }
18+ const adapter = await navigator.gpu.requestAdapter();
19+ if (!adapter) {
20+ console.error('no WebGPU adapter\n' + NO_ADAPTER_HINT);
21+ process.exit(skipWithoutGpu ? 0 : 1);
22+ }
23+ const device = await adapter.requestDevice({
24+ requiredLimits: {
25+ maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize,
26+ maxBufferSize: adapter.limits.maxBufferSize,
27+ maxStorageBuffersPerShaderStage: adapter.limits.maxStorageBuffersPerShaderStage,
28+ },
29+ });
30+
31+ const failures = await runCases(device, (line) => console.log(line));
32+ console.log(failures ? `\n${failures} failure(s)` : '\nall tests passed');
33+ process.exit(failures ? 1 : 0);
34+}
35+
36+main().catch((e) => {
37+ console.error(e);
38+ process.exit(1);
39+});
src/cpu/cpuRunner.tsadded+76−0View file
@@ -0,0 +1,76 @@
1+/**
2+ * The optional CPU comparison: run the very same script through numbl's
3+ * normal engine (interpreter + its own JIT) in a worker numbl manages.
4+ *
5+ * tic/toc, fprintf and disp behave natively there, so the output — including
6+ * the "Elapsed time is ..." lines — is directly comparable to the GPU pane
7+ * and to real MATLAB. numbl computes in f64, which is also a useful
8+ * cross-check on the GPU's f32 results.
9+ *
10+ * A fresh session per run keeps workspace state from leaking between runs.
11+ * mip is disabled: sandbox scripts are self-contained MATLAB, and skipping
12+ * the bootstrap keeps first-run latency down.
13+ */
14+import { createNumblSession } from 'numbl/browser';
15+
16+export interface CpuRunResult {
17+ output: string;
18+ /** Wall time of the whole script, as seen from the host. */
19+ totalSeconds: number;
20+ error?: string;
21+ aborted?: boolean;
22+}
23+
24+export interface CpuRunHandle {
25+ result: Promise<CpuRunResult>;
26+ /** Cooperative stop (needs cross-origin isolation to preempt loops). */
27+ cancel: () => void;
28+}
29+
30+export function runOnCpu(
31+ source: string,
32+ onOutput?: (text: string) => void,
33+): CpuRunHandle {
34+ let disposed = false;
35+ let sessionRef: { dispose(): void; interrupt(): void } | null = null;
36+
37+ const result = (async (): Promise<CpuRunResult> => {
38+ let output = '';
39+ const session = await createNumblSession({
40+ mip: false,
41+ persistSystem: false,
42+ displayResults: true,
43+ onOutput: (text: string) => {
44+ output += text;
45+ onOutput?.(text);
46+ },
47+ });
48+ sessionRef = session;
49+ if (disposed) {
50+ session.dispose();
51+ return { output: '', totalSeconds: 0, aborted: true };
52+ }
53+ const t0 = performance.now();
54+ try {
55+ const res = await session.execute(source);
56+ const totalSeconds = (performance.now() - t0) / 1000;
57+ return {
58+ output,
59+ totalSeconds,
60+ error: res.ok ? undefined : (res.error ?? 'numbl error'),
61+ aborted: res.aborted === true,
62+ };
63+ } finally {
64+ session.dispose();
65+ }
66+ })();
67+
68+ return {
69+ result,
70+ cancel: () => {
71+ disposed = true;
72+ sessionRef?.interrupt();
73+ sessionRef?.dispose();
74+ },
75+ };
76+}
src/editor/codeEditor.tsadded+92−0View file
@@ -0,0 +1,92 @@
1+/**
2+ * A textarea with syntax highlighting, by overlay.
3+ *
4+ * A textarea cannot colour its own text, so the highlighted source is rendered
5+ * into a <pre> underneath and the textarea sits on top with transparent text and
6+ * a visible caret. The two must agree on every metric that affects layout —
7+ * font, line height, padding, tab size, wrapping — and their scroll offsets are
8+ * kept in sync, or the colours drift away from the characters.
9+ */
10+import { highlightMatlab } from './matlab.ts';
11+
12+export interface CodeEditorOptions {
13+ textarea: HTMLTextAreaElement;
14+ /** The <pre> behind it, holding the highlighted copy. */
15+ overlay: HTMLElement;
16+ /** Names to mark as host-provided operations. */
17+ external?: ReadonlySet<string>;
18+ /** Called on every edit. */
19+ onInput?: (value: string) => void;
20+}
21+
22+export class CodeEditor {
23+ #textarea: HTMLTextAreaElement;
24+ #overlay: HTMLElement;
25+ #external: ReadonlySet<string>;
26+
27+ constructor(opts: CodeEditorOptions) {
28+ this.#textarea = opts.textarea;
29+ this.#overlay = opts.overlay;
30+ this.#external = opts.external ?? new Set();
31+
32+ this.#textarea.addEventListener('input', () => {
33+ this.#repaint();
34+ opts.onInput?.(this.#textarea.value);
35+ });
36+ // Keep the colours under the characters while scrolling.
37+ this.#textarea.addEventListener('scroll', () => this.#syncScroll());
38+ // Tab should indent rather than leave the editor.
39+ this.#textarea.addEventListener('keydown', (e) => this.#onKeyDown(e));
40+ this.#repaint();
41+ }
42+
43+ get value(): string {
44+ return this.#textarea.value;
45+ }
46+
47+ set value(next: string) {
48+ this.#textarea.value = next;
49+ this.#repaint();
50+ }
51+
52+ focus(): void {
53+ this.#textarea.focus();
54+ }
55+
56+ /** Select a character range, scrolling it into view. */
57+ select(start: number, end: number): void {
58+ this.#textarea.focus();
59+ this.#textarea.setSelectionRange(start, end);
60+ // setSelectionRange does not always scroll; nudge the line into view.
61+ const line = this.#textarea.value.slice(0, start).split('\n').length - 1;
62+ const lineHeight = this.#textarea.scrollHeight / Math.max(1, this.#lineCount());
63+ const target = line * lineHeight - this.#textarea.clientHeight / 2;
64+ this.#textarea.scrollTop = Math.max(0, target);
65+ this.#syncScroll();
66+ }
67+
68+ #lineCount(): number {
69+ return this.#textarea.value.split('\n').length + 1; // +1 for the trailing line
70+ }
71+
72+ #onKeyDown(e: KeyboardEvent): void {
73+ if (e.key !== 'Tab' || e.ctrlKey || e.metaKey || e.altKey) return;
74+ e.preventDefault();
75+ const el = this.#textarea;
76+ const { selectionStart: s, selectionEnd: t, value } = el;
77+ el.value = `${value.slice(0, s)} ${value.slice(t)}`;
78+ el.selectionStart = el.selectionEnd = s + 2;
79+ // Let the input listener repaint and notify, as for any other edit.
80+ el.dispatchEvent(new Event('input'));
81+ }
82+
83+ #repaint(): void {
84+ this.#overlay.innerHTML = highlightMatlab(this.#textarea.value, this.#external);
85+ this.#syncScroll();
86+ }
87+
88+ #syncScroll(): void {
89+ this.#overlay.scrollTop = this.#textarea.scrollTop;
90+ this.#overlay.scrollLeft = this.#textarea.scrollLeft;
91+ }
92+}
src/editor/matlab.tsadded+213−0View file
@@ -0,0 +1,213 @@
1+/**
2+ * A small MATLAB tokenizer, for syntax highlighting the model editor.
3+ *
4+ * Only what highlighting needs — comments, literals, numbers, keywords — and
5+ * deliberately not a parser: numbl does the real parsing, and reports errors
6+ * with positions. Tokens preserve the source text exactly, character for
7+ * character, because the highlighted output is overlaid on a textarea and any
8+ * dropped or added character would shift the two out of alignment.
9+ */
10+
11+export type TokenClass = 'com' | 'str' | 'num' | 'kw' | 'ext';
12+
13+export interface Token {
14+ text: string;
15+ cls: TokenClass | null;
16+}
17+
18+const KEYWORDS = new Set([
19+ 'break', 'case', 'catch', 'classdef', 'continue', 'else', 'elseif', 'end',
20+ 'for', 'function', 'global', 'if', 'otherwise', 'parfor', 'persistent',
21+ 'return', 'spmd', 'switch', 'try', 'while',
22+]);
23+
24+const isIdentStart = (c: string): boolean => /[A-Za-z_]/.test(c);
25+const isIdent = (c: string): boolean => /[A-Za-z0-9_]/.test(c);
26+const isDigit = (c: string): boolean => c >= '0' && c <= '9';
27+
28+/**
29+ * In MATLAB `'` is both the transpose operator and the char-literal delimiter.
30+ * It opens a literal unless it directly follows something that can be
31+ * transposed — a value, a closing bracket, or another transpose.
32+ */
33+function quoteIsTranspose(src: string, at: number): boolean {
34+ for (let i = at - 1; i >= 0; i--) {
35+ const c = src[i];
36+ if (c === ' ' || c === '\t') continue;
37+ return isIdent(c) || c === ')' || c === ']' || c === '}' || c === '.' || c === "'";
38+ }
39+ return false;
40+}
41+
42+/**
43+ * Tokenize `src`. `external` names (the operations the host provides, e.g.
44+ * `synth` / `analys`) get their own class so the boundary between the model and
45+ * what it is given is visible in the editor.
46+ */
47+export function tokenizeMatlab(
48+ src: string,
49+ external: ReadonlySet<string> = new Set(),
50+): Token[] {
51+ const out: Token[] = [];
52+ const push = (text: string, cls: TokenClass | null): void => {
53+ if (!text) return;
54+ const last = out[out.length - 1];
55+ if (last && last.cls === cls) last.text += text;
56+ else out.push({ text, cls });
57+ };
58+
59+ let i = 0;
60+ let atLineStart = true;
61+ let inBlockComment = false;
62+
63+ while (i < src.length) {
64+ const c = src[i];
65+
66+ // Block comments: `%{` and `%}` each alone on their line.
67+ if (atLineStart) {
68+ const eol = src.indexOf('\n', i);
69+ const lineEnd = eol === -1 ? src.length : eol;
70+ const line = src.slice(i, lineEnd);
71+ const trimmed = line.trim();
72+ if (!inBlockComment && trimmed === '%{') inBlockComment = true;
73+ else if (inBlockComment && trimmed === '%}') {
74+ push(line, 'com');
75+ i = lineEnd;
76+ inBlockComment = false;
77+ atLineStart = false;
78+ continue;
79+ }
80+ if (inBlockComment) {
81+ push(line, 'com');
82+ i = lineEnd;
83+ atLineStart = false;
84+ continue;
85+ }
86+ }
87+
88+ if (c === '\n') {
89+ push(c, null);
90+ i++;
91+ atLineStart = true;
92+ continue;
93+ }
94+ if (c === ' ' || c === '\t') {
95+ push(c, null);
96+ i++;
97+ continue;
98+ }
99+ atLineStart = false;
100+
101+ // Line comment, including MATLAB's `%%` section markers.
102+ if (c === '%') {
103+ const eol = src.indexOf('\n', i);
104+ const end = eol === -1 ? src.length : eol;
105+ push(src.slice(i, end), 'com');
106+ i = end;
107+ continue;
108+ }
109+
110+ // Line continuation is an operator, but any trailing text is a comment.
111+ if (c === '.' && src.startsWith('...', i)) {
112+ const eol = src.indexOf('\n', i);
113+ const end = eol === -1 ? src.length : eol;
114+ push('...', null);
115+ push(src.slice(i + 3, end), 'com');
116+ i = end;
117+ continue;
118+ }
119+
120+ // Char literal (or transpose).
121+ if (c === "'") {
122+ if (quoteIsTranspose(src, i)) {
123+ push("'", null);
124+ i++;
125+ continue;
126+ }
127+ let j = i + 1;
128+ while (j < src.length && src[j] !== '\n') {
129+ if (src[j] === "'") {
130+ if (src[j + 1] === "'") j += 2; // escaped quote
131+ else {
132+ j++;
133+ break;
134+ }
135+ } else j++;
136+ }
137+ push(src.slice(i, j), 'str');
138+ i = j;
139+ continue;
140+ }
141+
142+ // Double-quoted string.
143+ if (c === '"') {
144+ let j = i + 1;
145+ while (j < src.length && src[j] !== '\n') {
146+ if (src[j] === '"') {
147+ if (src[j + 1] === '"') j += 2;
148+ else {
149+ j++;
150+ break;
151+ }
152+ } else j++;
153+ }
154+ push(src.slice(i, j), 'str');
155+ i = j;
156+ continue;
157+ }
158+
159+ // Number: 12, 1.5, .5, 1e-3, 2i
160+ if (isDigit(c) || (c === '.' && isDigit(src[i + 1]))) {
161+ let j = i;
162+ while (j < src.length && isDigit(src[j])) j++;
163+ if (src[j] === '.') {
164+ j++;
165+ while (j < src.length && isDigit(src[j])) j++;
166+ }
167+ if (src[j] === 'e' || src[j] === 'E') {
168+ let k = j + 1;
169+ if (src[k] === '+' || src[k] === '-') k++;
170+ if (isDigit(src[k])) {
171+ k++;
172+ while (k < src.length && isDigit(src[k])) k++;
173+ j = k;
174+ }
175+ }
176+ if (src[j] === 'i' || src[j] === 'j') j++;
177+ push(src.slice(i, j), 'num');
178+ i = j;
179+ continue;
180+ }
181+
182+ // Identifier / keyword / external operation.
183+ if (isIdentStart(c)) {
184+ let j = i;
185+ while (j < src.length && isIdent(src[j])) j++;
186+ const word = src.slice(i, j);
187+ push(word, KEYWORDS.has(word) ? 'kw' : external.has(word) ? 'ext' : null);
188+ i = j;
189+ continue;
190+ }
191+
192+ push(c, null);
193+ i++;
194+ }
195+
196+ return out;
197+}
198+
199+const escapeHtml = (s: string): string =>
200+ s.replace(/&/g, '&').replace(/</g, '<').replace(/>/g, '>');
201+
202+/** Highlighted HTML for `src`, safe to assign to innerHTML. */
203+export function highlightMatlab(
204+ src: string,
205+ external: ReadonlySet<string> = new Set(),
206+): string {
207+ const html = tokenizeMatlab(src, external)
208+ .map((t) => (t.cls ? `<span class="tok-${t.cls}">${escapeHtml(t.text)}</span>` : escapeHtml(t.text)))
209+ .join('');
210+ // A trailing newline keeps the last line's box height stable, so the overlay
211+ // and the textarea scroll to the same extent.
212+ return `${html}\n`;
213+}
src/main.tsadded+220−0View file
@@ -0,0 +1,220 @@
1+/**
2+ * The sandbox page: editor on the left, output on the right, a timing table
3+ * underneath that lines up GPU tic/toc segments with the CPU run's.
4+ */
5+import { CodeEditor } from './editor/codeEditor.ts';
6+import { requestSandboxDevice, runScript, type SandboxGpu } from './mgpu/session.ts';
7+import { formatFailure } from './mgpu/errors.ts';
8+import type { CpuRunHandle } from './cpu/cpuRunner.ts';
9+import fusedLoop from '../examples/fused_loop.m?raw';
10+import matmul from '../examples/matmul.m?raw';
11+import monteCarloPi from '../examples/monte_carlo_pi.m?raw';
12+import logisticEnsemble from '../examples/logistic_ensemble.m?raw';
13+import reductions from '../examples/reductions.m?raw';
14+
15+const EXAMPLES: { key: string; title: string; source: string }[] = [
16+ { key: 'fused_loop', title: 'fused elementwise loop', source: fusedLoop },
17+ { key: 'matmul', title: 'matrix multiply (GEMM)', source: matmul },
18+ { key: 'monte_carlo_pi', title: 'Monte Carlo pi (masks)', source: monteCarloPi },
19+ { key: 'logistic_ensemble', title: 'logistic map ensemble', source: logisticEnsemble },
20+ { key: 'reductions', title: 'fused reductions', source: reductions },
21+];
22+
23+const $ = <T extends HTMLElement>(id: string): T => {
24+ const el = document.getElementById(id);
25+ if (!el) throw new Error(`missing #${id}`);
26+ return el as T;
27+};
28+
29+const consoleEl = $<HTMLPreElement>('console');
30+const cpuConsoleEl = $<HTMLPreElement>('cpuconsole');
31+const runBtn = $<HTMLButtonElement>('run');
32+const runCpuBtn = $<HTMLButtonElement>('runcpu');
33+const copyBtn = $<HTMLButtonElement>('copy');
34+const exampleSel = $<HTMLSelectElement>('example');
35+const timingsEl = $<HTMLTableElement>('timings');
36+const planDetails = $<HTMLElement>('plandetails');
37+const planEl = $<HTMLPreElement>('plan');
38+const runState = $<HTMLElement>('runstate');
39+const deviceEl = $<HTMLElement>('device');
40+const gpuWarn = $<HTMLElement>('gpuwarn');
41+
42+const editor = new CodeEditor({
43+ textarea: $<HTMLTextAreaElement>('source'),
44+ overlay: $('highlight'),
45+ onInput: () => {
46+ $('editstate').textContent = '';
47+ },
48+});
49+
50+for (const ex of EXAMPLES) {
51+ const opt = document.createElement('option');
52+ opt.value = ex.key;
53+ opt.textContent = ex.title;
54+ exampleSel.appendChild(opt);
55+}
56+exampleSel.addEventListener('change', () => {
57+ const ex = EXAMPLES.find((e) => e.key === exampleSel.value);
58+ if (ex) editor.value = ex.source;
59+});
60+editor.value = EXAMPLES[0].source;
61+
62+const append = (el: HTMLElement, text: string, cls?: string): void => {
63+ if (cls) {
64+ const span = document.createElement('span');
65+ span.className = cls;
66+ span.textContent = text;
67+ el.appendChild(span);
68+ } else {
69+ el.appendChild(document.createTextNode(text));
70+ }
71+ el.scrollTop = el.scrollHeight;
72+};
73+
74+/** Latest timing results, merged into the table by tic..toc pair. */
75+const lastSeconds: { gpu: number[]; cpu: number[] } = { gpu: [], cpu: [] };
76+let gpuTotal: number | null = null;
77+let cpuTotal: number | null = null;
78+
79+function renderTimings(): void {
80+ const n = Math.max(lastSeconds.gpu.length, lastSeconds.cpu.length);
81+ if (n === 0 && gpuTotal === null && cpuTotal === null) {
82+ timingsEl.style.display = 'none';
83+ return;
84+ }
85+ const fmt = (s: number | undefined | null): string =>
86+ s === undefined || s === null ? '—' : s >= 0.1 ? `${s.toFixed(3)} s` : `${(s * 1000).toFixed(2)} ms`;
87+ const rows: string[] = [
88+ '<tr><th></th><th>GPU</th><th>CPU (numbl)</th><th>CPU / GPU</th></tr>',
89+ ];
90+ for (let i = 0; i < n; i++) {
91+ const g = lastSeconds.gpu[i];
92+ const c = lastSeconds.cpu[i];
93+ const ratio = g !== undefined && c !== undefined ? `${(c / g).toFixed(1)}×` : '—';
94+ rows.push(
95+ `<tr><td>tic…toc #${i + 1}</td><td>${fmt(g)}</td><td>${fmt(c)}</td><td>${ratio}</td></tr>`,
96+ );
97+ }
98+ const totalRatio =
99+ gpuTotal !== null && cpuTotal !== null ? `${(cpuTotal / gpuTotal).toFixed(1)}×` : '—';
100+ rows.push(
101+ `<tr><td>whole script</td><td>${fmt(gpuTotal)}</td><td>${fmt(cpuTotal)}</td><td>${totalRatio}</td></tr>`,
102+ );
103+ timingsEl.innerHTML = rows.join('');
104+ timingsEl.style.display = 'table';
105+}
106+
107+let gpu: SandboxGpu | null = null;
108+(async () => {
109+ try {
110+ gpu = await requestSandboxDevice();
111+ deviceEl.textContent = `GPU: ${gpu.description}`;
112+ gpu.device.lost.then((info) => {
113+ gpu = null;
114+ gpuWarn.hidden = false;
115+ gpuWarn.textContent = `The GPU device was lost (${info.message}); reload the page.`;
116+ runBtn.disabled = true;
117+ });
118+ } catch (e) {
119+ gpuWarn.hidden = false;
120+ gpuWarn.innerHTML =
121+ `<b>No WebGPU here.</b> ${e instanceof Error ? e.message : String(e)} — ` +
122+ `Chrome/Edge have it on by default; Firefox and Safari are rolling it out. ` +
123+ `The CPU run still works.`;
124+ runBtn.disabled = true;
125+ }
126+})();
127+
128+runBtn.addEventListener('click', () => {
129+ void runGpu();
130+});
131+async function runGpu(): Promise<void> {
132+ if (!gpu) return;
133+ runBtn.disabled = true;
134+ runState.textContent = 'compiling…';
135+ consoleEl.textContent = '';
136+ lastSeconds.gpu = [];
137+ gpuTotal = null;
138+ renderTimings();
139+ planDetails.hidden = true;
140+ const source = editor.value;
141+ try {
142+ const t0 = performance.now();
143+ const run = await runScript(gpu.device, source, (text) => {
144+ runState.textContent = 'running…';
145+ append(consoleEl, text);
146+ });
147+ void t0;
148+ if (run.result.error) {
149+ append(consoleEl, `\n${run.result.error}\n`, 'err');
150+ }
151+ append(
152+ consoleEl,
153+ `\n[compile ${(run.compileSeconds * 1000).toFixed(0)} ms · ` +
154+ `run ${run.result.totalSeconds.toFixed(3)} s · f32 · ${gpu.description}]\n`,
155+ 'meta',
156+ );
157+ lastSeconds.gpu = run.result.segments.map((s) => s.seconds);
158+ gpuTotal = run.result.totalSeconds;
159+ planEl.textContent = run.planDescription.join('\n') || '(no GPU ops)';
160+ planDetails.hidden = false;
161+ renderTimings();
162+ runState.textContent = '';
163+ } catch (e) {
164+ append(consoleEl, formatFailure(e, source) + '\n', 'err');
165+ runState.textContent = '';
166+ } finally {
167+ runBtn.disabled = !gpu;
168+ }
169+}
170+
171+let cpuHandle: CpuRunHandle | null = null;
172+runCpuBtn.addEventListener('click', () => {
173+ if (cpuHandle) {
174+ cpuHandle.cancel();
175+ cpuHandle = null;
176+ runCpuBtn.textContent = 'Run on CPU (numbl)';
177+ return;
178+ }
179+ void runCpu();
180+});
181+async function runCpu(): Promise<void> {
182+ cpuConsoleEl.style.display = 'block';
183+ cpuConsoleEl.textContent = '';
184+ append(cpuConsoleEl, '[CPU · numbl engine in a worker · f64]\n', 'meta');
185+ lastSeconds.cpu = [];
186+ cpuTotal = null;
187+ renderTimings();
188+ runCpuBtn.textContent = 'Stop CPU run';
189+ // numbl's browser engine is a couple of megabytes; load it on first use.
190+ const { runOnCpu } = await import('./cpu/cpuRunner.ts');
191+ const handle = runOnCpu(editor.value, (text) => append(cpuConsoleEl, text));
192+ cpuHandle = handle;
193+ try {
194+ const res = await handle.result;
195+ if (res.error) append(cpuConsoleEl, `\n${res.error}\n`, 'err');
196+ if (res.aborted) append(cpuConsoleEl, `\n[stopped]\n`, 'meta');
197+ else {
198+ append(cpuConsoleEl, `[total ${res.totalSeconds.toFixed(3)} s]\n`, 'meta');
199+ // Line the CPU's tic..toc pairs up with the GPU's by order.
200+ lastSeconds.cpu = [...res.output.matchAll(/Elapsed time is ([0-9.eE+-]+) seconds/g)]
201+ .map((m) => Number(m[1]));
202+ cpuTotal = res.totalSeconds;
203+ renderTimings();
204+ }
205+ } catch (e) {
206+ append(cpuConsoleEl, `${e instanceof Error ? e.message : String(e)}\n`, 'err');
207+ } finally {
208+ if (cpuHandle === handle) cpuHandle = null;
209+ runCpuBtn.textContent = 'Run on CPU (numbl)';
210+ }
211+}
212+
213+copyBtn.addEventListener('click', () => {
214+ void navigator.clipboard.writeText(editor.value).then(() => {
215+ copyBtn.textContent = 'Copied ✓';
216+ setTimeout(() => {
217+ copyBtn.textContent = 'Copy script';
218+ }, 1200);
219+ });
220+});
src/mgpu/compile.tsadded+72−0View file
@@ -0,0 +1,72 @@
1+/**
2+ * MATLAB script source -> numbl's JIT IR, ready for the WGSL planner.
3+ *
4+ * Unlike turing-surface — which specializes named functions against
5+ * host-supplied argument types — the sandbox lowers a whole *script*: shapes
6+ * come from the script itself (`n = 2048; A = rand(n);`), pinned static by
7+ * numbl's exact-value propagation through the type lattice.
8+ *
9+ * Two numbl passes matter here:
10+ * - `lowerProgram` lowers the top-level statements to IR, one statement per
11+ * operation (ANF), with every node's type fixed.
12+ * - `inlinePass` then folds single-use temps back into their consumer, so a
13+ * source line like `y = a - u + u.*u.*v` becomes ONE statement whose RHS is
14+ * an expression tree — i.e. one fused GPU kernel instead of four.
15+ */
16+import { parseMFile } from 'numbl-src/numbl-core/parser/index.ts';
17+import { Workspace, Lowerer } from 'numbl-src/numbl-core/jit/index.ts';
18+import { inlinePass } from 'numbl-src/numbl-core/jit/codegen/inlinePass.ts';
19+import type { IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
20+import { applyBuiltinPatches } from './patches.ts';
21+import { fuseTemps } from './fuse.ts';
22+import { inScript, ScriptCompileError } from './errors.ts';
23+
24+export interface CompiledScript {
25+ /** The lowered, inline-folded top-level statements. */
26+ stmts: IRStmt[];
27+ /**
28+ * Should the statement covering source offset `at` echo its result,
29+ * MATLAB-style? True exactly when the source statement has no trailing
30+ * semicolon. Compiler temps sit inside their source statement's span, so
31+ * the caller must additionally skip `_mtoc2_*` names.
32+ */
33+ isEchoed(at: number): boolean;
34+}
35+
36+export function compileScript(source: string, fileName = 'script.m'): CompiledScript {
37+ applyBuiltinPatches();
38+
39+ const ast = inScript(() => parseMFile(source, fileName));
40+
41+ // Statements the parser marked unsuppressed (no `;`), by source range.
42+ const echoed: { start: number; end: number }[] = [];
43+ for (const s of ast.body) {
44+ if (s.suppressed === false) echoed.push({ start: s.span.start, end: s.span.end });
45+ }
46+
47+ if (ast.body.some((s) => s.type === 'Function')) {
48+ // A script may syntactically end with local functions, but nothing here
49+ // compiles calls to them — say so up front rather than at the call site.
50+ throw new ScriptCompileError(
51+ 'local functions are not supported in the sandbox yet — inline their bodies',
52+ );
53+ }
54+
55+ const prog = inScript(() => {
56+ const ws = new Workspace(fileName, []);
57+ ws.addFile({ name: fileName, source, ast });
58+ ws.finalize();
59+ const lowered = new Lowerer(ws).lowerProgram(ast);
60+ inlinePass(lowered);
61+ // numbl's pass stops at what its C backend fuses; fold the rest of the
62+ // single-use temps the WGSL emitter can absorb (sin/exp, comparisons,
63+ // logicals, generators) so one source line is one kernel.
64+ fuseTemps(lowered.topLevelStmts);
65+ return lowered;
66+ });
67+
68+ return {
69+ stmts: prog.topLevelStmts,
70+ isEchoed: (at) => echoed.some((r) => at >= r.start && at <= r.end),
71+ };
72+}
src/mgpu/errors.tsadded+87−0View file
@@ -0,0 +1,87 @@
1+/**
2+ * Compile failures, reported in coordinates of the script the user edits.
3+ *
4+ * Failures arrive from three places, each with its own idea of position:
5+ * numbl's parser (a `position` offset), numbl's lowerer (`UnsupportedConstruct`
6+ * / `JitTypeError`, with a `span`), and this project's planner and WGSL
7+ * emitter (`UnsupportedOnGpu`, carrying the numbl span it was given). All of
8+ * them are offsets into the script, so they need only be turned into a line
9+ * and column for display.
10+ */
11+
12+/** A compile failure located in the script source. */
13+export class ScriptCompileError extends Error {
14+ /** Offset into the script, when the failure has a position. */
15+ readonly start?: number;
16+ readonly end?: number;
17+
18+ constructor(
19+ message: string,
20+ opts: { start?: number; end?: number; cause?: unknown } = {},
21+ ) {
22+ super(message, { cause: opts.cause });
23+ this.name = 'ScriptCompileError';
24+ this.start = opts.start;
25+ this.end = opts.end;
26+ }
27+}
28+
29+/** Raised for a construct the WGSL backend cannot express. Mirrors numbl's
30+ * own decline discipline: fail at compile time with a source span, never
31+ * silently produce something that computes the wrong thing. */
32+export class UnsupportedOnGpu extends Error {
33+ readonly span?: unknown;
34+ constructor(message: string, span?: unknown) {
35+ super(message);
36+ this.name = 'UnsupportedOnGpu';
37+ this.span = span;
38+ }
39+}
40+
41+/** Extract whatever position information an error carries. */
42+function positionOf(e: unknown): { start?: number; end?: number } {
43+ const span = (e as { span?: { start?: unknown; end?: unknown } }).span;
44+ if (span && typeof span.start === 'number') {
45+ return {
46+ start: span.start,
47+ end: typeof span.end === 'number' ? span.end : undefined,
48+ };
49+ }
50+ // numbl's parser SyntaxError reports a bare offset.
51+ const position = (e as { position?: unknown }).position;
52+ if (typeof position === 'number') return { start: position };
53+ return {};
54+}
55+
56+/** Normalize any thrown value into a located `ScriptCompileError`. */
57+export function asCompileError(e: unknown): ScriptCompileError {
58+ if (e instanceof ScriptCompileError) return e;
59+ const { start, end } = positionOf(e);
60+ const raw = e instanceof Error ? e.message : String(e);
61+ // numbl's parse errors read as bare token complaints out of context.
62+ const message =
63+ (e as Error)?.name === 'SyntaxError' ? `MATLAB syntax error: ${raw}` : raw;
64+ return new ScriptCompileError(message, { start, end, cause: e });
65+}
66+
67+/** Run `fn`, locating any compile failure in the script. */
68+export function inScript<T>(fn: () => T): T {
69+ try {
70+ return fn();
71+ } catch (e) {
72+ throw asCompileError(e);
73+ }
74+}
75+
76+/** Render a failure for display: message and 1-based line/column. */
77+export function formatFailure(e: unknown, source: string): string {
78+ const message = e instanceof Error ? e.message : String(e);
79+ if (!(e instanceof ScriptCompileError)) return message;
80+ if (e.start !== undefined && e.start <= source.length) {
81+ const before = source.slice(0, e.start);
82+ const line = before.split('\n').length;
83+ const column = e.start - before.lastIndexOf('\n');
84+ return `${message} (line ${line}, column ${column})`;
85+ }
86+ return message;
87+}
src/mgpu/fuse.tsadded+163−0View file
@@ -0,0 +1,163 @@
1+/**
2+ * Sandbox-side fusion: fold the single-use ANF temps numbl's inline pass
3+ * left behind.
4+ *
5+ * numbl's pass (`inlinePass`) only folds producers its C backend can fuse —
6+ * no tensor-producing Calls (`sin(x)`), no logical results, no ranges. The
7+ * WGSL emitter fuses all of those, so without this pass a line like
8+ * `y = y + 0.1*sin(x + k) .* exp(-x)` becomes five kernels instead of one.
9+ *
10+ * Same shape as numbl's pass, deliberately narrower where it matters:
11+ * - only compiler temps (`_mtoc2_*`) are folded, so every user-named
12+ * variable still materializes — one source line stays one kernel, and
13+ * anything the user might echo or reuse keeps its buffer;
14+ * - the producer's RHS must be fusable on the GPU (isGpuFusableExpr);
15+ * - the single use must be in a following Assign/ExprStmt at the same body
16+ * level, with no intervening write to the producer's operands and no
17+ * control-flow statement in between.
18+ *
19+ * `for` bodies are processed as their own levels; nothing folds across the
20+ * loop boundary.
21+ */
22+import type { Assign, ExprStmt, IRExpr, IRStmt } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
23+import { isMultiElement, type Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
24+import { isGpuFusableExpr } from './wgsl.ts';
25+
26+const isMulti = (t: Type): boolean => t.kind === 'Numeric' && isMultiElement(t);
27+
28+export function fuseTemps(stmts: IRStmt[]): void {
29+ for (const s of stmts) {
30+ if (s.kind === 'For') fuseTemps(s.body);
31+ }
32+ for (let iter = 0; iter < 32; iter++) {
33+ if (!fuseOnePass(stmts)) break;
34+ }
35+}
36+
37+const isTemp = (cName: string): boolean => cName.startsWith('_mtoc2_');
38+
39+function fuseOnePass(stmts: IRStmt[]): boolean {
40+ const uses = useCounts(stmts);
41+ for (let i = 0; i < stmts.length; i++) {
42+ const p = stmts[i];
43+ if (p.kind !== 'Assign') continue;
44+ if (!isTemp(p.cName)) continue;
45+ if (uses.get(p.cName) !== 1) continue;
46+ if (!isGpuFusableExpr(p.expr)) continue;
47+
48+ const reads = new Set<string>();
49+ walkVars(p.expr, (c) => reads.add(c));
50+
51+ for (let j = i + 1; j < stmts.length; j++) {
52+ const c = stmts[j];
53+ if (c.kind !== 'Assign' && c.kind !== 'ExprStmt') break; // control flow
54+ const holder = c as Assign | ExprStmt;
55+ if (c.kind === 'ExprStmt' && isMulti(p.ty) && countIn(holder.expr, p.cName) === 1) {
56+ break; // disp/fprintf read tensors by name, not as expressions
57+ }
58+ if (countIn(holder.expr, p.cName) === 1) {
59+ holder.expr = substitute(holder.expr, p.cName, p.expr);
60+ stmts.splice(i, 1);
61+ return true;
62+ }
63+ if (c.kind === 'Assign' && (c.cName === p.cName || reads.has(c.cName))) {
64+ break; // intervening write invalidates the fold window
65+ }
66+ }
67+ }
68+ return false;
69+}
70+
71+/** Var-read counts by cName, across this level AND nested bodies (a use
72+ * inside a nested loop must keep the temp alive at this level). */
73+function useCounts(stmts: IRStmt[]): Map<string, number> {
74+ const counts = new Map<string, number>();
75+ const bump = (c: string): void => {
76+ counts.set(c, (counts.get(c) ?? 0) + 1);
77+ };
78+ const walkStmts = (list: IRStmt[]): void => {
79+ for (const s of list) {
80+ if (s.kind === 'Assign' || s.kind === 'ExprStmt') walkVars(s.expr, bump);
81+ else if (s.kind === 'For') {
82+ walkVars(s.start, bump);
83+ walkVars(s.end, bump);
84+ walkStmts(s.body);
85+ }
86+ }
87+ };
88+ walkStmts(stmts);
89+ return counts;
90+}
91+
92+function walkVars(e: IRExpr, visit: (cName: string) => void): void {
93+ const walk = (x: IRExpr): void => {
94+ switch (x.kind) {
95+ case 'Var':
96+ visit(x.cName);
97+ return;
98+ case 'Binary':
99+ walk(x.left);
100+ walk(x.right);
101+ return;
102+ case 'Unary':
103+ walk(x.operand);
104+ return;
105+ case 'Call':
106+ x.args.forEach(walk);
107+ return;
108+ case 'IndexSlice':
109+ walk(x.base);
110+ x.index.forEach((a) => {
111+ const inner = (a as { expr?: IRExpr }).expr;
112+ if (inner) walk(inner);
113+ });
114+ return;
115+ case 'MakeRange':
116+ walk(x.start);
117+ walk(x.step);
118+ walk(x.end);
119+ return;
120+ default:
121+ return;
122+ }
123+ };
124+ walk(e);
125+}
126+
127+function countIn(e: IRExpr, cName: string): number {
128+ let n = 0;
129+ walkVars(e, (c) => {
130+ if (c === cName) n++;
131+ });
132+ return n;
133+}
134+
135+/** Replace the (single) `Var` read of `cName` with `replacement`. */
136+function substitute(e: IRExpr, cName: string, replacement: IRExpr): IRExpr {
137+ const sub = (x: IRExpr): IRExpr => {
138+ if (x.kind === 'Var' && x.cName === cName) return replacement;
139+ switch (x.kind) {
140+ case 'Binary':
141+ x.left = sub(x.left);
142+ x.right = sub(x.right);
143+ return x;
144+ case 'Unary':
145+ x.operand = sub(x.operand);
146+ return x;
147+ case 'Call':
148+ for (let i = 0; i < x.args.length; i++) x.args[i] = sub(x.args[i]);
149+ return x;
150+ case 'IndexSlice':
151+ x.base = sub(x.base);
152+ return x;
153+ case 'MakeRange':
154+ x.start = sub(x.start);
155+ x.step = sub(x.step);
156+ x.end = sub(x.end);
157+ return x;
158+ default:
159+ return x;
160+ }
161+ };
162+ return sub(e);
163+}
src/mgpu/kernels.tsadded+247−0View file
@@ -0,0 +1,247 @@
1+/**
2+ * The non-elementwise kernels: tiled GEMM, tiled transpose, and reductions.
3+ *
4+ * Everything is column-major f32, matching MATLAB's layout, so `A(:)` walks
5+ * the same linear buffer the kernels index. All shapes are compile-time
6+ * constants baked into the WGSL — nothing here reads a dims uniform, which is
7+ * what lets the planner prebuild every pipeline and bind group.
8+ *
9+ * The GEMM is the classic 16x16 shared-memory tile (adapted to column-major
10+ * from matmul-bench's row-major sgemm). It will not beat a tuned BLAS, but it
11+ * is the honest baseline for "what does A*B cost in a browser".
12+ *
13+ * Reductions take a *fused loader*: the per-element expression emitted by
14+ * wgsl.ts, so `sum(a .* b + c)` reads its operands exactly once, in one pass.
15+ */
16+import type { FusedLoader } from './wgsl.ts';
17+
18+const TILE = 16;
19+
20+/** C(m x n) = A(m x k) * B(k x n), column-major. Bindings: 0=C, 1=A, 2=B. */
21+export function gemmKernel(m: number, k: number, n: number): string {
22+ return `
23+@group(0) @binding(0) var<storage, read_write> C: array<f32>;
24+@group(0) @binding(1) var<storage, read> A: array<f32>;
25+@group(0) @binding(2) var<storage, read> B: array<f32>;
26+
27+var<workgroup> tileA: array<array<f32, ${TILE}>, ${TILE}>;
28+var<workgroup> tileB: array<array<f32, ${TILE}>, ${TILE}>;
29+
30+@compute @workgroup_size(${TILE}, ${TILE})
31+fn main(
32+ @builtin(global_invocation_id) gid: vec3<u32>,
33+ @builtin(local_invocation_id) lid: vec3<u32>,
34+) {
35+ let row = gid.x;
36+ let col = gid.y;
37+ let lx = lid.x;
38+ let ly = lid.y;
39+ var acc: f32 = 0.0;
40+ let numTiles = ${Math.ceil(k / TILE)}u;
41+ for (var t: u32 = 0u; t < numTiles; t = t + 1u) {
42+ // Consecutive lx reads consecutive addresses in both loads (column-major).
43+ let aCol = t * ${TILE}u + ly;
44+ let bRow = t * ${TILE}u + lx;
45+ tileA[lx][ly] = select(0.0, A[row + aCol * ${m}u], row < ${m}u && aCol < ${k}u);
46+ tileB[lx][ly] = select(0.0, B[bRow + col * ${k}u], bRow < ${k}u && col < ${n}u);
47+ workgroupBarrier();
48+ for (var p: u32 = 0u; p < ${TILE}u; p = p + 1u) {
49+ acc = acc + tileA[lx][p] * tileB[p][ly];
50+ }
51+ workgroupBarrier();
52+ }
53+ if (row < ${m}u && col < ${n}u) {
54+ C[row + col * ${m}u] = acc;
55+ }
56+}
57+`;
58+}
59+
60+export const gemmDispatch = (m: number, n: number): [number, number] => [
61+ Math.ceil(m / TILE),
62+ Math.ceil(n / TILE),
63+];
64+
65+/** out(n x m) = in(m x n)', column-major, staged through a tile so both the
66+ * read and the write are coalesced. Bindings: 0=out, 1=in. */
67+export function transposeKernel(m: number, n: number): string {
68+ return `
69+@group(0) @binding(0) var<storage, read_write> out: array<f32>;
70+@group(0) @binding(1) var<storage, read> src: array<f32>;
71+
72+var<workgroup> tile: array<array<f32, ${TILE}>, ${TILE}>;
73+
74+@compute @workgroup_size(${TILE}, ${TILE})
75+fn main(
76+ @builtin(workgroup_id) wg: vec3<u32>,
77+ @builtin(local_invocation_id) lid: vec3<u32>,
78+) {
79+ let lx = lid.x;
80+ let ly = lid.y;
81+ // Read block (wg.x, wg.y) of src: rows wg.x*T.., cols wg.y*T..
82+ let sr = wg.x * ${TILE}u + lx;
83+ let sc = wg.y * ${TILE}u + ly;
84+ if (sr < ${m}u && sc < ${n}u) {
85+ tile[ly][lx] = src[sr + sc * ${m}u];
86+ }
87+ workgroupBarrier();
88+ // Write the transposed block: rows of out are cols of src.
89+ let dr = wg.y * ${TILE}u + lx;
90+ let dc = wg.x * ${TILE}u + ly;
91+ if (dr < ${n}u && dc < ${m}u) {
92+ out[dr + dc * ${n}u] = tile[lx][ly];
93+ }
94+}
95+`;
96+}
97+
98+export const transposeDispatch = (m: number, n: number): [number, number] => [
99+ Math.ceil(m / TILE),
100+ Math.ceil(n / TILE),
101+];
102+
103+// ── Reductions ──────────────────────────────────────────────────────────
104+
105+export type Combine = 'add' | 'mul' | 'max' | 'min';
106+/** Applied to each loaded element before combining. `sq` serves norm/dot. */
107+export type MapKind = 'id' | 'sq';
108+/** Applied to the final value. `scale` divides (mean); `sqrt` closes norm. */
109+export type Epilogue = { kind: 'none' } | { kind: 'scale'; by: number } | { kind: 'sqrt' };
110+
111+const REDUCE_WG = 256;
112+
113+const IDENT: Record<Combine, string> = {
114+ add: '0.0',
115+ mul: '1.0',
116+ max: '-3.4028234663852886e+38',
117+ min: '3.4028234663852886e+38',
118+};
119+
120+const comb = (c: Combine, a: string, b: string): string =>
121+ c === 'add' ? `${a} + ${b}`
122+ : c === 'mul' ? `${a} * ${b}`
123+ : `${c}(${a}, ${b})`;
124+
125+const mapped = (m: MapKind, v: string): string => (m === 'sq' ? `${v} * ${v}` : v);
126+
127+const epilogued = (e: Epilogue, v: string): string =>
128+ e.kind === 'scale' ? `(${v}) * ${e.by}` : e.kind === 'sqrt' ? `sqrt(${v})` : v;
129+
130+const treeReduce = (c: Combine): string => `
131+ sdata[li] = acc;
132+ workgroupBarrier();
133+ var stride = ${REDUCE_WG / 2}u;
134+ while (stride > 0u) {
135+ if (li < stride) {
136+ sdata[li] = ${comb(c, 'sdata[li]', 'sdata[li + stride]')};
137+ }
138+ workgroupBarrier();
139+ stride = stride / 2u;
140+ }`;
141+
142+/** Number of pass-1 partials for a full reduction over `count` elements. */
143+export function reducePartials(count: number): number {
144+ return Math.max(1, Math.min(1024, Math.ceil(count / (REDUCE_WG * 8))));
145+}
146+
147+/**
148+ * Full reduction, pass 1: `numWg` workgroups grid-stride over `count`
149+ * elements of the fused loader, leaving one partial each. Loader binding
150+ * declarations (`decls`) come from wgsl.ts's bindingDecls, whose binding 0
151+ * ("out") is the partials buffer here.
152+ */
153+export function reduceFullPass1(
154+ decls: string[],
155+ loader: FusedLoader,
156+ count: number,
157+ numWg: number,
158+ combine: Combine,
159+ map: MapKind,
160+): string {
161+ return `${decls.join('\n')}
162+${loader.helpers}
163+var<workgroup> sdata: array<f32, ${REDUCE_WG}>;
164+
165+@compute @workgroup_size(${REDUCE_WG})
166+fn main(
167+ @builtin(workgroup_id) wg: vec3<u32>,
168+ @builtin(local_invocation_id) lid: vec3<u32>,
169+) {
170+ let li = lid.x;
171+ var acc: f32 = ${IDENT[combine]};
172+ var i = wg.x * ${REDUCE_WG}u + li;
173+ while (i < ${count}u) {
174+ let v = ${loader.body};
175+ acc = ${comb(combine, 'acc', mapped(map, 'v'))};
176+ i = i + ${numWg * REDUCE_WG}u;
177+ }
178+${treeReduce(combine)}
179+ if (li == 0u) { out[wg.x] = sdata[0]; }
180+}
181+`;
182+}
183+
184+/** Full reduction, pass 2: one workgroup folds the partials and applies the
185+ * epilogue. Bindings: 0=result (1 element), 1=partials. */
186+export function reduceFullPass2(
187+ numPartials: number,
188+ combine: Combine,
189+ epilogue: Epilogue,
190+): string {
191+ return `
192+@group(0) @binding(0) var<storage, read_write> out: array<f32>;
193+@group(0) @binding(1) var<storage, read> partials: array<f32>;
194+var<workgroup> sdata: array<f32, ${REDUCE_WG}>;
195+
196+@compute @workgroup_size(${REDUCE_WG})
197+fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
198+ let li = lid.x;
199+ var acc: f32 = ${IDENT[combine]};
200+ var i = li;
201+ while (i < ${numPartials}u) {
202+ acc = ${comb(combine, 'acc', 'partials[i]')};
203+ i = i + ${REDUCE_WG}u;
204+ }
205+${treeReduce(combine)}
206+ if (li == 0u) { out[0] = ${epilogued(epilogue, 'sdata[0]')}; }
207+}
208+`;
209+}
210+
211+/**
212+ * Column-wise reduction of an (m x n) input: one workgroup per column,
213+ * threads grid-stride down the column (contiguous in column-major), leaving
214+ * out[column]. Dispatch n workgroups.
215+ */
216+export function reduceColumns(
217+ decls: string[],
218+ loader: FusedLoader,
219+ m: number,
220+ combine: Combine,
221+ map: MapKind,
222+ epilogue: Epilogue,
223+): string {
224+ return `${decls.join('\n')}
225+${loader.helpers}
226+var<workgroup> sdata: array<f32, ${REDUCE_WG}>;
227+
228+@compute @workgroup_size(${REDUCE_WG})
229+fn main(
230+ @builtin(workgroup_id) wg: vec3<u32>,
231+ @builtin(local_invocation_id) lid: vec3<u32>,
232+) {
233+ let li = lid.x;
234+ let col = wg.x;
235+ var acc: f32 = ${IDENT[combine]};
236+ var r = li;
237+ while (r < ${m}u) {
238+ let i = r + col * ${m}u;
239+ let v = ${loader.body};
240+ acc = ${comb(combine, 'acc', mapped(map, 'v'))};
241+ r = r + ${REDUCE_WG}u;
242+ }
243+${treeReduce(combine)}
244+ if (li == 0u) { out[col] = ${epilogued(epilogue, 'sdata[0]')}; }
245+}
246+`;
247+}
src/mgpu/numbl.d.tsadded+275−0View file
@@ -0,0 +1,275 @@
1+/**
2+ * The numbl compiler surface this project depends on.
3+ *
4+ * We reach past numbl's published entry points into its JIT internals (parser,
5+ * lowerer, IR, inline pass, builtin registry), which its package `exports` map
6+ * does not expose. Those imports resolve through the `numbl-src` alias in
7+ * vite.config.ts; these declarations are what TypeScript checks against.
8+ *
9+ * This follows turing-surface's arrangement (see its src/mgpu/numbl.d.ts for
10+ * the rationale); the surface here is wider because the sandbox lowers whole
11+ * scripts rather than individual functions, and patches builtin type rules.
12+ *
13+ * Only the nodes the WGSL backend actually walks are spelled out; every other
14+ * IR kind is collapsed into a catch-all so that unhandled constructs are
15+ * rejected with a message instead of being silently mis-compiled.
16+ */
17+
18+declare module 'numbl-src/numbl-core/jit/lowering/types.ts' {
19+ export type Sign =
20+ | 'positive' | 'nonneg' | 'negative' | 'nonpositive'
21+ | 'zero' | 'nonzero' | 'unknown';
22+
23+ export type DimInfo = { kind: 'exact'; value: number } | { kind: 'unknown' };
24+
25+ export type NumericExact =
26+ | number
27+ | Float64Array
28+ | { re: number; im: number }
29+ | { re: Float64Array; im: Float64Array };
30+
31+ export interface NumericType {
32+ kind: 'Numeric';
33+ elem: 'double' | 'logical' | 'char' | string;
34+ isComplex: boolean;
35+ dims: DimInfo[];
36+ /** Present iff every dim is exact. */
37+ shape?: number[];
38+ sign: Sign;
39+ exact?: NumericExact;
40+ }
41+
42+ /** Everything the WGSL backend rejects. */
43+ export interface NonNumericType {
44+ kind: 'Void' | 'Unknown' | 'String' | 'Handle' | 'Struct' | 'Class' | 'Cell';
45+ }
46+
47+ export type Type = NumericType | NonNumericType;
48+
49+ export function isMultiElement(t: NumericType): boolean;
50+ export function isNumeric(t: Type): t is NumericType;
51+ export function isScalar(t: NumericType): boolean;
52+ export function tensorDouble(shape: number[], exact?: Float64Array): NumericType;
53+ export function scalarDouble(sign?: Sign, exact?: number): NumericType;
54+}
55+
56+declare module 'numbl-src/numbl-core/jit/lowering/ir.ts' {
57+ import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
58+
59+ export interface Span {
60+ file: string;
61+ start: number;
62+ end: number;
63+ }
64+
65+ export interface NumLit {
66+ kind: 'NumLit';
67+ value: number;
68+ ty: Type;
69+ span: Span;
70+ }
71+ export interface StringLit {
72+ kind: 'StringLit';
73+ value: string;
74+ ty: Type;
75+ span: Span;
76+ }
77+ export interface Var {
78+ kind: 'Var';
79+ name: string;
80+ cName: string;
81+ ty: Type;
82+ span: Span;
83+ }
84+ export interface Binary {
85+ kind: 'Binary';
86+ builtin: string;
87+ left: IRExpr;
88+ right: IRExpr;
89+ ty: Type;
90+ span: Span;
91+ }
92+ export interface Unary {
93+ kind: 'Unary';
94+ builtin: string;
95+ operand: IRExpr;
96+ ty: Type;
97+ span: Span;
98+ }
99+ export interface Call {
100+ kind: 'Call';
101+ cName: string;
102+ name: string;
103+ args: IRExpr[];
104+ ty: Type;
105+ span: Span;
106+ }
107+ /** One slot of an `IndexSlice`'s subscript list. Only the full colon
108+ * (`A(:)`) is executed by this backend; the other shapes exist so they can
109+ * be rejected with a source location. */
110+ export type IndexSliceArg =
111+ | { kind: 'Colon' }
112+ | { kind: 'Range' | 'Scalar' | 'Gather'; [k: string]: unknown };
113+ export interface IndexSlice {
114+ kind: 'IndexSlice';
115+ base: IRExpr;
116+ index: ReadonlyArray<IndexSliceArg>;
117+ ty: Type;
118+ span: Span;
119+ }
120+ export interface MakeRange {
121+ kind: 'MakeRange';
122+ start: IRExpr;
123+ step: IRExpr;
124+ end: IRExpr;
125+ ty: Type;
126+ span: Span;
127+ }
128+ /** Any other IR expression kind — rejected by the WGSL emitter. */
129+ export interface OtherExpr {
130+ kind:
131+ | 'ImagLit' | 'TensorBuild' | 'TensorConcat' | 'CellLit'
132+ | 'CellEmpty' | 'CellIndexLoad' | 'HandleLit' | 'HandleCaptureLoad'
133+ | 'StructLit' | 'MemberLoad' | 'IndexLoad' | 'EndRef';
134+ ty: Type;
135+ span: Span;
136+ }
137+
138+ export type IRExpr =
139+ | NumLit | StringLit | Var | Binary | Unary | Call
140+ | IndexSlice | MakeRange | OtherExpr;
141+
142+ export interface Assign {
143+ kind: 'Assign';
144+ name: string;
145+ cName: string;
146+ ty: Type;
147+ expr: IRExpr;
148+ span: Span;
149+ }
150+ export interface ExprStmt {
151+ kind: 'ExprStmt';
152+ expr: IRExpr;
153+ span: Span;
154+ }
155+ /** A counted loop. `step` is already a literal number in the IR — numbl
156+ * rejects a non-literal step during lowering — while `start` and `end` are
157+ * expressions that must carry an exact value for the planner to accept the
158+ * loop. */
159+ export interface For {
160+ kind: 'For';
161+ /** Loop variable, as written in the .m. */
162+ varName: string;
163+ /** Loop variable's cName, the key kernels read its value under. */
164+ cVar: string;
165+ start: IRExpr;
166+ step: number;
167+ end: IRExpr;
168+ body: IRStmt[];
169+ span: Span;
170+ }
171+ /** Any other IR statement kind — rejected by the planner. */
172+ export interface OtherStmt {
173+ kind:
174+ | 'If' | 'While' | 'ReturnFromFunction' | 'Break'
175+ | 'Continue' | 'TypeComment' | 'MemberStore' | 'MultiAssignCall'
176+ | 'IndexStore' | 'IndexSliceStore' | 'CellIndexStore';
177+ span: Span;
178+ }
179+
180+ export type IRStmt = Assign | ExprStmt | For | OtherStmt;
181+
182+ export interface IRFunc {
183+ name: string;
184+ cName: string;
185+ params: string[];
186+ cParams: string[];
187+ paramTypes: Type[];
188+ outputs: string[];
189+ cOutputs: string[];
190+ outputTypes: Type[];
191+ body: IRStmt[];
192+ span: Span;
193+ }
194+
195+ export interface IRProgram {
196+ topLevelStmts: IRStmt[];
197+ functions: Map<string, IRFunc>;
198+ }
199+}
200+
201+declare module 'numbl-src/numbl-core/parser/index.ts' {
202+ /** Parser statement. Only the fields the sandbox reads (statement span and
203+ * the `;` suppression flag, which decides MATLAB-style display) are
204+ * declared. Control-flow statements (`for`, `if`, ...) carry no
205+ * `suppressed` flag. */
206+ export interface AstStatement {
207+ type: string;
208+ suppressed?: boolean;
209+ span: { start: number; end: number };
210+ }
211+ export interface AbstractSyntaxTree {
212+ body: AstStatement[];
213+ }
214+ export function parseMFile(input: string, fileName?: string): AbstractSyntaxTree;
215+ export class SyntaxError extends Error {}
216+}
217+
218+declare module 'numbl-src/numbl-core/jit/index.ts' {
219+ import type { AbstractSyntaxTree } from 'numbl-src/numbl-core/parser/index.ts';
220+ import type { IRProgram, Span } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
221+ import type { NumericType, Sign } from 'numbl-src/numbl-core/jit/lowering/types.ts';
222+
223+ export interface WorkspaceFile {
224+ name: string;
225+ source: string;
226+ ast?: AbstractSyntaxTree;
227+ }
228+
229+ export class Workspace {
230+ constructor(mainFile: string, searchPaths?: ReadonlyArray<string>);
231+ addFile(file: WorkspaceFile): void;
232+ finalize(): void;
233+ }
234+
235+ export class Lowerer {
236+ constructor(workspace: Workspace);
237+ lowerProgram(ast: AbstractSyntaxTree): IRProgram;
238+ }
239+
240+ /** Thrown for MATLAB the JIT pipeline cannot lower; carries a source span. */
241+ export class UnsupportedConstruct extends Error {
242+ span?: Span;
243+ }
244+ export class JitTypeError extends Error {
245+ span?: Span;
246+ }
247+
248+ export function tensorDouble(shape: number[], exact?: Float64Array): NumericType;
249+ export function scalarDouble(sign?: Sign, exact?: number): NumericType;
250+ export function isMultiElement(t: NumericType): boolean;
251+}
252+
253+declare module 'numbl-src/numbl-core/jit/codegen/inlinePass.ts' {
254+ import type { IRProgram } from 'numbl-src/numbl-core/jit/lowering/ir.ts';
255+ /** Folds single-use ANF temps into their consumer, in place. */
256+ export function inlinePass(prog: IRProgram): void;
257+}
258+
259+declare module 'numbl-src/numbl-core/jit/builtins/index.ts' {
260+ import type { Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
261+
262+ /** A JIT builtin. Beyond `name` and `transfer` the object carries emit hooks
263+ * and interpreter call hooks the sandbox never invokes (only numbl's own
264+ * backends do); the index signature carries them through when patching. */
265+ export interface Builtin {
266+ name: string;
267+ /** Safe to evaluate one output element from one input element per slot. */
268+ elementwise?: boolean;
269+ transfer(argTypes: Type[], nargout: number): Type[];
270+ [k: string]: unknown;
271+ }
272+ export function getBuiltin(name: string): Builtin | undefined;
273+ /** Overwrites any existing builtin of the same name (by design). */
274+ export function registerBuiltin(b: Builtin): void;
275+}
src/mgpu/patches.tsadded+186−0View file
@@ -0,0 +1,186 @@
1+/**
2+ * Type-rule patches for numbl JIT builtins, applied to the in-process builtin
3+ * registry before lowering a script.
4+ *
5+ * numbl's JIT declines a handful of constructs its C/JS backends do not
6+ * implement yet (tensor `&`/`|`/`~`, two-arg `max`/`min` on tensors, matrix
7+ * `rand`/`randn`), and its comparison builtins deliberately return a
8+ * shape-unknown logical tensor because the lowerer has no static broadcast
9+ * helper. This backend executes the IR on WebGPU — numbl's own emitters never
10+ * run here — so all this sandbox needs from those builtins is a precise
11+ * *type*: everything else is supplied by the WGSL kernels.
12+ *
13+ * `registerBuiltin` overwrites by name (by design, for HMR), which makes it
14+ * the sanctioned way to swap a builtin's rules. The patches only ever widen:
15+ * every case the original transfer accepts is delegated to it unchanged, so
16+ * scalar behavior (including exact-value folding) is untouched.
17+ *
18+ * IMPORTANT: the patch is process-global to this module instance of numbl.
19+ * The CPU comparison runner executes scripts through numbl's normal engine in
20+ * a separate Web Worker, whose module instances are its own — it never sees
21+ * these patches.
22+ */
23+import {
24+ getBuiltin,
25+ registerBuiltin,
26+ type Builtin,
27+} from 'numbl-src/numbl-core/jit/builtins/index.ts';
28+import {
29+ isNumeric,
30+ isMultiElement,
31+ type NumericType,
32+ type Type,
33+} from 'numbl-src/numbl-core/jit/lowering/types.ts';
34+import { UnsupportedOnGpu } from './errors.ts';
35+
36+const isRealTensor = (t: Type): t is NumericType =>
37+ isNumeric(t) && !t.isComplex && isMultiElement(t);
38+const isRealScalarish = (t: Type): t is NumericType =>
39+ isNumeric(t) && !t.isComplex && !isMultiElement(t);
40+
41+const logicalTensor = (shape: number[]): NumericType => ({
42+ kind: 'Numeric',
43+ elem: 'logical',
44+ isComplex: false,
45+ dims: shape.map((n) => ({ kind: 'exact', value: n })),
46+ shape: [...shape],
47+ sign: 'nonneg',
48+});
49+
50+const doubleTensor = (shape: number[], sign: NumericType['sign']): NumericType => ({
51+ kind: 'Numeric',
52+ elem: 'double',
53+ isComplex: false,
54+ dims: shape.map((n) => ({ kind: 'exact', value: n })),
55+ shape: [...shape],
56+ sign,
57+});
58+
59+const sameShape = (a: number[], b: number[]): boolean =>
60+ a.length === b.length && a.every((d, i) => d === b[i]);
61+
62+/**
63+ * The elementwise result shape for two operands, under the same rule the
64+ * kernels implement: operands are either scalars or tensors of one common
65+ * shape. Implicit expansion (`n x 1` with `1 x m`) is rejected here, matching
66+ * the emitter's refusal to broadcast.
67+ */
68+function elementwiseShape(name: string, argTypes: Type[]): number[] | null {
69+ let shape: number[] | null = null;
70+ for (const t of argTypes) {
71+ if (!isNumeric(t) || t.isComplex) {
72+ throw new UnsupportedOnGpu(`'${name}': operands must be real numeric`);
73+ }
74+ if (!isMultiElement(t)) continue;
75+ if (!t.shape) {
76+ throw new UnsupportedOnGpu(
77+ `'${name}': operand shape is not known at compile time`,
78+ );
79+ }
80+ if (shape && !sameShape(shape, t.shape)) {
81+ throw new UnsupportedOnGpu(
82+ `'${name}': operands are ${shape.join('x')} and ${t.shape.join('x')}; ` +
83+ `the GPU kernels do not broadcast — expand explicitly`,
84+ );
85+ }
86+ shape = t.shape;
87+ }
88+ return shape;
89+}
90+
91+/** Wrap `transfer` so tensor operands get an elementwise result type instead
92+ * of the original's decline (or, for comparisons, its shape-unknown type). */
93+function widenElementwise(
94+ name: string,
95+ makeResult: (shape: number[]) => NumericType,
96+ /** Arity the elementwise form requires. `max(x)` with one tensor arg is the
97+ * reduction form and must stay with the original transfer. */
98+ arity?: number,
99+): void {
100+ const orig = getBuiltin(name);
101+ if (!orig) throw new Error(`patch: no builtin named '${name}'`);
102+ const origTransfer = orig.transfer.bind(orig);
103+ registerBuiltin({
104+ ...orig,
105+ transfer(argTypes: Type[], nargout: number): Type[] {
106+ const anyTensor = argTypes.some((t) => isNumeric(t) && isMultiElement(t));
107+ if (!anyTensor || (arity !== undefined && argTypes.length !== arity)) {
108+ return origTransfer(argTypes, nargout);
109+ }
110+ if (nargout > 1) {
111+ throw new UnsupportedOnGpu(`'${name}' returns one value`);
112+ }
113+ const shape = elementwiseShape(name, argTypes);
114+ // anyTensor guaranteed a tensor operand, so shape is non-null.
115+ return [makeResult(shape!)];
116+ },
117+ } as Builtin);
118+}
119+
120+/** Comparisons: `lt` et al. accept tensors already but return unknown dims;
121+ * replace the tensor case with the precise shape. */
122+const COMPARISONS = ['lt', 'le', 'gt', 'ge', 'eq', 'ne'];
123+
124+/** Eager elementwise `&` / `|` / `xor` / `~`. (`&&`/`||` stay scalar-only —
125+ * that is MATLAB semantics, not a backend gap.) */
126+const LOGICALS = ['and', 'or', 'xor', 'not'];
127+
128+/** Two-arg elementwise forms of `max`/`min`. The one-arg reduction form is
129+ * typed by the original transfer. */
130+const MINMAX = ['max', 'min'];
131+
132+/** `rand(n)` / `rand(m, n)` / `randn(...)`: numbl's JS-JIT supports only the
133+ * scalar form; here any exact-sized matrix form is a fill kernel. */
134+function widenRand(name: string, sign: NumericType['sign']): void {
135+ // `randn` has no JIT builtin at all in numbl yet; register it whole.
136+ const orig = getBuiltin(name);
137+ const origTransfer = orig
138+ ? orig.transfer.bind(orig)
139+ : (): Type[] => [{
140+ kind: 'Numeric', elem: 'double', isComplex: false,
141+ dims: [{ kind: 'exact', value: 1 }, { kind: 'exact', value: 1 }],
142+ shape: [1, 1], sign,
143+ } satisfies NumericType];
144+ registerBuiltin({
145+ ...(orig ?? { name }),
146+ name,
147+ transfer(argTypes: Type[], nargout: number): Type[] {
148+ if (argTypes.length === 0) return origTransfer(argTypes, nargout);
149+ if (nargout > 1) throw new UnsupportedOnGpu(`'${name}' returns one value`);
150+ if (argTypes.length > 2) {
151+ throw new UnsupportedOnGpu(
152+ `'${name}': only the 2-D forms ${name}(n) / ${name}(m, n) are supported`,
153+ );
154+ }
155+ const dims = argTypes.map((t) => {
156+ if (!isRealScalarish(t) || typeof t.exact !== 'number') {
157+ throw new UnsupportedOnGpu(
158+ `'${name}': array sizes must be known at compile time (assign the ` +
159+ `size from a literal, e.g. n = 1024)`,
160+ );
161+ }
162+ if (!Number.isInteger(t.exact) || t.exact < 0) {
163+ throw new UnsupportedOnGpu(`'${name}': sizes must be whole numbers`);
164+ }
165+ return t.exact;
166+ });
167+ const shape = dims.length === 1 ? [dims[0], dims[0]] : dims;
168+ return [doubleTensor(shape, sign)];
169+ },
170+ } as Builtin);
171+}
172+
173+let applied = false;
174+
175+/** Apply every patch, once per module instance. */
176+export function applyBuiltinPatches(): void {
177+ if (applied) return;
178+ applied = true;
179+ for (const name of COMPARISONS) widenElementwise(name, logicalTensor);
180+ for (const name of LOGICALS) widenElementwise(name, logicalTensor);
181+ for (const name of MINMAX) {
182+ widenElementwise(name, (shape) => doubleTensor(shape, 'unknown'), 2);
183+ }
184+ widenRand('rand', 'positive');
185+ widenRand('randn', 'unknown');
186+}
src/mgpu/plan.tsadded+1221−0View file
@@ -0,0 +1,1221 @@
1+/**
2+ * Lowered script -> a replayable sequence of GPU + host operations.
3+ *
4+ * Everything expensive happens once, here: pipeline compilation, buffer
5+ * allocation, bind-group construction. Because numbl fixes every type and
6+ * shape at lowering time, the op sequence is fully static; executing it is
7+ * pure command recording plus the host ops (tic/toc/printing) the script asked
8+ * for, which are the only synchronization points.
9+ *
10+ * Statement routing:
11+ * - elementwise trees -> one fused kernel (buildKernel)
12+ * - `A * B` (both tensors) -> tiled GEMM; non-variable operands are
13+ * materialized into scratch first
14+ * - matrix `A'` -> tiled transpose kernel
15+ * - vector `x'`, `X(:)`, -> a view of the same buffer when the source
16+ * `reshape` is never reassigned, else a plain copy
17+ * - sum/mean/prod/max/min/ -> reduction kernels over a fused loader
18+ * norm/dot
19+ * - tic/toc/disp/fprintf and -> host ops: the executor flushes GPU work,
20+ * unsuppressed echoes then times/reads/prints
21+ * - `for` with exact bounds -> body planned ONCE; the loop variable lives
22+ * in a dynamic-offset uniform (one 256-byte
23+ * slot per iteration) and the executor
24+ * re-encodes the body per iteration
25+ */
26+import type {
27+ Assign,
28+ Call,
29+ ExprStmt,
30+ For,
31+ IRExpr,
32+ IRStmt,
33+ Span,
34+} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
35+import {
36+ isMultiElement,
37+ type NumericType,
38+ type Type,
39+} from 'numbl-src/numbl-core/jit/lowering/types.ts';
40+import type { CompiledScript } from './compile.ts';
41+import { UnsupportedOnGpu } from './errors.ts';
42+import {
43+ buildKernel,
44+ bindingDecls,
45+ dispatchFor,
46+ emitLoader,
47+ exactValue,
48+ fullColonBase,
49+ numel,
50+ REDUCTIONS,
51+ type KernelInputs,
52+} from './wgsl.ts';
53+import {
54+ gemmDispatch,
55+ gemmKernel,
56+ reduceColumns,
57+ reduceFullPass1,
58+ reduceFullPass2,
59+ reducePartials,
60+ transposeDispatch,
61+ transposeKernel,
62+ type Combine,
63+ type Epilogue,
64+ type MapKind,
65+} from './kernels.ts';
66+
67+const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
68+const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
69+const isVectorish = (t: Type): boolean =>
70+ isNumeric(t) && (t.shape ?? []).filter((d) => d !== 1).length <= 1;
71+
72+/** Iterations a single `for` may replay. Each costs one 256-byte uniform slot
73+ * and a re-encode of the body's dispatches. */
74+const MAX_TRIPS = 65536;
75+/** Total dispatches one run may encode, across all loops. */
76+const MAX_DISPATCHES = 2_000_000;
77+/** Loop-variable uniform slot stride (minUniformBufferOffsetAlignment). */
78+const LV_STRIDE = 256;
79+
80+export interface Slot {
81+ buffer: GPUBuffer;
82+ count: number;
83+}
84+
85+/** What a printed/displayed value reads from. */
86+export type ValueRef =
87+ | { kind: 'literal'; value: number }
88+ | { kind: 'buffer'; slot: Slot; count: number }
89+ | { kind: 'host'; cName: string };
90+
91+/** One piece of an fprintf: fixed text or a formatted value. */
92+export type EmitPart =
93+ | { kind: 'text'; text: string }
94+ | { kind: 'value'; ref: ValueRef; spec: string };
95+
96+export type Op =
97+ | {
98+ kind: 'kernel';
99+ pipeline: GPUComputePipeline;
100+ bindGroup: GPUBindGroup;
101+ dispatch: [number, number];
102+ /** cVars whose dynamic offsets must be passed, in binding order. */
103+ loops: string[];
104+ label: string;
105+ copyBack?: { from: GPUBuffer; to: GPUBuffer; bytes: number };
106+ }
107+ | { kind: 'copy'; from: GPUBuffer; to: GPUBuffer; bytes: number; label: string }
108+ | { kind: 'write'; slot: Slot; data: Float32Array; label: string }
109+ | { kind: 'loop'; cVar: string; trips: number; uniform: GPUBuffer; body: Op[]; label: string }
110+ | { kind: 'tic'; assignTo?: { cName: string; slot: Slot } }
111+ | {
112+ kind: 'toc';
113+ print: boolean;
114+ sinceCName?: string;
115+ assignTo?: { cName: string; slot: Slot };
116+ /** 1-based tic..toc pair index, for the timing table. */
117+ seq: number;
118+ }
119+ | { kind: 'emit'; parts: EmitPart[] }
120+ | {
121+ kind: 'display';
122+ /** Variable name, or null for disp() (which prints no name). */
123+ label: string | null;
124+ ref: ValueRef;
125+ shape: number[];
126+ };
127+
128+export interface ScriptPlan {
129+ ops: Op[];
130+ /** Human-readable op sequence — what the script actually compiled to. */
131+ describe(): string[];
132+ destroy(): void;
133+}
134+
135+interface PlannerVar {
136+ slot: Slot;
137+ shape: number[];
138+}
139+
140+export async function planScript(
141+ device: GPUDevice,
142+ compiled: CompiledScript,
143+): Promise<ScriptPlan> {
144+ const owned: GPUBuffer[] = [];
145+ /** cName -> buffer-backed variable (tensors and runtime scalars). */
146+ const vars = new Map<string, PlannerVar>();
147+ /** cName -> cName it is a view of. */
148+ const aliases = new Map<string, string>();
149+ /** cName -> char value (format strings). */
150+ const chars = new Map<string, string>();
151+ /** cNames whose value the executor knows on the host (tic/toc results). */
152+ const hostScalars = new Set<string>();
153+ /** Loop-variable uniform buffers, by cVar. */
154+ const loopUniforms = new Map<string, GPUBuffer>();
155+ const pipelines = new Map<string, GPUComputePipeline>();
156+ const describeLines: string[] = [];
157+
158+ let seedCounter = 1;
159+ let tempCounter = 0;
160+ let tocCounter = 0;
161+ let dispatchBudget = MAX_DISPATCHES;
162+
163+ // How many times each cName is assigned, anywhere. Decides when a variable
164+ // may be a view of another's buffer, and when an exact scalar still needs a
165+ // real buffer (a later assignment makes its uses flow-dependent).
166+ const assignCounts = new Map<string, number>();
167+ {
168+ const walkCounts = (stmts: IRStmt[]): void => {
169+ for (const s of stmts) {
170+ if (s.kind === 'Assign') {
171+ assignCounts.set(s.cName, (assignCounts.get(s.cName) ?? 0) + 1);
172+ } else if (s.kind === 'For') {
173+ walkCounts(s.body);
174+ }
175+ }
176+ };
177+ walkCounts(compiled.stmts);
178+ }
179+
180+ const resolve = (cName: string): string => {
181+ let c = cName;
182+ while (aliases.has(c)) c = aliases.get(c)!;
183+ return c;
184+ };
185+
186+ const maxBytes = device.limits.maxStorageBufferBindingSize;
187+
188+ const makeSlot = (label: string, count: number): Slot => {
189+ const bytes = Math.max(4, 4 * count);
190+ if (bytes > maxBytes) {
191+ throw new UnsupportedOnGpu(
192+ `'${label}' needs ${(bytes / 1e6).toFixed(0)} MB, over this device's ` +
193+ `storage-buffer limit of ${(maxBytes / 1e6).toFixed(0)} MB`,
194+ );
195+ }
196+ const buffer = device.createBuffer({
197+ label: `mgpu-${label}`,
198+ size: bytes,
199+ usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
200+ });
201+ owned.push(buffer);
202+ return { buffer, count };
203+ };
204+
205+ /** The variable's slot, allocating on first assignment. */
206+ const slotFor = (cName: string, name: string, ty: NumericType, span: Span): Slot => {
207+ const key = resolve(cName);
208+ const count = numel(ty);
209+ const existing = vars.get(key);
210+ if (existing) {
211+ if (existing.slot.count !== count) {
212+ throw new UnsupportedOnGpu(
213+ `'${name}' changes size between assignments (${existing.slot.count} ` +
214+ `-> ${count} elements); the sandbox fixes each variable's storage once`,
215+ span,
216+ );
217+ }
218+ existing.shape = ty.shape ?? [count, 1];
219+ return existing.slot;
220+ }
221+ const slot = makeSlot(name, count);
222+ vars.set(key, { slot, shape: ty.shape ?? [count, 1] });
223+ return slot;
224+ };
225+
226+ const readSlot = (cName: string, name: string, span: Span): Slot => {
227+ const v = vars.get(resolve(cName));
228+ if (!v) {
229+ throw new UnsupportedOnGpu(`'${name}' is read before it has a value`, span);
230+ }
231+ return v.slot;
232+ };
233+
234+ async function pipeline(code: string, label: string, layout: GPUBindGroupLayout): Promise<GPUComputePipeline> {
235+ const hit = pipelines.get(code);
236+ if (hit) return hit;
237+ device.pushErrorScope('validation');
238+ const module = device.createShaderModule({ code, label });
239+ const info = await module.getCompilationInfo();
240+ const errors = info.messages.filter((m) => m.type === 'error');
241+ if (errors.length) {
242+ throw new UnsupportedOnGpu(
243+ `generated WGSL failed to compile for '${label}':\n` +
244+ errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
245+ `\n--- shader ---\n${code}`,
246+ );
247+ }
248+ const p = await device.createComputePipelineAsync({
249+ layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }),
250+ compute: { module, entryPoint: 'main' },
251+ label,
252+ });
253+ const err = await device.popErrorScope();
254+ if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
255+ pipelines.set(code, p);
256+ return p;
257+ }
258+
259+ /** Bind group layout: out at 0, `inputs` read-only buffers, then `loops`
260+ * dynamic-offset uniforms. Explicit so unused bindings still match. */
261+ function kernelLayout(inputs: number, loops: number): GPUBindGroupLayout {
262+ const entries: GPUBindGroupLayoutEntry[] = [
263+ { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
264+ ];
265+ for (let i = 0; i < inputs; i++) {
266+ entries.push({
267+ binding: i + 1,
268+ visibility: GPUShaderStage.COMPUTE,
269+ buffer: { type: 'read-only-storage' },
270+ });
271+ }
272+ for (let k = 0; k < loops; k++) {
273+ entries.push({
274+ binding: inputs + 1 + k,
275+ visibility: GPUShaderStage.COMPUTE,
276+ buffer: { type: 'uniform', hasDynamicOffset: true },
277+ });
278+ }
279+ return device.createBindGroupLayout({ entries });
280+ }
281+
282+ /** Ops are appended to the innermost list; loops nest via this stack. */
283+ const opStack: Op[][] = [[]];
284+ const ops = (): Op[] => opStack[opStack.length - 1];
285+ const enclosingLoops: string[] = [];
286+ const inLoop = (): boolean => enclosingLoops.length > 0;
287+
288+ const spendDispatches = (n: number, span: Span): void => {
289+ // Inside loops the body is re-encoded per iteration; multiply out.
290+ let mult = 1;
291+ for (const f of loopTrips) mult *= f;
292+ dispatchBudget -= n * mult;
293+ if (dispatchBudget < 0) {
294+ throw new UnsupportedOnGpu(
295+ `this script would encode more than ${MAX_DISPATCHES.toLocaleString()} ` +
296+ `GPU dispatches; shrink loop counts`,
297+ span,
298+ );
299+ }
300+ };
301+ const loopTrips: number[] = [];
302+
303+ /** KernelInputs seeded with fresh maps; buffers registered on demand. */
304+ const freshIo = (): KernelInputs => ({
305+ buffers: new Map(),
306+ loopVars: new Map(),
307+ nextSeed: () => seedCounter++,
308+ });
309+
310+ /** Register every buffer-backed read in `expr` into `io.buffers`. */
311+ function collectBuffers(expr: IRExpr, io: KernelInputs, span: Span): void {
312+ const walk = (x: IRExpr): void => {
313+ // A constant-folded subtree emits as a literal; nothing under it is read.
314+ if (exactValue(x) !== undefined) return;
315+ switch (x.kind) {
316+ case 'Var': {
317+ if (enclosingLoops.includes(x.cName)) {
318+ if (!io.loopVars.has(x.cName)) io.loopVars.set(x.cName, io.loopVars.size);
319+ return;
320+ }
321+ if (!isNumeric(x.ty)) {
322+ throw new UnsupportedOnGpu(`'${x.name}' is not numeric`, x.span);
323+ }
324+ if (typeof x.ty.exact === 'number') return; // folds to a literal
325+ const key = resolve(x.cName);
326+ if (!io.buffers.has(key)) io.buffers.set(key, io.buffers.size);
327+ // Rewrite so the emitter sees the resolved name.
328+ x.cName = key;
329+ return;
330+ }
331+ case 'Binary':
332+ walk(x.left);
333+ walk(x.right);
334+ return;
335+ case 'Unary':
336+ walk(x.operand);
337+ return;
338+ case 'IndexSlice':
339+ walk(x.base);
340+ return;
341+ case 'MakeRange':
342+ walk(x.start);
343+ walk(x.step);
344+ return;
345+ case 'Call':
346+ if (!['zeros', 'ones', 'eye', 'rand', 'randn'].includes(x.name)) {
347+ x.args.forEach(walk);
348+ }
349+ return;
350+ default:
351+ return;
352+ }
353+ };
354+ walk(expr);
355+ void span;
356+ }
357+
358+ /** Bind-group entries for a kernel built from `io`, writing `target`. */
359+ function kernelBindGroup(
360+ io: KernelInputs,
361+ bufferOrder: string[],
362+ loopOrder: string[],
363+ target: GPUBuffer,
364+ layout: GPUBindGroupLayout,
365+ span: Span,
366+ ): GPUBindGroup {
367+ const entries: GPUBindGroupEntry[] = [{ binding: 0, resource: { buffer: target } }];
368+ bufferOrder.forEach((cName, i) => {
369+ const v = vars.get(resolve(cName));
370+ if (!v) throw new UnsupportedOnGpu(`a value read here has no buffer`, span);
371+ entries.push({ binding: i + 1, resource: { buffer: v.slot.buffer } });
372+ });
373+ loopOrder.forEach((cVar, k) => {
374+ const u = loopUniforms.get(cVar);
375+ if (!u) throw new UnsupportedOnGpu(`internal: no uniform for loop '${cVar}'`, span);
376+ entries.push({
377+ binding: bufferOrder.length + 1 + k,
378+ resource: { buffer: u, offset: 0, size: 8 },
379+ });
380+ });
381+ return device.createBindGroup({ layout, entries });
382+ }
383+
384+ /** Materialize an arbitrary tensor expression into a slot, planning
385+ * whatever ops that takes. A plain Var is returned as-is. */
386+ async function materialize(expr: IRExpr): Promise<{ cName: string; slot: Slot }> {
387+ if (expr.kind === 'Var' && isTensor(expr.ty) && typeof (expr.ty as NumericType).exact !== 'object') {
388+ return { cName: resolve(expr.cName), slot: readSlot(expr.cName, expr.name, expr.span) };
389+ }
390+ if (expr.kind === 'IndexSlice') {
391+ const base = fullColonBase(expr);
392+ if (base && base.kind === 'Var') {
393+ return { cName: resolve(base.cName), slot: readSlot(base.cName, base.name, base.span) };
394+ }
395+ }
396+ if (!isNumeric(expr.ty)) {
397+ throw new UnsupportedOnGpu(`expression is not numeric`, expr.span);
398+ }
399+ const cName = `%tmp${tempCounter++}`;
400+ await planValue(cName, cName, expr.ty, expr, expr.span);
401+ return { cName, slot: vars.get(resolve(cName))!.slot };
402+ }
403+
404+ /** Replace non-elementwise subtrees (GEMM, matrix transpose, reductions)
405+ * with materialized temps, so what remains is one fused kernel. */
406+ async function hoistNonElementwise(expr: IRExpr): Promise<IRExpr> {
407+ const hoist = async (x: IRExpr): Promise<IRExpr> => {
408+ if (isHoistRoot(x)) {
409+ const { cName, slot } = await materialize(x);
410+ void slot;
411+ return {
412+ kind: 'Var',
413+ name: cName,
414+ cName,
415+ ty: x.ty,
416+ span: x.span,
417+ };
418+ }
419+ switch (x.kind) {
420+ case 'Binary':
421+ x.left = await hoist(x.left);
422+ x.right = await hoist(x.right);
423+ return x;
424+ case 'Unary':
425+ x.operand = await hoist(x.operand);
426+ return x;
427+ case 'Call':
428+ for (let i = 0; i < x.args.length; i++) x.args[i] = await hoist(x.args[i]);
429+ return x;
430+ case 'IndexSlice':
431+ x.base = await hoist(x.base);
432+ return x;
433+ default:
434+ return x;
435+ }
436+ };
437+ // The root itself was already routed by planValue; only hoist children.
438+ switch (expr.kind) {
439+ case 'Binary':
440+ expr.left = await hoist(expr.left);
441+ expr.right = await hoist(expr.right);
442+ return expr;
443+ case 'Unary':
444+ expr.operand = await hoist(expr.operand);
445+ return expr;
446+ case 'Call':
447+ for (let i = 0; i < expr.args.length; i++) expr.args[i] = await hoist(expr.args[i]);
448+ return expr;
449+ case 'IndexSlice':
450+ expr.base = await hoist(expr.base);
451+ return expr;
452+ default:
453+ return expr;
454+ }
455+ }
456+
457+ function isHoistRoot(x: IRExpr): boolean {
458+ if (x.kind === 'Binary' && x.builtin === 'mtimes' && isTensor(x.left.ty) && isTensor(x.right.ty)) {
459+ return true;
460+ }
461+ if (x.kind === 'Unary' && x.builtin === 'transpose' && isTensor(x.operand.ty) && !isVectorish(x.operand.ty)) {
462+ return true;
463+ }
464+ if (x.kind === 'Call' && REDUCTIONS.has(x.name)) {
465+ // Two-arg max/min is elementwise, not a reduction.
466+ if ((x.name === 'max' || x.name === 'min') && x.args.length === 2) return false;
467+ return x.args.some((a) => isTensor(a.ty));
468+ }
469+ return false;
470+ }
471+
472+ /** Plan `dest = expr`, routing to the right kind of op. */
473+ async function planValue(
474+ cName: string,
475+ name: string,
476+ ty: NumericType,
477+ expr: IRExpr,
478+ span: Span,
479+ ): Promise<void> {
480+ // View-or-copy forms: X(:), reshape, vector transpose, plain `B = A`.
481+ const viewOf = viewSource(expr);
482+ if (viewOf) {
483+ const srcSlot = readSlot(viewOf.cName, viewOf.name, span);
484+ const srcKey = resolve(viewOf.cName);
485+ if (numel(ty) !== srcSlot.count) {
486+ throw new UnsupportedOnGpu(`'${name}' and '${viewOf.name}' differ in size`, span);
487+ }
488+ const destAssigns = assignCounts.get(cName) ?? 1;
489+ const srcAssigns = assignCounts.get(viewOf.cName) ?? 1;
490+ if (destAssigns === 1 && srcAssigns === 1 && !inLoop() && !vars.has(cName)) {
491+ aliases.set(cName, srcKey);
492+ // Track the new orientation under the alias's own name via vars of
493+ // the base: display uses the Assign's ty directly, so nothing to do.
494+ describeLines.push(`view ${name} -> ${viewOf.name}`);
495+ return;
496+ }
497+ const dest = slotFor(cName, name, ty, span);
498+ if (dest.buffer !== srcSlot.buffer) {
499+ ops().push({
500+ kind: 'copy',
501+ from: srcSlot.buffer,
502+ to: dest.buffer,
503+ bytes: 4 * srcSlot.count,
504+ label: `${name} = ${viewOf.name}`,
505+ });
506+ describeLines.push(`copy ${name} = ${viewOf.name}`);
507+ }
508+ return;
509+ }
510+
511+ // GEMM: A * B with both sides tensors.
512+ if (expr.kind === 'Binary' && expr.builtin === 'mtimes' && isTensor(expr.left.ty) && isTensor(expr.right.ty)) {
513+ return planGemm(cName, name, ty, expr, span);
514+ }
515+ // Matrix transpose.
516+ if (expr.kind === 'Unary' && expr.builtin === 'transpose' && isTensor(expr.operand.ty) && !isVectorish(expr.operand.ty)) {
517+ return planTranspose(cName, name, ty, expr, span);
518+ }
519+ // Reductions.
520+ if (expr.kind === 'Call' && REDUCTIONS.has(expr.name) &&
521+ !((expr.name === 'max' || expr.name === 'min') && expr.args.length === 2) &&
522+ expr.args.some((a) => isTensor(a.ty))) {
523+ return planReduce(cName, name, ty, expr as Call, span);
524+ }
525+
526+ // Fused elementwise kernel (with any non-elementwise subtrees hoisted).
527+ const hoisted = await hoistNonElementwise(expr);
528+ const io = freshIo();
529+ collectBuffers(hoisted, io, span);
530+ const label = `${name} = <${numel(ty)} elem>`;
531+ const kernel = buildKernel({ name, cName, ty, expr: hoisted, span }, io, [...enclosingLoops], label);
532+ const layout = kernelLayout(kernel.buffers.length, kernel.loops.length);
533+ const pipe = await pipeline(kernel.code, label, layout);
534+
535+ const dest = slotFor(cName, name, ty, span);
536+ // WebGPU forbids aliasing a writable binding with a readable one, so an
537+ // in-place update (`u = u + 1`) writes scratch and copies back.
538+ const aliased = kernel.buffers.some((c) => vars.get(resolve(c))?.slot.buffer === dest.buffer);
539+ const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;
540+
541+ const bindGroup = kernelBindGroup(io, kernel.buffers, kernel.loops, target.buffer, layout, span);
542+ spendDispatches(1, span);
543+ ops().push({
544+ kind: 'kernel',
545+ pipeline: pipe,
546+ bindGroup,
547+ dispatch: dispatchFor(kernel.count),
548+ loops: kernel.loops,
549+ label,
550+ copyBack: aliased
551+ ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
552+ : undefined,
553+ });
554+ describeLines.push(`kernel ${label}${aliased ? ' (+copy-back)' : ''}`);
555+ }
556+
557+ /** `B = A`, `X(:)`, `reshape(A, ...)`, vector `x'` — pure views/copies. */
558+ function viewSource(expr: IRExpr): { cName: string; name: string } | null {
559+ if (expr.kind === 'Var' && isTensor(expr.ty)) {
560+ return { cName: expr.cName, name: expr.name };
561+ }
562+ if (expr.kind === 'IndexSlice') {
563+ const base = fullColonBase(expr);
564+ if (base && base.kind === 'Var' && isTensor(base.ty)) {
565+ return { cName: base.cName, name: base.name };
566+ }
567+ return null;
568+ }
569+ if (expr.kind === 'Call' && expr.name === 'reshape' && expr.args.length >= 1 &&
570+ expr.args[0].kind === 'Var' && isTensor(expr.args[0].ty)) {
571+ return { cName: expr.args[0].cName, name: expr.args[0].name };
572+ }
573+ if (expr.kind === 'Unary' && expr.builtin === 'transpose' &&
574+ isVectorish(expr.operand.ty) && expr.operand.kind === 'Var' && isTensor(expr.operand.ty)) {
575+ return { cName: expr.operand.cName, name: expr.operand.name };
576+ }
577+ return null;
578+ }
579+
580+ async function planGemm(
581+ cName: string,
582+ name: string,
583+ ty: NumericType,
584+ expr: IRExpr & { kind: 'Binary' },
585+ span: Span,
586+ ): Promise<void> {
587+ const a = await materialize(expr.left);
588+ const b = await materialize(expr.right);
589+ const [m, k] = shapeOf(expr.left.ty, span);
590+ const [, n] = shapeOf(expr.right.ty, span);
591+ const dest = slotFor(cName, name, ty, span);
592+ const aliased = dest.buffer === a.slot.buffer || dest.buffer === b.slot.buffer;
593+ const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;
594+
595+ const label = `${name} = ${m}x${k} * ${k}x${n}`;
596+ const layout = kernelLayout(2, 0);
597+ const pipe = await pipeline(gemmKernel(m, k, n), label, layout);
598+ const bindGroup = device.createBindGroup({
599+ layout,
600+ entries: [
601+ { binding: 0, resource: { buffer: target.buffer } },
602+ { binding: 1, resource: { buffer: a.slot.buffer } },
603+ { binding: 2, resource: { buffer: b.slot.buffer } },
604+ ],
605+ });
606+ spendDispatches(1, span);
607+ ops().push({
608+ kind: 'kernel',
609+ pipeline: pipe,
610+ bindGroup,
611+ dispatch: gemmDispatch(m, n),
612+ loops: [],
613+ label,
614+ copyBack: aliased
615+ ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
616+ : undefined,
617+ });
618+ describeLines.push(`gemm ${label}`);
619+ }
620+
621+ async function planTranspose(
622+ cName: string,
623+ name: string,
624+ ty: NumericType,
625+ expr: IRExpr & { kind: 'Unary' },
626+ span: Span,
627+ ): Promise<void> {
628+ const src = await materialize(expr.operand);
629+ const [m, n] = shapeOf(expr.operand.ty, span);
630+ const dest = slotFor(cName, name, ty, span);
631+ const aliased = dest.buffer === src.slot.buffer;
632+ const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;
633+
634+ const label = `${name} = (${m}x${n})'`;
635+ const layout = kernelLayout(1, 0);
636+ const pipe = await pipeline(transposeKernel(m, n), label, layout);
637+ const bindGroup = device.createBindGroup({
638+ layout,
639+ entries: [
640+ { binding: 0, resource: { buffer: target.buffer } },
641+ { binding: 1, resource: { buffer: src.slot.buffer } },
642+ ],
643+ });
644+ spendDispatches(1, span);
645+ ops().push({
646+ kind: 'kernel',
647+ pipeline: pipe,
648+ bindGroup,
649+ dispatch: transposeDispatch(m, n),
650+ loops: [],
651+ label,
652+ copyBack: aliased
653+ ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
654+ : undefined,
655+ });
656+ describeLines.push(`transp ${label}`);
657+ }
658+
659+ async function planReduce(
660+ cName: string,
661+ name: string,
662+ ty: NumericType,
663+ call: Call,
664+ span: Span,
665+ ): Promise<void> {
666+ const fn = call.name;
667+ if (fn !== 'dot' && call.args.length !== 1) {
668+ throw new UnsupportedOnGpu(
669+ `'${fn}' supports only the one-argument form here (no dim/'all' ` +
670+ `arguments — use ${fn}(X(:)) for the whole array)`,
671+ span,
672+ );
673+ }
674+
675+ // Build the loader expression: the (fused) element the reduction eats.
676+ let loaderExpr: IRExpr;
677+ let map: MapKind = 'id';
678+ if (fn === 'dot') {
679+ if (call.args.length !== 2) {
680+ throw new UnsupportedOnGpu(`'dot' takes two vectors`, span);
681+ }
682+ loaderExpr = {
683+ kind: 'Binary',
684+ builtin: 'times',
685+ left: call.args[0],
686+ right: call.args[1],
687+ ty: call.args[0].ty,
688+ span,
689+ };
690+ } else {
691+ loaderExpr = call.args[0];
692+ if (fn === 'norm') {
693+ if (!isVectorish(loaderExpr.ty)) {
694+ throw new UnsupportedOnGpu(
695+ `'norm' of a matrix is the spectral norm, which the sandbox does ` +
696+ `not compute; norm(v) for vectors only`,
697+ span,
698+ );
699+ }
700+ map = 'sq';
701+ }
702+ }
703+ // See through X(:) so the loader reads the base buffer directly.
704+ if (loaderExpr.kind === 'IndexSlice') {
705+ const base = fullColonBase(loaderExpr);
706+ if (base) loaderExpr = { ...base, ty: loaderExpr.ty } as IRExpr;
707+ }
708+
709+ const inputTy = loaderExpr.ty as NumericType;
710+ const inCount = numel(inputTy);
711+ const outCount = numel(ty);
712+ const full = outCount === 1;
713+ if (!full) {
714+ const [m, n] = shapeOf(inputTy, span);
715+ const [om, on] = shapeOf(ty, span);
716+ if (om !== 1 || on !== n) {
717+ throw new UnsupportedOnGpu(
718+ `'${fn}' along that dimension is not supported — transpose first, ` +
719+ `or reduce the whole array with ${fn}(X(:))`,
720+ span,
721+ );
722+ }
723+ void m;
724+ }
725+
726+ const combine: Combine =
727+ fn === 'prod' ? 'mul' : fn === 'max' ? 'max' : fn === 'min' ? 'min' : 'add';
728+ const epilogue: Epilogue =
729+ fn === 'mean'
730+ ? { kind: 'scale', by: 1 / (full ? inCount : shapeOf(inputTy, span)[0]) }
731+ : fn === 'norm'
732+ ? { kind: 'sqrt' }
733+ : { kind: 'none' };
734+
735+ // Hoist nested non-elementwise pieces, then emit the fused loader.
736+ loaderExpr = await hoistNonElementwise(
737+ loaderExpr.kind === 'Binary' || loaderExpr.kind === 'Unary' || loaderExpr.kind === 'Call' || loaderExpr.kind === 'IndexSlice'
738+ ? loaderExpr
739+ : loaderExpr,
740+ );
741+ if (isHoistRoot(loaderExpr)) {
742+ const { cName: mc } = await materialize(loaderExpr);
743+ loaderExpr = { kind: 'Var', name: mc, cName: mc, ty: loaderExpr.ty, span };
744+ }
745+ const io = freshIo();
746+ collectBuffers(loaderExpr, io, span);
747+ const loader = emitLoader(loaderExpr, io, [...enclosingLoops]);
748+ const { decls, buffers, loops } = bindingDecls(io);
749+
750+ const dest = slotFor(cName, name, ty, span);
751+ const aliased = buffers.some((c) => vars.get(resolve(c))?.slot.buffer === dest.buffer);
752+ const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;
753+
754+ if (full) {
755+ const numWg = reducePartials(inCount);
756+ const partials = makeSlot(`${name}-partials`, numWg);
757+ const label1 = `${name} = ${fn}(<${inCount} elem>) pass1`;
758+ const layout1 = kernelLayout(buffers.length, loops.length);
759+ const pipe1 = await pipeline(
760+ reduceFullPass1(decls, loader, inCount, numWg, combine, map),
761+ label1,
762+ layout1,
763+ );
764+ const bg1 = kernelBindGroup(io, buffers, loops, partials.buffer, layout1, span);
765+ const label2 = `${name} = ${fn}(...) pass2`;
766+ const layout2 = kernelLayout(1, 0);
767+ const pipe2 = await pipeline(reduceFullPass2(numWg, combine, epilogue), label2, layout2);
768+ const bg2 = device.createBindGroup({
769+ layout: layout2,
770+ entries: [
771+ { binding: 0, resource: { buffer: target.buffer } },
772+ { binding: 1, resource: { buffer: partials.buffer } },
773+ ],
774+ });
775+ spendDispatches(2, span);
776+ ops().push({
777+ kind: 'kernel', pipeline: pipe1, bindGroup: bg1,
778+ dispatch: [numWg, 1], loops, label: label1,
779+ });
780+ ops().push({
781+ kind: 'kernel', pipeline: pipe2, bindGroup: bg2,
782+ dispatch: [1, 1], loops: [], label: label2,
783+ copyBack: aliased
784+ ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
785+ : undefined,
786+ });
787+ describeLines.push(`reduce ${name} = ${fn}(<${inCount} elem>)`);
788+ } else {
789+ const [m, n] = shapeOf(inputTy, span);
790+ if (n > 65535) {
791+ throw new UnsupportedOnGpu(`'${fn}' over ${n} columns exceeds the dispatch limit`, span);
792+ }
793+ const label = `${name} = ${fn}(${m}x${n} by columns)`;
794+ const layout = kernelLayout(buffers.length, loops.length);
795+ const pipe = await pipeline(
796+ reduceColumns(decls, loader, m, combine, map, epilogue),
797+ label,
798+ layout,
799+ );
800+ const bg = kernelBindGroup(io, buffers, loops, target.buffer, layout, span);
801+ spendDispatches(1, span);
802+ ops().push({
803+ kind: 'kernel', pipeline: pipe, bindGroup: bg,
804+ dispatch: [n, 1], loops, label,
805+ copyBack: aliased
806+ ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
807+ : undefined,
808+ });
809+ describeLines.push(`reduce ${label}`);
810+ }
811+ }
812+
813+ function shapeOf(t: Type, span: Span): [number, number] {
814+ if (!isNumeric(t) || !t.shape) {
815+ throw new UnsupportedOnGpu(`shape is not known at compile time`, span);
816+ }
817+ if (t.shape.length !== 2) {
818+ throw new UnsupportedOnGpu(`only 2-D arrays are supported (got ${t.shape.length}-D)`, span);
819+ }
820+ return [t.shape[0], t.shape[1]];
821+ }
822+
823+ // ── Host-op planning ──────────────────────────────────────────────────
824+
825+ const noHostOpsInLoops = (what: string, span: Span): void => {
826+ if (inLoop()) {
827+ throw new UnsupportedOnGpu(
828+ `'${what}' inside a for loop is not supported: the loop body is ` +
829+ `compiled once and replayed on the GPU, so per-iteration host I/O ` +
830+ `has nowhere to run. Hoist it out of the loop (or time the whole loop).`,
831+ span,
832+ );
833+ }
834+ };
835+
836+ /** A ValueRef for a scalar-valued expression the host wants to print. The
837+ * inline pass folds argument temps into the call, so this routinely sees
838+ * whole expressions — they get computed into a 1-element buffer. */
839+ async function scalarRef(e: IRExpr, span: Span): Promise<ValueRef> {
840+ if (isNumeric(e.ty) && isTensor(e.ty)) {
841+ throw new UnsupportedOnGpu(
842+ `printing an array here is not supported (MATLAB would recycle the ` +
843+ `format); print a scalar, or use disp`,
844+ span,
845+ );
846+ }
847+ const exact = exactValue(e);
848+ if (exact !== undefined) return { kind: 'literal', value: exact };
849+ if (e.kind === 'Var') {
850+ if (hostScalars.has(e.cName)) return { kind: 'host', cName: e.cName };
851+ const slot = readSlot(e.cName, e.name, span);
852+ return { kind: 'buffer', slot, count: 1 };
853+ }
854+ if (!isNumeric(e.ty)) {
855+ throw new UnsupportedOnGpu(`only numeric values can be printed`, span);
856+ }
857+ const cName = `%tmp${tempCounter++}`;
858+ await planValue(cName, cName, e.ty, e, span);
859+ return { kind: 'buffer', slot: vars.get(resolve(cName))!.slot, count: 1 };
860+ }
861+
862+ async function planDisp(call: Call, span: Span): Promise<void> {
863+ noHostOpsInLoops('disp', span);
864+ if (call.args.length !== 1) {
865+ throw new UnsupportedOnGpu(`'disp' takes one argument`, span);
866+ }
867+ const a = call.args[0];
868+ if (a.kind === 'StringLit') {
869+ ops().push({ kind: 'emit', parts: [{ kind: 'text', text: a.value + '\n' }] });
870+ return;
871+ }
872+ if (a.kind === 'Var' && chars.has(a.cName)) {
873+ ops().push({ kind: 'emit', parts: [{ kind: 'text', text: chars.get(a.cName)! + '\n' }] });
874+ return;
875+ }
876+ if (!isNumeric(a.ty)) throw new UnsupportedOnGpu(`'disp' argument is not numeric`, span);
877+ if (isTensor(a.ty)) {
878+ if (a.kind !== 'Var') {
879+ throw new UnsupportedOnGpu(`'disp' of an expression — give it a name first`, span);
880+ }
881+ const slot = readSlot(a.cName, a.name, span);
882+ ops().push({
883+ kind: 'display',
884+ label: null,
885+ ref: { kind: 'buffer', slot, count: slot.count },
886+ shape: a.ty.shape ?? [slot.count, 1],
887+ });
888+ return;
889+ }
890+ ops().push({ kind: 'display', label: null, ref: await scalarRef(a, span), shape: [1, 1] });
891+ }
892+
893+ async function planFprintf(call: Call, span: Span): Promise<void> {
894+ noHostOpsInLoops('fprintf', span);
895+ let args = call.args;
896+ // fprintf(1, fmt, ...) — MATLAB's stdout file id.
897+ if (args.length >= 2 && args[0].kind === 'NumLit' && args[0].value === 1) {
898+ args = args.slice(1);
899+ }
900+ if (args.length === 0) throw new UnsupportedOnGpu(`'fprintf' needs a format string`, span);
901+ const fmtArg = args[0];
902+ const fmt =
903+ fmtArg.kind === 'StringLit'
904+ ? fmtArg.value
905+ : fmtArg.kind === 'Var' && chars.has(fmtArg.cName)
906+ ? chars.get(fmtArg.cName)!
907+ : null;
908+ if (fmt === null) {
909+ throw new UnsupportedOnGpu(`'fprintf' format must be a literal string`, span);
910+ }
911+ const parts = parseFormat(fmt, span);
912+ const specs = parts.filter((p) => p.kind === 'spec');
913+ const values = args.slice(1);
914+ if (specs.length !== values.length) {
915+ throw new UnsupportedOnGpu(
916+ `'fprintf' format has ${specs.length} conversion(s) but ${values.length} ` +
917+ `value(s); the sandbox does not recycle the format over arrays`,
918+ span,
919+ );
920+ }
921+ let vi = 0;
922+ const emitParts: EmitPart[] = [];
923+ for (const p of parts) {
924+ if (p.kind === 'text') emitParts.push({ kind: 'text', text: p.text });
925+ else emitParts.push({ kind: 'value', ref: await scalarRef(values[vi++], span), spec: p.spec });
926+ }
927+ ops().push({ kind: 'emit', parts: emitParts });
928+ }
929+
930+ /** printf-format split: text runs (escapes decoded) and % conversions. */
931+ function parseFormat(
932+ fmt: string,
933+ span: Span,
934+ ): ({ kind: 'text'; text: string } | { kind: 'spec'; spec: string })[] {
935+ const out: ({ kind: 'text'; text: string } | { kind: 'spec'; spec: string })[] = [];
936+ let text = '';
937+ for (let i = 0; i < fmt.length; i++) {
938+ const c = fmt[i];
939+ if (c === '\\') {
940+ const n = fmt[i + 1];
941+ if (n === 'n') { text += '\n'; i++; }
942+ else if (n === 't') { text += '\t'; i++; }
943+ else if (n === '\\') { text += '\\'; i++; }
944+ else text += c;
945+ } else if (c === '%') {
946+ if (fmt[i + 1] === '%') { text += '%'; i++; continue; }
947+ const m = /^%[-+ 0#]*\d*(?:\.\d+)?[diufeEgGs]/.exec(fmt.slice(i));
948+ if (!m) {
949+ throw new UnsupportedOnGpu(
950+ `'fprintf': unsupported conversion at "${fmt.slice(i, i + 6)}"`,
951+ span,
952+ );
953+ }
954+ if (text) { out.push({ kind: 'text', text }); text = ''; }
955+ out.push({ kind: 'spec', spec: m[0] });
956+ i += m[0].length - 1;
957+ } else {
958+ text += c;
959+ }
960+ }
961+ if (text) out.push({ kind: 'text', text });
962+ return out;
963+ }
964+
965+ // ── Statement walk ────────────────────────────────────────────────────
966+
967+ async function planStmt(stmt: IRStmt): Promise<void> {
968+ switch (stmt.kind) {
969+ case 'Assign':
970+ return planAssign(stmt);
971+ case 'ExprStmt':
972+ return planExprStmt(stmt);
973+ case 'For':
974+ return planFor(stmt);
975+ default:
976+ throw new UnsupportedOnGpu(
977+ `'${stmtName(stmt.kind)}' is not supported in the sandbox`,
978+ stmt.span,
979+ );
980+ }
981+ }
982+
983+ function stmtName(kind: string): string {
984+ return (
985+ {
986+ If: 'if', While: 'while', Break: 'break', Continue: 'continue',
987+ IndexStore: 'indexed assignment', IndexSliceStore: 'indexed assignment',
988+ MultiAssignCall: 'multiple assignment', MemberStore: 'struct assignment',
989+ CellIndexStore: 'cell assignment',
990+ }[kind] ?? kind
991+ );
992+ }
993+
994+ const isTemp = (name: string): boolean => name.startsWith('_mtoc2_') || name.startsWith('%tmp');
995+
996+ async function planAssign(stmt: Assign): Promise<void> {
997+ // Char values (format strings) ride along on the host. Their type kind
998+ // is numbl's 'Char'/'String', not Numeric.
999+ if (stmt.expr.kind === 'StringLit') {
1000+ chars.set(stmt.cName, stmt.expr.value);
1001+ return;
1002+ }
1003+ if (!isNumeric(stmt.ty)) {
1004+ throw new UnsupportedOnGpu(
1005+ `'${stmt.name}' is not a numeric value (only numbers, strings for ` +
1006+ `printing, and numeric arrays exist in the sandbox)`,
1007+ stmt.span,
1008+ );
1009+ }
1010+
1011+ // tic/toc as values.
1012+ if (stmt.expr.kind === 'Call' && (stmt.expr.name === 'tic' || stmt.expr.name === 'toc')) {
1013+ noHostOpsInLoops(stmt.expr.name, stmt.span);
1014+ const slot = slotFor(stmt.cName, stmt.name, stmt.ty, stmt.span);
1015+ hostScalars.add(stmt.cName);
1016+ if (stmt.expr.name === 'tic') {
1017+ ops().push({ kind: 'tic', assignTo: { cName: stmt.cName, slot } });
1018+ } else {
1019+ const since = tocSince(stmt.expr, stmt.span);
1020+ ops().push({
1021+ kind: 'toc', print: false, sinceCName: since,
1022+ assignTo: { cName: stmt.cName, slot }, seq: ++tocCounter,
1023+ });
1024+ }
1025+ maybeEcho(stmt);
1026+ return;
1027+ }
1028+
1029+ // Exact scalar: folds into every consumer. It only needs real storage if
1030+ // the variable is reassigned elsewhere (later uses then read the buffer).
1031+ if (typeof stmt.ty.exact === 'number') {
1032+ if ((assignCounts.get(stmt.cName) ?? 1) > 1) {
1033+ const slot = slotFor(stmt.cName, stmt.name, stmt.ty, stmt.span);
1034+ ops().push({
1035+ kind: 'write', slot,
1036+ data: new Float32Array([stmt.ty.exact]),
1037+ label: `${stmt.name} = ${stmt.ty.exact}`,
1038+ });
1039+ }
1040+ maybeEcho(stmt);
1041+ return;
1042+ }
1043+ // Exact tensor (a literal like [1 2 3]): upload the data.
1044+ if (stmt.ty.exact instanceof Float64Array) {
1045+ const slot = slotFor(stmt.cName, stmt.name, stmt.ty, stmt.span);
1046+ ops().push({
1047+ kind: 'write', slot,
1048+ data: Float32Array.from(stmt.ty.exact),
1049+ label: `${stmt.name} = <literal ${slot.count} elem>`,
1050+ });
1051+ maybeEcho(stmt);
1052+ return;
1053+ }
1054+ if (stmt.ty.exact !== undefined) {
1055+ throw new UnsupportedOnGpu(`complex values are not supported (f32 backend)`, stmt.span);
1056+ }
1057+
1058+ await planValue(stmt.cName, stmt.name, stmt.ty, stmt.expr, stmt.span);
1059+ maybeEcho(stmt);
1060+ }
1061+
1062+ /** MATLAB-style echo for statements without a trailing semicolon. */
1063+ function maybeEcho(stmt: Assign): void {
1064+ if (isTemp(stmt.name)) return;
1065+ if (!compiled.isEchoed(stmt.span.start)) return;
1066+ noHostOpsInLoops(`echo of '${stmt.name}' (add a semicolon)`, stmt.span);
1067+ if (!isNumeric(stmt.ty)) return;
1068+ const exact = typeof stmt.ty.exact === 'number' ? stmt.ty.exact : undefined;
1069+ const shape = stmt.ty.shape ?? [1, 1];
1070+ if (exact !== undefined) {
1071+ ops().push({
1072+ kind: 'display', label: stmt.name,
1073+ ref: { kind: 'literal', value: exact }, shape,
1074+ });
1075+ return;
1076+ }
1077+ if (hostScalars.has(stmt.cName)) {
1078+ ops().push({
1079+ kind: 'display', label: stmt.name,
1080+ ref: { kind: 'host', cName: stmt.cName }, shape,
1081+ });
1082+ return;
1083+ }
1084+ const v = vars.get(resolve(stmt.cName));
1085+ if (!v) return;
1086+ ops().push({
1087+ kind: 'display', label: stmt.name,
1088+ ref: { kind: 'buffer', slot: v.slot, count: v.slot.count }, shape,
1089+ });
1090+ }
1091+
1092+ function tocSince(call: IRExpr & { kind: 'Call' }, span: Span): string | undefined {
1093+ if (call.args.length === 0) return undefined;
1094+ const a = call.args[0];
1095+ if (a.kind === 'Var' && hostScalars.has(a.cName)) return a.cName;
1096+ throw new UnsupportedOnGpu(
1097+ `'toc(t)' needs a value produced by 't = tic'`,
1098+ span,
1099+ );
1100+ }
1101+
1102+ async function planExprStmt(stmt: ExprStmt): Promise<void> {
1103+ const e = stmt.expr;
1104+ if (e.kind === 'Call') {
1105+ switch (e.name) {
1106+ case 'tic':
1107+ noHostOpsInLoops('tic', stmt.span);
1108+ ops().push({ kind: 'tic' });
1109+ return;
1110+ case 'toc':
1111+ case 'toc_print': // numbl's lowering of a bare `toc` statement
1112+ noHostOpsInLoops('toc', stmt.span);
1113+ ops().push({
1114+ kind: 'toc', print: true,
1115+ sinceCName: tocSince(e, stmt.span), seq: ++tocCounter,
1116+ });
1117+ return;
1118+ case 'disp':
1119+ return planDisp(e, stmt.span);
1120+ case 'fprintf':
1121+ return planFprintf(e, stmt.span);
1122+ case 'rng':
1123+ throw new UnsupportedOnGpu(
1124+ `'rng' is not supported: the sandbox's rand/randn streams are ` +
1125+ `deterministic per run already`,
1126+ stmt.span,
1127+ );
1128+ default:
1129+ break;
1130+ }
1131+ }
1132+ // A bare expression with a value: numbl assigns display-relevant results
1133+ // to `ans` as an Assign, so a leftover ExprStmt is side-effect-free.
1134+ if (e.kind === 'Call') {
1135+ throw new UnsupportedOnGpu(`'${e.name}' is not supported in the sandbox`, stmt.span);
1136+ }
1137+ }
1138+
1139+ async function planFor(stmt: For): Promise<void> {
1140+ const from = exactValue(stmt.start);
1141+ const to = exactValue(stmt.end);
1142+ if (from === undefined || to === undefined) {
1143+ throw new UnsupportedOnGpu(
1144+ `a for loop's bounds must be known when the script is compiled ` +
1145+ `(assign them from literals)`,
1146+ stmt.span,
1147+ );
1148+ }
1149+ const trips = Math.floor((to - from) / stmt.step + 1e-9) + 1;
1150+ if (trips <= 0) return; // never executes
1151+ if (trips > MAX_TRIPS) {
1152+ throw new UnsupportedOnGpu(
1153+ `'for ${stmt.varName}' runs ${trips} iterations, over the sandbox ` +
1154+ `limit of ${MAX_TRIPS}`,
1155+ stmt.span,
1156+ );
1157+ }
1158+
1159+ // One 256-byte uniform slot per iteration: { v: f32, it: u32 }.
1160+ const uniform = device.createBuffer({
1161+ label: `mgpu-loop-${stmt.varName}`,
1162+ size: trips * LV_STRIDE,
1163+ usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
1164+ });
1165+ owned.push(uniform);
1166+ const data = new ArrayBuffer(trips * LV_STRIDE);
1167+ const f32 = new Float32Array(data);
1168+ const u32 = new Uint32Array(data);
1169+ for (let i = 0; i < trips; i++) {
1170+ f32[(i * LV_STRIDE) / 4] = from + i * stmt.step;
1171+ u32[(i * LV_STRIDE) / 4 + 1] = i;
1172+ }
1173+ device.queue.writeBuffer(uniform, 0, data);
1174+ loopUniforms.set(stmt.cVar, uniform);
1175+
1176+ const body: Op[] = [];
1177+ opStack.push(body);
1178+ enclosingLoops.push(stmt.cVar);
1179+ loopTrips.push(trips);
1180+ try {
1181+ for (const s of stmt.body) await planStmt(s);
1182+ } finally {
1183+ loopTrips.pop();
1184+ enclosingLoops.pop();
1185+ opStack.pop();
1186+ }
1187+ ops().push({
1188+ kind: 'loop', cVar: stmt.cVar, trips, uniform, body,
1189+ label: `for ${stmt.varName} = ${from}:${stmt.step}:${to}`,
1190+ });
1191+ describeLines.push(`loop for ${stmt.varName} (${trips} iterations, body above)`);
1192+
1193+ // MATLAB leaves the loop variable holding its final value; make that
1194+ // readable afterwards.
1195+ const finalValue = from + (trips - 1) * stmt.step;
1196+ const slot = slotFor(
1197+ stmt.cVar, stmt.varName,
1198+ { kind: 'Numeric', elem: 'double', isComplex: false,
1199+ dims: [{ kind: 'exact', value: 1 }, { kind: 'exact', value: 1 }],
1200+ shape: [1, 1], sign: 'unknown' },
1201+ stmt.span,
1202+ );
1203+ ops().push({
1204+ kind: 'write', slot, data: new Float32Array([finalValue]),
1205+ label: `${stmt.varName} = ${finalValue} (final)`,
1206+ });
1207+ }
1208+
1209+ for (const stmt of compiled.stmts) {
1210+ await planStmt(stmt);
1211+ }
1212+
1213+ return {
1214+ ops: opStack[0],
1215+ describe: () => describeLines,
1216+ destroy: () => {
1217+ for (const b of owned) b.destroy();
1218+ owned.length = 0;
1219+ },
1220+ };
1221+}
src/mgpu/run.tsadded+306−0View file
@@ -0,0 +1,306 @@
1+/**
2+ * Execute a ScriptPlan on a GPUDevice.
3+ *
4+ * GPU ops stream into a command encoder and are submitted in batches; the
5+ * host ops the script asked for (tic/toc/disp/fprintf/echo) are the only
6+ * synchronization points. `tic` and `toc` flush pending work and await
7+ * `onSubmittedWorkDone`, so `toc` reports wall-clock time for work that has
8+ * actually finished — the same thing MATLAB's synchronous tic/toc measures,
9+ * which is what makes the number comparable when the script is pasted there.
10+ *
11+ * A `loop` op re-encodes its body per iteration, switching the loop-variable
12+ * uniform's dynamic offset; everything still lands in one submit.
13+ */
14+import { WORKGROUP_SIZE } from './wgsl.ts';
15+import type { EmitPart, Op, ScriptPlan, Slot, ValueRef } from './plan.ts';
16+
17+void WORKGROUP_SIZE;
18+
19+export interface TimingSegment {
20+ /** 1-based tic..toc pair index. */
21+ seq: number;
22+ seconds: number;
23+}
24+
25+export interface RunResult {
26+ /** Everything the script printed, in order. */
27+ output: string;
28+ segments: TimingSegment[];
29+ /** Wall time of the whole execution (excluding compilation). */
30+ totalSeconds: number;
31+ error?: string;
32+}
33+
34+const LV_STRIDE = 256;
35+
36+export async function executePlan(
37+ device: GPUDevice,
38+ plan: ScriptPlan,
39+ onOutput?: (text: string) => void,
40+): Promise<RunResult> {
41+ let output = '';
42+ const segments: TimingSegment[] = [];
43+ const print = (text: string): void => {
44+ output += text;
45+ onOutput?.(text);
46+ };
47+
48+ let encoder: GPUCommandEncoder | null = null;
49+ let pass: GPUComputePassEncoder | null = null;
50+ const inEncoder = (): GPUCommandEncoder => {
51+ if (!encoder) encoder = device.createCommandEncoder();
52+ return encoder;
53+ };
54+ const inPass = (): GPUComputePassEncoder => {
55+ if (!pass) pass = inEncoder().beginComputePass();
56+ return pass;
57+ };
58+ const endPass = (): void => {
59+ if (pass) {
60+ pass.end();
61+ pass = null;
62+ }
63+ };
64+ const flush = (): void => {
65+ endPass();
66+ if (encoder) {
67+ device.queue.submit([encoder.finish()]);
68+ encoder = null;
69+ }
70+ };
71+ const sync = async (): Promise<void> => {
72+ flush();
73+ await device.queue.onSubmittedWorkDone();
74+ };
75+
76+ /** Values of tic/toc-produced variables, in seconds. */
77+ const hostVals = new Map<string, number>();
78+ let ticStartMs: number | null = null;
79+
80+ const readBuffer = async (slot: Slot, count: number): Promise<Float32Array> => {
81+ flush();
82+ const bytes = Math.max(4, 4 * count);
83+ const staging = device.createBuffer({
84+ size: bytes,
85+ usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
86+ });
87+ const e = device.createCommandEncoder();
88+ e.copyBufferToBuffer(slot.buffer, 0, staging, 0, bytes);
89+ device.queue.submit([e.finish()]);
90+ await staging.mapAsync(GPUMapMode.READ);
91+ const data = new Float32Array(staging.getMappedRange().slice(0));
92+ staging.unmap();
93+ staging.destroy();
94+ return data.subarray(0, count);
95+ };
96+
97+ const refValue = async (ref: ValueRef): Promise<number> => {
98+ switch (ref.kind) {
99+ case 'literal':
100+ return ref.value;
101+ case 'host':
102+ return hostVals.get(ref.cName) ?? NaN;
103+ case 'buffer':
104+ return (await readBuffer(ref.slot, 1))[0];
105+ }
106+ };
107+
108+ const writeScalar = (slot: Slot, value: number): void => {
109+ device.queue.writeBuffer(slot.buffer, 0, new Float32Array([value]) as Float32Array<ArrayBuffer>);
110+ };
111+
112+ async function execOps(ops: Op[], offsets: Map<string, number>): Promise<void> {
113+ for (const op of ops) {
114+ switch (op.kind) {
115+ case 'kernel': {
116+ const p = inPass();
117+ p.setPipeline(op.pipeline);
118+ if (op.loops.length) {
119+ p.setBindGroup(0, op.bindGroup, op.loops.map((cv) => offsets.get(cv) ?? 0));
120+ } else {
121+ p.setBindGroup(0, op.bindGroup);
122+ }
123+ p.dispatchWorkgroups(op.dispatch[0], op.dispatch[1]);
124+ if (op.copyBack) {
125+ endPass();
126+ inEncoder().copyBufferToBuffer(
127+ op.copyBack.from, 0, op.copyBack.to, 0, op.copyBack.bytes,
128+ );
129+ }
130+ break;
131+ }
132+ case 'copy':
133+ endPass();
134+ inEncoder().copyBufferToBuffer(op.from, 0, op.to, 0, op.bytes);
135+ break;
136+ case 'write':
137+ // Queue writes execute before any later submit; flush pending
138+ // encodes first so ordering matches program order.
139+ flush();
140+ device.queue.writeBuffer(op.slot.buffer, 0, op.data as Float32Array<ArrayBuffer>);
141+ break;
142+ case 'loop': {
143+ for (let it = 0; it < op.trips; it++) {
144+ offsets.set(op.cVar, it * LV_STRIDE);
145+ await execOps(op.body, offsets);
146+ }
147+ offsets.delete(op.cVar);
148+ break;
149+ }
150+ case 'tic': {
151+ await sync();
152+ ticStartMs = performance.now();
153+ if (op.assignTo) {
154+ const seconds = ticStartMs / 1000;
155+ hostVals.set(op.assignTo.cName, seconds);
156+ writeScalar(op.assignTo.slot, seconds);
157+ }
158+ break;
159+ }
160+ case 'toc': {
161+ await sync();
162+ const now = performance.now();
163+ const baseMs =
164+ op.sinceCName !== undefined
165+ ? (hostVals.get(op.sinceCName) ?? 0) * 1000
166+ : ticStartMs;
167+ if (baseMs === null) {
168+ print(`Error: toc without a preceding tic\n`);
169+ break;
170+ }
171+ const seconds = (now - baseMs) / 1000;
172+ segments.push({ seq: op.seq, seconds });
173+ if (op.print) {
174+ print(`Elapsed time is ${seconds.toFixed(6)} seconds.\n`);
175+ }
176+ if (op.assignTo) {
177+ hostVals.set(op.assignTo.cName, seconds);
178+ writeScalar(op.assignTo.slot, seconds);
179+ }
180+ break;
181+ }
182+ case 'emit': {
183+ const parts: string[] = [];
184+ for (const part of op.parts) {
185+ parts.push(await formatPart(part));
186+ }
187+ print(parts.join(''));
188+ break;
189+ }
190+ case 'display': {
191+ print(await formatDisplay(op));
192+ break;
193+ }
194+ }
195+ }
196+ }
197+
198+ async function formatPart(part: EmitPart): Promise<string> {
199+ if (part.kind === 'text') return part.text;
200+ return formatSpec(part.spec, await refValue(part.ref));
201+ }
202+
203+ async function formatDisplay(op: Op & { kind: 'display' }): Promise<string> {
204+ const head = op.label !== null ? `${op.label} =\n\n` : '';
205+ const count = op.shape.reduce((a, b) => a * b, 1);
206+ if (count === 1 || op.ref.kind !== 'buffer') {
207+ const v = await refValue(op.ref);
208+ return `${head} ${formatShort(v)}\n\n`;
209+ }
210+ const [m, n] = op.shape.length === 2 ? op.shape : [count, 1];
211+ if (count > 400) {
212+ // MATLAB would print all of it; that is unreadable in a sandbox pane.
213+ const data = await readBuffer(op.ref.slot, Math.min(count, 4));
214+ const preview = Array.from(data).map(formatShort).join(' ');
215+ return `${head} [${m}x${n}] ${preview} ... (display truncated)\n\n`;
216+ }
217+ const data = await readBuffer(op.ref.slot, count);
218+ const lines: string[] = [];
219+ for (let r = 0; r < m; r++) {
220+ const cells: string[] = [];
221+ for (let c = 0; c < n; c++) {
222+ cells.push(formatShort(data[r + c * m]).padStart(12));
223+ }
224+ lines.push(' ' + cells.join(''));
225+ }
226+ return `${head}${lines.join('\n')}\n\n`;
227+ }
228+
229+ const start = performance.now();
230+ let error: string | undefined;
231+ device.pushErrorScope('out-of-memory');
232+ device.pushErrorScope('validation');
233+ try {
234+ await execOps(plan.ops, new Map());
235+ await sync();
236+ } catch (e) {
237+ error = e instanceof Error ? e.message : String(e);
238+ }
239+ const validation = await device.popErrorScope();
240+ const oom = await device.popErrorScope();
241+ if (!error && validation) error = `GPU validation error: ${validation.message}`;
242+ if (!error && oom) error = `GPU out of memory: ${oom.message}`;
243+ const totalSeconds = (performance.now() - start) / 1000;
244+
245+ return { output, segments, totalSeconds, error };
246+}
247+
248+/** MATLAB `format short`-flavored scalar rendering. */
249+export function formatShort(v: number): string {
250+ if (!Number.isFinite(v)) return v > 0 ? 'Inf' : v < 0 ? '-Inf' : 'NaN';
251+ if (v === 0) return '0';
252+ if (Number.isInteger(v) && Math.abs(v) < 1e10) return String(v);
253+ const a = Math.abs(v);
254+ if (a >= 1e5 || a < 1e-3) return v.toExponential(4);
255+ return v.toFixed(4);
256+}
257+
258+/** One printf-style conversion. */
259+function formatSpec(spec: string, v: number): string {
260+ const m = /^%([-+ 0#]*)(\d*)(?:\.(\d+))?([diufeEgGs])$/.exec(spec);
261+ if (!m) return String(v);
262+ const [, flags, widthS, precS, conv] = m;
263+ const width = widthS ? parseInt(widthS, 10) : 0;
264+ const prec = precS !== undefined ? parseInt(precS, 10) : undefined;
265+ let s: string;
266+ switch (conv) {
267+ case 'd':
268+ case 'i':
269+ case 'u':
270+ s = Number.isInteger(v) ? String(v) : v.toExponential(prec ?? 6);
271+ break;
272+ case 'f':
273+ s = v.toFixed(prec ?? 6);
274+ break;
275+ case 'e':
276+ case 'E': {
277+ s = v.toExponential(prec ?? 6);
278+ if (conv === 'E') s = s.toUpperCase();
279+ break;
280+ }
281+ case 'g':
282+ case 'G': {
283+ const p = prec === undefined || prec === 0 ? 6 : prec;
284+ const a = Math.abs(v);
285+ s = a !== 0 && (a < 1e-5 || a >= 10 ** p)
286+ ? v.toExponential(Math.max(0, p - 1)).replace(/\.?0+e/, 'e')
287+ : String(Number(v.toPrecision(p)));
288+ if (conv === 'G') s = s.toUpperCase();
289+ break;
290+ }
291+ case 's':
292+ s = String(v);
293+ break;
294+ default:
295+ s = String(v);
296+ }
297+ if (flags.includes('+') && v >= 0 && 'dfeg'.includes(conv.toLowerCase())) s = '+' + s;
298+ if (width > s.length) {
299+ s = flags.includes('-')
300+ ? s.padEnd(width)
301+ : flags.includes('0') && !flags.includes('-')
302+ ? (s.startsWith('-') ? '-' + s.slice(1).padStart(width - 1, '0') : s.padStart(width, '0'))
303+ : s.padStart(width);
304+ }
305+ return s;
306+}
src/mgpu/session.tsadded+69−0View file
@@ -0,0 +1,69 @@
1+/**
2+ * The one entry point: compile a MATLAB script and run it on a GPUDevice.
3+ *
4+ * Compilation (numbl lowering + WGSL pipeline builds) is timed separately
5+ * from execution — the execution numbers are the ones to compare with MATLAB.
6+ */
7+import { compileScript } from './compile.ts';
8+import { planScript, type ScriptPlan } from './plan.ts';
9+import { executePlan, type RunResult } from './run.ts';
10+import { asCompileError } from './errors.ts';
11+
12+export interface ScriptRun {
13+ compileSeconds: number;
14+ /** Human-readable op sequence, for the "what did this compile to" pane. */
15+ planDescription: string[];
16+ result: RunResult;
17+}
18+
19+export interface SandboxGpu {
20+ device: GPUDevice;
21+ /** Human-readable adapter description, for the "ran on" line. */
22+ description: string;
23+}
24+
25+export async function requestSandboxDevice(): Promise<SandboxGpu> {
26+ if (!navigator.gpu) {
27+ throw new Error('WebGPU is not available in this browser');
28+ }
29+ const adapter = await navigator.gpu.requestAdapter();
30+ if (!adapter) throw new Error('No WebGPU adapter available');
31+ const want = (name: keyof GPUSupportedLimits): Record<string, number> => ({
32+ [name]: adapter.limits[name] as number,
33+ });
34+ const device = await adapter.requestDevice({
35+ requiredLimits: {
36+ ...want('maxStorageBufferBindingSize'),
37+ ...want('maxBufferSize'),
38+ ...want('maxStorageBuffersPerShaderStage'),
39+ },
40+ });
41+ const info = adapter.info;
42+ const description =
43+ [info?.description || info?.architecture, info?.vendor]
44+ .filter(Boolean)
45+ .join(' — ') || 'unknown adapter';
46+ return { device, description };
47+}
48+
49+export async function runScript(
50+ device: GPUDevice,
51+ source: string,
52+ onOutput?: (text: string) => void,
53+): Promise<ScriptRun> {
54+ const t0 = performance.now();
55+ let plan: ScriptPlan;
56+ try {
57+ const compiled = compileScript(source);
58+ plan = await planScript(device, compiled);
59+ } catch (e) {
60+ throw asCompileError(e);
61+ }
62+ const compileSeconds = (performance.now() - t0) / 1000;
63+ try {
64+ const result = await executePlan(device, plan, onOutput);
65+ return { compileSeconds, planDescription: plan.describe(), result };
66+ } finally {
67+ plan.destroy();
68+ }
69+}
src/mgpu/wgsl.tsadded+734−0View file
@@ -0,0 +1,734 @@
1+/**
2+ * IR expression tree -> one WGSL compute kernel.
3+ *
4+ * The elementwise core follows turing-surface's emitter: for an `Assign` whose
5+ * right-hand side is purely element-wise over operands of the target's shape,
6+ * emit a single kernel that computes one output element per invocation.
7+ * Because numbl's inline pass has already folded the ANF temps back together,
8+ * one source line of MATLAB becomes one kernel.
9+ *
10+ * The sandbox extends it with:
11+ * - comparisons and eager logicals (`<`, `&`, `~`, ...), carried as f32 0/1;
12+ * - inline *generators* — `rand`, `randn`, `linspace`, ranges, `zeros`,
13+ * `ones`, `eye` — evaluated per element from the linear index, so
14+ * `x = 2*rand(n,1) - 1` is one kernel and touches no other buffer;
15+ * - runtime scalars read from 1-element storage buffers (`in3[0]`);
16+ * - loop variables read from a per-loop dynamic-offset uniform, so a `for`
17+ * body compiles once and replays.
18+ *
19+ * Everything is f32 — WGSL has no f64. Arrays are column-major linear buffers,
20+ * matching MATLAB, so `A(:)` and `reshape` are views of the same buffer.
21+ */
22+import { isMultiElement } from 'numbl-src/numbl-core/jit/lowering/types.ts';
23+import type {
24+ IRExpr,
25+ Assign,
26+ IndexSlice,
27+ Span,
28+} from 'numbl-src/numbl-core/jit/lowering/ir.ts';
29+import type { NumericType, Type } from 'numbl-src/numbl-core/jit/lowering/types.ts';
30+import { UnsupportedOnGpu } from './errors.ts';
31+
32+export const WORKGROUP_SIZE = 64;
33+
34+const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
35+const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
36+export const numel = (t: NumericType): number =>
37+ (t.shape ?? []).reduce((a, b) => a * b, 1);
38+
39+/** The compile-time value of a scalar expression, if it has one. */
40+export const exactValue = (e: IRExpr): number | undefined => {
41+ if (isNumeric(e.ty) && typeof e.ty.exact === 'number') return e.ty.exact;
42+ return e.kind === 'NumLit' ? e.value : undefined;
43+};
44+
45+/** Element-wise binary builtins -> WGSL infix operator. */
46+const BINARY_OPS: Record<string, string> = {
47+ plus: '+',
48+ minus: '-',
49+ times: '*',
50+ rdivide: '/',
51+ // Degenerate to element-wise when at least one side is a scalar; the
52+ // both-tensor (true matrix) case never reaches here — the planner routes it
53+ // to the GEMM kernel (mtimes) or rejects it (mrdivide).
54+ mtimes: '*',
55+ mrdivide: '/',
56+};
57+
58+/** Comparison builtins -> WGSL comparison; result carried as f32 0/1. */
59+const COMPARE_OPS: Record<string, string> = {
60+ lt: '<',
61+ le: '<=',
62+ gt: '>',
63+ ge: '>=',
64+ eq: '==',
65+ ne: '!=',
66+};
67+
68+/** Element-wise unary builtins -> WGSL prefix operator. */
69+const UNARY_OPS: Record<string, string> = { uminus: '-', uplus: '+' };
70+
71+/** Element-wise builtin calls -> WGSL builtin of the same arity. */
72+const CALL_FNS: Record<string, string> = {
73+ abs: 'abs',
74+ acos: 'acos',
75+ asin: 'asin',
76+ atan: 'atan',
77+ atan2: 'atan2',
78+ ceil: 'ceil',
79+ cos: 'cos',
80+ cosh: 'cosh',
81+ exp: 'exp',
82+ fix: 'trunc',
83+ floor: 'floor',
84+ log: 'log',
85+ log2: 'log2',
86+ round: 'round',
87+ sign: 'sign',
88+ sin: 'sin',
89+ sinh: 'sinh',
90+ sqrt: 'sqrt',
91+ tan: 'tan',
92+ tanh: 'tanh',
93+};
94+
95+/** Reductions the planner materializes before a kernel is built. Their 1-arg
96+ * (and for dot, 2-arg) tensor forms never reach the elementwise emitter. */
97+export const REDUCTIONS = new Set([
98+ 'sum', 'mean', 'prod', 'max', 'min', 'norm', 'dot',
99+]);
100+
101+/** WGSL f32 literal. Must always carry a decimal point or exponent, or WGSL
102+ * infers AbstractInt and rejects the mixed-type arithmetic. */
103+function f32Lit(v: number): string {
104+ if (!Number.isFinite(v)) {
105+ // WGSL has no NaN/Inf literal, and 0.0/0.0 is a const-eval error;
106+ // bitcast the IEEE pattern instead.
107+ if (Number.isNaN(v)) return 'bitcast<f32>(0x7fc00000u)';
108+ return v > 0 ? 'bitcast<f32>(0x7f800000u)' : 'bitcast<f32>(0xff800000u)';
109+ }
110+ return Number.isInteger(v) && Math.abs(v) < 1e21
111+ ? `${v}.0`
112+ : String(v).includes('e')
113+ ? `${v}f`
114+ : String(v);
115+}
116+
117+/** `A(:)` — the only IndexSlice this backend executes. Returns the base Var
118+ * expression, which reads the same buffer at the same linear index (a
119+ * column-major flatten is the identity on the linear buffer). */
120+export function fullColonBase(e: IndexSlice): IRExpr | null {
121+ if (e.index.length !== 1 || e.index[0].kind !== 'Colon') return null;
122+ return e.base;
123+}
124+
125+/** How operands are read inside a kernel. */
126+export interface KernelInputs {
127+ /** cName -> storage binding slot, for tensors AND runtime scalars (a
128+ * runtime scalar is a 1-element buffer, read as `inN[0]`). */
129+ buffers: Map<string, number>;
130+ /** cVar -> dense per-kernel index of enclosing loop variables, each bound
131+ * as a dynamic-offset uniform (`lvK`). */
132+ loopVars: Map<string, number>;
133+ /** Distinct-per-call-site seeds for `rand`/`randn`. Global to the plan, so
134+ * two kernels never share a stream. */
135+ nextSeed: () => number;
136+}
137+
138+interface Ctx {
139+ io: KernelInputs;
140+ /** Emitted-helper flags, gathered during emission. */
141+ usesHash: boolean;
142+ usedPows: Set<number>;
143+ usesMod: boolean;
144+ usesRem: boolean;
145+}
146+
147+/** Mix every enclosing loop's iteration counter into a hash lane, so a
148+ * generator inside a replayed loop draws fresh values each iteration. */
149+function seedExpr(seed: number, ctx: Ctx): string {
150+ let s = `${seed >>> 0}u`;
151+ for (const [, k] of ctx.io.loopVars) {
152+ s += ` ^ (lv${k}.it * ${[2654435761, 2246822519, 3266489917, 668265263][k % 4]}u)`;
153+ }
154+ return s;
155+}
156+
157+/** Is `t` a value with more than one element? (logical/double both count) */
158+const multi = (t: Type): boolean => isTensor(t);
159+
160+/**
161+ * Emit the per-element WGSL expression for `e`. `i` is the element index
162+ * variable in scope. Comparisons/logicals produce f32 0/1 so any consumer
163+ * can treat them as numbers, exactly like MATLAB's logicals.
164+ */
165+function emitExpr(e: IRExpr, ctx: Ctx): string {
166+ // Anything numbl constant-folded (pi, 2*pi, n-1, ...) is a literal, no
167+ // matter what expression kind computed it.
168+ const exact = exactValue(e);
169+ if (exact !== undefined) return f32Lit(exact);
170+ const io = ctx.io;
171+ switch (e.kind) {
172+ case 'NumLit':
173+ return f32Lit(e.value);
174+
175+ case 'Var': {
176+ const lv = io.loopVars.get(e.cName);
177+ if (lv !== undefined) return `lv${lv}.v`;
178+ if (isNumeric(e.ty) && typeof e.ty.exact === 'number') {
179+ return f32Lit(e.ty.exact);
180+ }
181+ const slot = io.buffers.get(e.cName);
182+ if (slot === undefined) {
183+ throw new UnsupportedOnGpu(`no buffer bound for '${e.name}'`, e.span);
184+ }
185+ return multi(e.ty) ? `in${slot}[i]` : `in${slot}[0]`;
186+ }
187+
188+ case 'IndexSlice': {
189+ const base = fullColonBase(e);
190+ if (!base || base.kind !== 'Var') {
191+ throw new UnsupportedOnGpu(
192+ `only the full linearization 'X(:)' of a variable is supported — ` +
193+ `general indexing/slicing is not implemented on the GPU yet`,
194+ e.span,
195+ );
196+ }
197+ return emitExpr(base, ctx);
198+ }
199+
200+ case 'MakeRange': {
201+ // start + i*step; the count is already fixed in the node's type.
202+ const start = emitScalar(e.start, ctx, `range start`);
203+ const stepV = exactValue(e.step);
204+ if (stepV === undefined) {
205+ throw new UnsupportedOnGpu(`a range's step must be a compile-time value`, e.span);
206+ }
207+ return `(${start} + f32(i) * ${f32Lit(stepV)})`;
208+ }
209+
210+ case 'Binary': {
211+ if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
212+ multi(e.left.ty) && multi(e.right.ty)) {
213+ // The planner materializes tensor mtimes into a GEMM before building
214+ // the kernel; reaching here means it could not (mrdivide) or a
215+ // planner bug (mtimes).
216+ throw new UnsupportedOnGpu(
217+ e.builtin === 'mrdivide'
218+ ? `matrix '/' (mrdivide) is not supported; use './' or a factorization`
219+ : `internal: tensor '*' was not materialized as a GEMM`,
220+ e.span,
221+ );
222+ }
223+ if (e.builtin === 'power' || e.builtin === 'mpower') {
224+ return emitPower(e.left, e.right, ctx, e.span);
225+ }
226+ const cmp = COMPARE_OPS[e.builtin];
227+ if (cmp) {
228+ return `select(0.0, 1.0, ${emitExpr(e.left, ctx)} ${cmp} ${emitExpr(e.right, ctx)})`;
229+ }
230+ if (e.builtin === 'and' || e.builtin === 'or' || e.builtin === 'andand' || e.builtin === 'oror') {
231+ const op = e.builtin === 'and' || e.builtin === 'andand' ? '&&' : '||';
232+ return `select(0.0, 1.0, (${emitExpr(e.left, ctx)} != 0.0) ${op} (${emitExpr(e.right, ctx)} != 0.0))`;
233+ }
234+ const op = BINARY_OPS[e.builtin];
235+ if (!op) {
236+ throw new UnsupportedOnGpu(`operator '${e.builtin}' is not supported`, e.span);
237+ }
238+ return `(${emitExpr(e.left, ctx)} ${op} ${emitExpr(e.right, ctx)})`;
239+ }
240+
241+ case 'Unary': {
242+ if (e.builtin === 'not') {
243+ return `select(1.0, 0.0, ${emitExpr(e.operand, ctx)} != 0.0)`;
244+ }
245+ if (e.builtin === 'transpose') {
246+ // A vector transpose changes orientation only — the linear buffer is
247+ // identical. A matrix transpose is materialized by the planner.
248+ if (isVectorish(e.operand.ty)) return emitExpr(e.operand, ctx);
249+ throw new UnsupportedOnGpu(
250+ `internal: matrix transpose was not materialized`,
251+ e.span,
252+ );
253+ }
254+ const op = UNARY_OPS[e.builtin];
255+ if (!op) {
256+ throw new UnsupportedOnGpu(`unary '${e.builtin}' is not supported`, e.span);
257+ }
258+ return `(${op}${emitExpr(e.operand, ctx)})`;
259+ }
260+
261+ case 'Call':
262+ return emitCall(e, ctx);
263+
264+ default:
265+ throw new UnsupportedOnGpu(`'${e.kind}' is not supported on the GPU`, e.span);
266+ }
267+}
268+
269+const isVectorish = (t: Type): boolean =>
270+ isNumeric(t) && (t.shape ?? []).filter((d) => d !== 1).length <= 1;
271+
272+/** A scalar-valued subexpression (range endpoints, linspace args). */
273+function emitScalar(e: IRExpr, ctx: Ctx, what: string): string {
274+ if (multi(e.ty)) {
275+ throw new UnsupportedOnGpu(`${what} must be a scalar`, e.span);
276+ }
277+ return emitExpr(e, ctx);
278+}
279+
280+function emitCall(e: IRExpr & { kind: 'Call' }, ctx: Ctx): string {
281+ switch (e.name) {
282+ case 'zeros':
283+ return '0.0';
284+ case 'ones':
285+ return '1.0';
286+ case 'eye': {
287+ const t = e.ty;
288+ const m = isNumeric(t) && t.shape ? t.shape[0] : undefined;
289+ if (m === undefined) {
290+ throw new UnsupportedOnGpu(`'eye' needs a compile-time size`, e.span);
291+ }
292+ return `select(0.0, 1.0, (i % ${m}u) == (i / ${m}u))`;
293+ }
294+ case 'rand': {
295+ ctx.usesHash = true;
296+ return `rand01(i, ${seedExpr(ctx.io.nextSeed(), ctx)})`;
297+ }
298+ case 'randn': {
299+ ctx.usesHash = true;
300+ const a = seedExpr(ctx.io.nextSeed(), ctx);
301+ const b = seedExpr(ctx.io.nextSeed(), ctx);
302+ // Box–Muller; rand01 returns (0,1) so the log is finite.
303+ return `(sqrt(-2.0 * log(rand01(i, ${a}))) * cos(6.283185307179586 * rand01(i, ${b})))`;
304+ }
305+ case 'linspace': {
306+ if (e.args.length !== 3) {
307+ throw new UnsupportedOnGpu(`'linspace' needs 3 arguments here`, e.span);
308+ }
309+ const n = exactValue(e.args[2]);
310+ if (n === undefined) {
311+ throw new UnsupportedOnGpu(`'linspace' count must be a compile-time value`, e.span);
312+ }
313+ const a = emitScalar(e.args[0], ctx, `'linspace' start`);
314+ const b = emitScalar(e.args[1], ctx, `'linspace' end`);
315+ if (n <= 1) return b; // MATLAB: linspace(a, b, 1) == b
316+ return `(${a} + f32(i) * ((${b} - ${a}) * ${f32Lit(1 / (n - 1))}))`;
317+ }
318+ case 'mod':
319+ ctx.usesMod = true;
320+ return `mod_m(${emitExpr(e.args[0], ctx)}, ${emitExpr(e.args[1], ctx)})`;
321+ case 'rem':
322+ ctx.usesRem = true;
323+ return `rem_m(${emitExpr(e.args[0], ctx)}, ${emitExpr(e.args[1], ctx)})`;
324+ case 'max':
325+ case 'min': {
326+ // Two-arg elementwise form; the one-arg reduction never reaches here.
327+ if (e.args.length !== 2) {
328+ throw new UnsupportedOnGpu(`internal: '${e.name}' reduction was not materialized`, e.span);
329+ }
330+ return `${e.name}(${emitExpr(e.args[0], ctx)}, ${emitExpr(e.args[1], ctx)})`;
331+ }
332+ case 'xor':
333+ return `select(0.0, 1.0, (${emitExpr(e.args[0], ctx)} != 0.0) != (${emitExpr(e.args[1], ctx)} != 0.0))`;
334+ case 'double':
335+ case 'logical':
336+ // Representation is f32 either way.
337+ return emitExpr(e.args[0], ctx);
338+ default: {
339+ const fn = CALL_FNS[e.name];
340+ if (!fn) {
341+ const isUserFunction = e.cName !== e.name;
342+ throw new UnsupportedOnGpu(
343+ isUserFunction
344+ ? `'${e.name}' is a user-defined function — not supported in the sandbox; inline it`
345+ : REDUCTIONS.has(e.name)
346+ ? `internal: '${e.name}' reduction was not materialized`
347+ : `'${e.name}' cannot be evaluated element-wise on the GPU`,
348+ e.span,
349+ );
350+ }
351+ return `${fn}(${e.args.map((a) => emitExpr(a, ctx)).join(', ')})`;
352+ }
353+ }
354+}
355+
356+/**
357+ * `x.^k`. WGSL's `pow` is undefined for a negative base, so expand literal
358+ * integer exponents into repeated multiplication — which is also what makes
359+ * `u.^2` free. Non-integer exponents fall through to `pow`, defined only for
360+ * a non-negative base (as in MATLAB, where a negative base goes complex —
361+ * here it is NaN, and the result cross-check will show it).
362+ */
363+function emitPower(base: IRExpr, exponent: IRExpr, ctx: Ctx, span: Span): string {
364+ const k = exactValue(exponent);
365+ const b = emitExpr(base, ctx);
366+ if (k !== undefined && Number.isInteger(k) && k >= 0 && k <= 16) {
367+ if (k === 0) return '1.0';
368+ ctx.usedPows.add(k);
369+ return `pow_i${k}(${b})`;
370+ }
371+ if (k !== undefined && Number.isInteger(k) && k < 0 && k >= -16) {
372+ ctx.usedPows.add(-k);
373+ return `(1.0 / pow_i${-k}(${b}))`;
374+ }
375+ return `pow(${b}, ${emitExpr(exponent, ctx)})`;
376+}
377+
378+/** Fixed-exponent power helpers, emitted only when used. */
379+function powHelpers(used: Set<number>): string {
380+ const out: string[] = [];
381+ for (const k of [...used].sort((a, b) => a - b)) {
382+ const body = k === 1 ? 'x' : `x${' * x'.repeat(k - 1)}`;
383+ out.push(`fn pow_i${k}(x: f32) -> f32 { return ${body}; }`);
384+ }
385+ return out.join('\n');
386+}
387+
388+/** PCG-style hash -> (0,1). Counter-based: a call site's stream is a pure
389+ * function of (element index, seed), so runs are reproducible. */
390+const HASH_HELPERS = `
391+fn hash_u(x0: u32) -> u32 {
392+ var x = x0 * 747796405u + 2891336453u;
393+ x = ((x >> ((x >> 28u) + 4u)) ^ x) * 277803737u;
394+ return (x >> 22u) ^ x;
395+}
396+fn rand01(i: u32, seed: u32) -> f32 {
397+ return (f32(hash_u(i ^ (seed * 2654435769u)) & 0x00FFFFFFu) + 0.5) * (1.0 / 16777216.0);
398+}`;
399+
400+const MOD_HELPER = `
401+fn mod_m(a: f32, b: f32) -> f32 { return select(a - b * floor(a / b), a, b == 0.0); }`;
402+const REM_HELPER = `
403+fn rem_m(a: f32, b: f32) -> f32 { return select(a - b * trunc(a / b), a, b == 0.0); }`;
404+
405+/**
406+ * Reject implicit expansion (broadcasting).
407+ *
408+ * numbl's lowering permits it — `2x4096 .* 1x4096` lowers happily with MATLAB
409+ * expansion semantics — but a kernel that walks one linear index across every
410+ * operand would quietly compute the wrong thing. So every multi-element
411+ * operand must have exactly the target's number of elements (a flattened or
412+ * transposed vector reads the same linear buffer, so only numel must match).
413+ */
414+function checkShapes(e: IRExpr, target: NumericType, name: string): void {
415+ const want = numel(target);
416+ const walk = (x: IRExpr): void => {
417+ if (isNumeric(x.ty) && isMultiElement(x.ty)) {
418+ const got = x.ty.shape ? numel(x.ty) : undefined;
419+ if (got !== want) {
420+ throw new UnsupportedOnGpu(
421+ `'${name}' would need implicit expansion: an operand is ` +
422+ `${x.ty.shape?.join('x') ?? 'dynamic'} but the result has ${want} ` +
423+ `elements. Expand it explicitly (the GPU kernel walks one linear ` +
424+ `index across every operand).`,
425+ x.span,
426+ );
427+ }
428+ // Same-numel vectors of different orientation share a linear layout;
429+ // same-numel *matrices* of different shape do not (transpose is not a
430+ // relayout numbl would insert silently, so shapes agree here).
431+ }
432+ switch (x.kind) {
433+ case 'Binary':
434+ walk(x.left);
435+ walk(x.right);
436+ return;
437+ case 'Unary':
438+ walk(x.operand);
439+ return;
440+ case 'IndexSlice':
441+ return; // the base reads through the slice's own (checked) type
442+ case 'Call':
443+ // A generator's arguments are sizes/endpoints, not per-element data.
444+ if (!['zeros', 'ones', 'eye', 'rand', 'randn', 'linspace'].includes(x.name)) {
445+ x.args.forEach(walk);
446+ }
447+ return;
448+ default:
449+ return;
450+ }
451+ };
452+ walk(e);
453+}
454+
455+export interface Kernel {
456+ code: string;
457+ /** Number of output elements. */
458+ count: number;
459+ label: string;
460+ /** Buffer operand cNames in binding order (bindings 1..n). */
461+ buffers: string[];
462+ /** Loop-variable cVars in binding order (after the buffers). */
463+ loops: string[];
464+}
465+
466+/** True if the expression draws random numbers anywhere. */
467+function usesRandom(e: IRExpr): boolean {
468+ let found = false;
469+ const walk = (x: IRExpr): void => {
470+ if (found) return;
471+ switch (x.kind) {
472+ case 'Call':
473+ if (x.name === 'rand' || x.name === 'randn') found = true;
474+ else x.args.forEach(walk);
475+ return;
476+ case 'Binary':
477+ walk(x.left);
478+ walk(x.right);
479+ return;
480+ case 'Unary':
481+ walk(x.operand);
482+ return;
483+ case 'IndexSlice':
484+ walk(x.base);
485+ return;
486+ default:
487+ return;
488+ }
489+ };
490+ walk(e);
491+ return found;
492+}
493+
494+/** Every buffer-backed variable the expression reads (tensors and runtime
495+ * scalars), and every loop variable. */
496+export function collectReads(
497+ e: IRExpr,
498+ isLoopVar: (cName: string) => boolean,
499+ isExact: (x: IRExpr) => boolean,
500+ visitBuffer: (cName: string) => void,
501+ visitLoop: (cName: string) => void,
502+): void {
503+ const walk = (x: IRExpr): void => {
504+ switch (x.kind) {
505+ case 'Var':
506+ if (isLoopVar(x.cName)) visitLoop(x.cName);
507+ else if (!isExact(x)) visitBuffer(x.cName);
508+ return;
509+ case 'Binary':
510+ walk(x.left);
511+ walk(x.right);
512+ return;
513+ case 'Unary':
514+ walk(x.operand);
515+ return;
516+ case 'IndexSlice':
517+ walk(x.base);
518+ return;
519+ case 'MakeRange':
520+ walk(x.start);
521+ walk(x.step);
522+ return;
523+ case 'Call':
524+ if (!['zeros', 'ones', 'eye', 'rand', 'randn'].includes(x.name)) {
525+ x.args.forEach(walk);
526+ }
527+ return;
528+ default:
529+ return;
530+ }
531+ };
532+ walk(e);
533+}
534+
535+/**
536+ * Can this expression live inside one fused GPU kernel?
537+ *
538+ * Wider than numbl's own `isPureElementwiseExpr`: the WGSL emitter fuses
539+ * transcendental calls, comparisons/logicals, generators, ranges, `X(:)` and
540+ * vector transposes, all of which numbl's C-side pass declines. The sandbox's
541+ * fuse pass uses this to fold the temps numbl's inline pass left behind.
542+ */
543+export function isGpuFusableExpr(e: IRExpr): boolean {
544+ if (exactValue(e) !== undefined) return true;
545+ switch (e.kind) {
546+ case 'NumLit':
547+ return true;
548+ case 'Var':
549+ return isNumeric(e.ty);
550+ case 'Binary': {
551+ if ((e.builtin === 'mtimes' || e.builtin === 'mrdivide') &&
552+ multi(e.left.ty) && multi(e.right.ty)) {
553+ return false;
554+ }
555+ if (e.builtin === 'mpower' && multi(e.left.ty)) return false; // matrix power
556+ const known =
557+ e.builtin in BINARY_OPS || e.builtin in COMPARE_OPS ||
558+ ['and', 'or', 'andand', 'oror', 'power', 'mpower'].includes(e.builtin);
559+ return known && isGpuFusableExpr(e.left) && isGpuFusableExpr(e.right);
560+ }
561+ case 'Unary': {
562+ if (e.builtin === 'transpose') {
563+ return isVectorish(e.operand.ty) && isGpuFusableExpr(e.operand);
564+ }
565+ return (e.builtin in UNARY_OPS || e.builtin === 'not') && isGpuFusableExpr(e.operand);
566+ }
567+ case 'Call': {
568+ if (['zeros', 'ones', 'eye', 'rand', 'randn'].includes(e.name)) return true;
569+ if (e.name === 'linspace') return e.args.length === 3;
570+ if (['mod', 'rem', 'xor', 'atan2'].includes(e.name)) {
571+ return e.args.every(isGpuFusableExpr);
572+ }
573+ if ((e.name === 'max' || e.name === 'min') && e.args.length === 2) {
574+ return e.args.every(isGpuFusableExpr);
575+ }
576+ if (e.name === 'double' || e.name === 'logical') {
577+ return e.args.length === 1 && isGpuFusableExpr(e.args[0]);
578+ }
579+ return e.name in CALL_FNS && e.args.every(isGpuFusableExpr);
580+ }
581+ case 'IndexSlice': {
582+ const base = fullColonBase(e);
583+ return !!base && base.kind === 'Var';
584+ }
585+ case 'MakeRange':
586+ return exactValue(e.step) !== undefined;
587+ default:
588+ return false;
589+ }
590+}
591+
592+/** A fused elementwise subexpression, for embedding inside a non-elementwise
593+ * kernel (a reduction's per-element load). `body` reads element `i`. */
594+export interface FusedLoader {
595+ body: string;
596+ helpers: string;
597+}
598+
599+/**
600+ * Emit `e` as a per-element load for a reduction kernel. Same contract as
601+ * `buildKernel`: `io.buffers` maps operands to binding slots; loop variables
602+ * are bound on demand (all of them, if the expression draws random numbers).
603+ */
604+export function emitLoader(
605+ e: IRExpr,
606+ io: KernelInputs,
607+ enclosingLoops: string[],
608+): FusedLoader {
609+ if (usesRandom(e)) {
610+ for (const cVar of enclosingLoops) {
611+ if (!io.loopVars.has(cVar)) io.loopVars.set(cVar, io.loopVars.size);
612+ }
613+ }
614+ const ctx: Ctx = {
615+ io,
616+ usesHash: false,
617+ usedPows: new Set(),
618+ usesMod: false,
619+ usesRem: false,
620+ };
621+ const body = emitExpr(e, ctx);
622+ const helpers = [
623+ ctx.usesHash ? HASH_HELPERS : '',
624+ ctx.usesMod ? MOD_HELPER : '',
625+ ctx.usesRem ? REM_HELPER : '',
626+ powHelpers(ctx.usedPows),
627+ ].filter(Boolean).join('\n');
628+ return { body, helpers };
629+}
630+
631+/** Binding declarations shared by every kernel shape: output at 0, operand
632+ * buffers after it, loop-variable uniforms after those. */
633+export function bindingDecls(io: KernelInputs): {
634+ decls: string[];
635+ buffers: string[];
636+ loops: string[];
637+} {
638+ const decls = [`@group(0) @binding(0) var<storage, read_write> out: array<f32>;`];
639+ const buffers: string[] = [];
640+ for (const [cName, slot] of io.buffers) {
641+ buffers[slot] = cName;
642+ decls.push(
643+ `@group(0) @binding(${slot + 1}) var<storage, read> in${slot}: array<f32>;`,
644+ );
645+ }
646+ const loops: string[] = [];
647+ if (io.loopVars.size) {
648+ decls.push(`struct Lv { v: f32, it: u32 }`);
649+ for (const [cVar, k] of io.loopVars) {
650+ loops[k] = cVar;
651+ decls.push(
652+ `@group(0) @binding(${io.buffers.size + 1 + k}) var<uniform> lv${k}: Lv;`,
653+ );
654+ }
655+ }
656+ return { decls, buffers, loops };
657+}
658+
659+/**
660+ * Build the fused elementwise kernel for one `Assign`. `io.buffers` must
661+ * already map every buffer operand to a binding slot; the output is binding 0
662+ * and loop-variable uniforms follow the last input.
663+ *
664+ * `enclosingLoops` lists the cVars of the loops this statement sits inside
665+ * (outermost first). Loop variables the expression reads are bound; if the
666+ * expression draws random numbers, ALL enclosing loop counters are bound and
667+ * mixed into the stream so each replayed iteration draws fresh values.
668+ */
669+export function buildKernel(
670+ stmt: Pick<Assign, 'name' | 'cName' | 'ty' | 'expr' | 'span'>,
671+ io: KernelInputs,
672+ enclosingLoops: string[],
673+ label: string,
674+): Kernel {
675+ if (!isNumeric(stmt.ty)) {
676+ throw new UnsupportedOnGpu(`'${stmt.name}' is not a numeric array`, stmt.span);
677+ }
678+ if (stmt.ty.isComplex) {
679+ throw new UnsupportedOnGpu(
680+ `'${stmt.name}' is complex; the GPU backend is real-only (f32)`,
681+ stmt.span,
682+ );
683+ }
684+ const count = numel(stmt.ty);
685+ if (isMultiElement(stmt.ty)) checkShapes(stmt.expr, stmt.ty, stmt.name);
686+
687+ // Random draws need every enclosing loop counter; make sure they are bound
688+ // before emission asks for them.
689+ if (usesRandom(stmt.expr)) {
690+ for (const cVar of enclosingLoops) {
691+ if (!io.loopVars.has(cVar)) io.loopVars.set(cVar, io.loopVars.size);
692+ }
693+ }
694+
695+ const ctx: Ctx = {
696+ io,
697+ usesHash: false,
698+ usedPows: new Set(),
699+ usesMod: false,
700+ usesRem: false,
701+ };
702+ const body = emitExpr(stmt.expr, ctx);
703+ const { decls, buffers, loops } = bindingDecls(io);
704+
705+ const helpers = [
706+ ctx.usesHash ? HASH_HELPERS : '',
707+ ctx.usesMod ? MOD_HELPER : '',
708+ ctx.usesRem ? REM_HELPER : '',
709+ powHelpers(ctx.usedPows),
710+ ].filter(Boolean).join('\n');
711+
712+ // Dispatch is 2-D so counts past 65535 workgroups still fit: x rows of
713+ // ELEMENTS_PER_ROW elements each. dispatchFor() picks matching counts.
714+ const code = `${decls.join('\n')}
715+${helpers}
716+@compute @workgroup_size(${WORKGROUP_SIZE})
717+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
718+ let i = gid.x + gid.y * ${ELEMENTS_PER_ROW}u;
719+ if (i >= ${count}u) { return; }
720+ out[i] = ${body};
721+}
722+`;
723+ return { code, count, label, buffers, loops };
724+}
725+
726+/** Elements covered by one row of the 2-D elementwise dispatch. */
727+export const ELEMENTS_PER_ROW = 32768 * WORKGROUP_SIZE;
728+
729+/** Workgroup counts for an elementwise dispatch over `count` elements. */
730+export function dispatchFor(count: number): [number, number] {
731+ const rows = Math.ceil(count / ELEMENTS_PER_ROW);
732+ const x = rows === 1 ? Math.ceil(count / WORKGROUP_SIZE) : 32768;
733+ return [x, rows];
734+}
src/raw.d.tsadded+4−0View file
@@ -0,0 +1,4 @@
1+declare module '*.m?raw' {
2+ const source: string;
3+ export default source;
4+}
test.htmladded+12−0View file
@@ -0,0 +1,12 @@
1+<!doctype html>
2+<html lang="en">
3+ <head>
4+ <meta charset="utf-8" />
5+ <title>math-webgpu-sandbox validation</title>
6+ </head>
7+ <body>
8+ <h1>math-webgpu-sandbox validation suite</h1>
9+ <pre id="log">starting…</pre>
10+ <script type="module" src="/test/test-page.ts"></script>
11+ </body>
12+</html>
test/cases.tsadded+307−0View file
@@ -0,0 +1,307 @@
1+/**
2+ * The correctness suite, shared by both GPU stacks: `npm run test:node`
3+ * drives it through desktop Dawn, `npm run test:gpu` through headless
4+ * Chrome's own WebGPU (SwiftShader when there is no hardware).
5+ *
6+ * Each case compiles and runs a MATLAB script and checks the values it
7+ * prints (via fprintf) against references computed here in f64. Tolerances
8+ * are f32-scale: the GPU computes in single precision.
9+ */
10+import { runScript } from '../src/mgpu/session.ts';
11+import { formatFailure } from '../src/mgpu/errors.ts';
12+
13+export interface Case {
14+ name: string;
15+ source: string;
16+ /** Expected values for each %-printed number, with relative tolerance. */
17+ expect?: { value: number; rel?: number; abs?: number }[];
18+ /** Substrings that must appear in the output. */
19+ contains?: string[];
20+ /** Expected compile failure instead of a run. */
21+ fails?: RegExp;
22+}
23+
24+const N = (v: number, rel = 2e-5): { value: number; rel: number } => ({ value: v, rel });
25+
26+/** f64 reference for the fused-chain case. */
27+function refFused(): number {
28+ const n = 10000;
29+ let acc = 0;
30+ for (let i = 0; i < n; i++) {
31+ const x = i / (n - 1);
32+ acc += 3 * x * x - 2 * x + Math.sin(2 * Math.PI * x) / (1 + x * x);
33+ }
34+ return acc / n;
35+}
36+
37+export const cases: Case[] = [
38+ {
39+ name: 'fused elementwise chain',
40+ source: `
41+n = 10000;
42+x = linspace(0, 1, n);
43+y = 3*x.^2 - 2*x + sin(2*pi*x)./(1 + x.^2);
44+fprintf('%.6f\\n', sum(y(:)) / n);
45+`,
46+ // 3e-4 relative: SwiftShader's sin/exp are a touch less accurate than
47+ // hardware drivers'.
48+ expect: [N(refFused(), 3e-4)],
49+ },
50+ {
51+ name: 'gemm small vs f64 reference',
52+ source: `
53+n = 64;
54+A = zeros(n, n) + 1;
55+B = zeros(n, n) + 2;
56+C = A * B;
57+fprintf('%.1f %.1f\\n', C(:)'*C(:)/(n*n), sum(C(:)));
58+`,
59+ expect: [N(128 * 128), N(128 * 64 * 64)],
60+ },
61+ {
62+ name: 'reductions full and by columns',
63+ source: `
64+m = 300; n = 5;
65+A = zeros(m, n) + 3;
66+s = sum(A);
67+fprintf('%.1f %.1f %.1f %.1f\\n', s(:)'*s(:)/n, mean(A(:)), max(A(:)), prod(zeros(3,1)+2));
68+`,
69+ expect: [N(900 * 900), N(3), N(3), N(8)],
70+ },
71+ {
72+ name: 'indexing is a clear compile error',
73+ source: `
74+n = 1000;
75+x = zeros(n, 1);
76+for k = 1:50
77+ x = x + k;
78+end
79+fprintf('%.1f\\n', x(1 + 0*x(:)'*x(:)));
80+`,
81+ fails: /index|slic|supported/i,
82+ },
83+ {
84+ name: 'loop replay accumulates',
85+ source: `
86+n = 1000;
87+x = zeros(n, 1);
88+for k = 1:50
89+ x = x + k;
90+end
91+fprintf('%.1f\\n', mean(x));
92+`,
93+ expect: [N(1275)],
94+ },
95+ {
96+ name: 'rand statistics and loop-fresh draws',
97+ source: `
98+n = 200000;
99+a = rand(n, 1);
100+s = zeros(1, 1);
101+for k = 1:3
102+ s = s + mean(rand(n, 1));
103+end
104+fprintf('%.3f %.3f %.3f %.3f\\n', mean(a), mean(a.^2), s/3, mean(randn(n,1)));
105+`,
106+ expect: [
107+ { value: 0.5, abs: 0.01 },
108+ { value: 1 / 3, abs: 0.01 },
109+ { value: 0.5, abs: 0.01 },
110+ { value: 0, abs: 0.02 },
111+ ],
112+ },
113+ {
114+ name: 'comparisons, logicals, masks (monte carlo pi)',
115+ source: `
116+n = 400000;
117+x = 2*rand(n, 1) - 1;
118+y = 2*rand(n, 1) - 1;
119+inside = (x.^2 + y.^2) <= 1;
120+fprintf('%.3f\\n', 4*mean(inside));
121+`,
122+ expect: [{ value: Math.PI, abs: 0.03 }],
123+ },
124+ {
125+ name: 'transpose and matrix identities',
126+ source: `
127+m = 33; n = 17;
128+A = rand(m, n);
129+B = A';
130+d1 = sum(A(:).^2);
131+d2 = sum(B(:).^2);
132+E = eye(m);
133+C = E * A;
134+d3 = sum(abs(C(:) - A(:)));
135+fprintf('%.6f %.6f\\n', d1 - d2, d3);
136+`,
137+ expect: [{ value: 0, abs: 1e-3 }, { value: 0, abs: 1e-4 }],
138+ },
139+ {
140+ name: 'gemm matrix-vector agrees with itself',
141+ source: `
142+n = 48;
143+A = rand(n, n);
144+v = rand(n, 1);
145+w = A * v;
146+d = sum(w) - sum(A * v);
147+fprintf('%.6f\\n', d);
148+`,
149+ expect: [{ value: 0, abs: 1e-4 }],
150+ },
151+ {
152+ name: 'tic toc segments and echo',
153+ source: `
154+n = 100000;
155+tic;
156+x = rand(n, 1);
157+s = mean(x);
158+t = toc;
159+tic
160+y = x + 1;
161+toc
162+z = 3.5
163+`,
164+ contains: ['Elapsed time is', 'z =', '3.5'],
165+ },
166+ {
167+ name: 'dot and norm',
168+ source: `
169+n = 5000;
170+a = zeros(n,1) + 2;
171+b = zeros(n,1) + 3;
172+fprintf('%.1f %.4f\\n', dot(a, b), norm(a) / sqrt(n));
173+`,
174+ expect: [N(30000), N(2)],
175+ },
176+ {
177+ name: 'min/max elementwise two-arg',
178+ source: `
179+n = 1000;
180+x = linspace(-1, 1, n);
181+y = max(x, 0) + min(x, 0);
182+fprintf('%.6f\\n', sum(abs(y - x)));
183+`,
184+ expect: [{ value: 0, abs: 1e-4 }],
185+ },
186+ {
187+ name: 'mod and integer powers',
188+ source: `
189+x = linspace(0, 10, 101);
190+y = mod(x, 3);
191+fprintf('%.4f %.4f\\n', max(y(:)), sum((0:4).^2));
192+`,
193+ expect: [{ value: 2.9, abs: 0.001 }, N(30)],
194+ },
195+ {
196+ name: 'literal row vector uploads',
197+ source: `
198+v = [1 2 3 4 5];
199+fprintf('%.1f\\n', sum(v));
200+`,
201+ expect: [N(15)],
202+ },
203+ {
204+ name: 'in-place update (aliased kernel) is correct',
205+ source: `
206+n = 100;
207+u = zeros(n, 1) + 1;
208+u = u + u.^2;
209+u = u * 2;
210+fprintf('%.1f\\n', mean(u));
211+`,
212+ expect: [N(4)],
213+ },
214+ {
215+ name: 'while is a clear compile error',
216+ source: `
217+x = 1;
218+while x < 10
219+ x = x + 1;
220+end
221+`,
222+ fails: /while/i,
223+ },
224+ {
225+ name: 'variable-size rand is a clear compile error',
226+ source: `
227+n = rand() * 100;
228+A = rand(n, 1);
229+`,
230+ fails: /compile time/i,
231+ },
232+];
233+
234+/** The %-formatted numbers a run printed (echo/timing lines stripped). */
235+function printedNumbers(output: string): number[] {
236+ const out: number[] = [];
237+ const kept = output
238+ .split('\n')
239+ .filter((l) => !/Elapsed time|=/.test(l))
240+ .join('\n');
241+ for (const m of kept.matchAll(/-?\d+(?:\.\d+)?(?:e[+-]?\d+)?/gi)) {
242+ out.push(Number(m[0]));
243+ }
244+ return out;
245+}
246+
247+const indent = (s: string): string => s.replace(/^/gm, ' ');
248+
249+/** Run every case; log a line per case; return the failure count. */
250+export async function runCases(
251+ device: GPUDevice,
252+ log: (line: string) => void,
253+): Promise<number> {
254+ let failures = 0;
255+ for (const c of cases) {
256+ try {
257+ const run = await runScript(device, c.source);
258+ if (c.fails) {
259+ log(`FAIL ${c.name}: expected a compile error, but it ran`);
260+ failures++;
261+ continue;
262+ }
263+ if (run.result.error) {
264+ log(`FAIL ${c.name}: runtime error: ${run.result.error}`);
265+ failures++;
266+ continue;
267+ }
268+ const nums = printedNumbers(run.result.output);
269+ let ok = true;
270+ (c.expect ?? []).forEach((e, i) => {
271+ const got = nums[i];
272+ const tol = e.abs ?? Math.abs(e.value) * (e.rel ?? 1e-5) + 1e-12;
273+ if (got === undefined || Math.abs(got - e.value) > tol) {
274+ log(`FAIL ${c.name}: printed[${i}] = ${got}, want ${e.value} ±${tol}`);
275+ ok = false;
276+ }
277+ });
278+ for (const s of c.contains ?? []) {
279+ if (!run.result.output.includes(s)) {
280+ log(`FAIL ${c.name}: output lacks ${JSON.stringify(s)}`);
281+ ok = false;
282+ }
283+ }
284+ if (!ok) {
285+ log(` output was:\n${indent(run.result.output)}`);
286+ log(` plan:\n${indent(run.planDescription.join('\n'))}`);
287+ failures++;
288+ } else {
289+ log(`ok ${c.name}`);
290+ }
291+ } catch (e) {
292+ if (c.fails) {
293+ const msg = formatFailure(e, c.source);
294+ if (c.fails.test(msg)) {
295+ log(`ok ${c.name} (declined: ${msg.split('\n')[0].slice(0, 90)})`);
296+ } else {
297+ log(`FAIL ${c.name}: wrong error: ${msg}`);
298+ failures++;
299+ }
300+ } else {
301+ log(`FAIL ${c.name}: ${formatFailure(e, c.source)}`);
302+ failures++;
303+ }
304+ }
305+ }
306+ return failures;
307+}
test/test-page.tsadded+32−0View file
@@ -0,0 +1,32 @@
1+/**
2+ * Browser validation, in the environment the sandbox actually ships to:
3+ * the same suite as `npm run test:node`, on the browser's own WebGPU.
4+ * Results are posted to window.__RESULTS__ for the headless runner.
5+ */
6+import { requestSandboxDevice } from '../src/mgpu/session.ts';
7+import { runCases } from './cases.ts';
8+
9+declare global {
10+ interface Window {
11+ __RESULTS__?: { ok: boolean; fatal?: string; lines: string[] };
12+ }
13+}
14+
15+const logEl = document.getElementById('log')!;
16+const lines: string[] = [];
17+const log = (line: string): void => {
18+ lines.push(line);
19+ logEl.textContent = lines.join('\n');
20+};
21+
22+(async () => {
23+ const { device, description } = await requestSandboxDevice();
24+ log(`adapter: ${description}`);
25+ const failures = await runCases(device, log);
26+ log(failures ? `${failures} failure(s)` : 'all tests passed');
27+ window.__RESULTS__ = { ok: failures === 0, lines };
28+})().catch((e) => {
29+ const fatal = e instanceof Error ? e.message : String(e);
30+ log(`fatal: ${fatal}`);
31+ window.__RESULTS__ = { ok: false, fatal, lines };
32+});
tsconfig.jsonadded+15−0View file
@@ -0,0 +1,15 @@
1+{
2+ "compilerOptions": {
3+ "target": "ES2022",
4+ "module": "ESNext",
5+ "moduleResolution": "bundler",
6+ "lib": ["ES2022", "DOM", "DOM.Iterable"],
7+ "types": ["@webgpu/types", "node"],
8+ "strict": true,
9+ "noEmit": true,
10+ "allowImportingTsExtensions": true,
11+ "verbatimModuleSyntax": true,
12+ "skipLibCheck": true
13+ },
14+ "include": ["src", "test", "scripts"]
15+}
vite.config.tsadded+36−0View file
@@ -0,0 +1,36 @@
1+import { defineConfig } from 'vite';
2+import { resolve } from 'node:path';
3+
4+// numbl is a local `file:` dependency, so node_modules/numbl is a symlink to
5+// the sibling checkout. Its package `exports` map only publishes the runtime
6+// entry points, not the compiler internals we need (parser + JIT lowering), so
7+// we reach them through a path alias — the same arrangement turing-surface
8+// uses, and documented there in src/mgpu/numbl.d.ts.
9+const numblSrc = resolve(import.meta.dirname, 'node_modules/numbl/src');
10+
11+export default defineConfig({
12+ base: './',
13+ resolve: {
14+ alias: { 'numbl-src': numblSrc },
15+ // Without this, the dev server canonicalizes the symlink for SOME import
16+ // chains (`/@fs/<realpath>?v=...`) but not others (raw absolute paths),
17+ // which (a) trips the fs allow list and (b) loads numbl's builtin
18+ // registry twice — so the type-rule patches in src/mgpu/patches.ts land
19+ // on one instance while the lowerer consults the other. Keeping the
20+ // symlinked path everywhere gives every module one canonical URL.
21+ preserveSymlinks: true,
22+ },
23+ server: {
24+ // the alias resolves through the symlink; allow both spellings
25+ fs: { allow: [import.meta.dirname] },
26+ },
27+ build: {
28+ target: 'es2022',
29+ rollupOptions: {
30+ input: {
31+ main: resolve(import.meta.dirname, 'index.html'),
32+ test: resolve(import.meta.dirname, 'test.html'),
33+ },
34+ },
35+ },
36+});