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 */
26import 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';
35import {
36 isMultiElement,
37 type NumericType,
38 type Type,
39} from 'numbl-src/numbl-core/jit/lowering/types.ts';
40import type { CompiledScript } from './compile.ts';
41import { UnsupportedOnGpu } from './errors.ts';
42import {
43 buildKernel,
44 bindingDecls,
45 dispatchFor,
46 emitLoader,
47 exactValue,
48 fullColonBase,
49 numel,
50 REDUCTIONS,
51 type KernelInputs,
52} from './wgsl.ts';
53import {
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';
67const isNumeric = (t: Type): t is NumericType => t.kind === 'Numeric';
68const isTensor = (t: Type): boolean => isNumeric(t) && isMultiElement(t);
69const isVectorish = (t: Type): boolean =>
70 isNumeric(t) && (t.shape ?? []).filter((d) => d !== 1).length <= 1;
72/** Iterations a single `for` may replay. Each costs one 256-byte uniform slot
73 * and a re-encode of the body's dispatches. */
74const MAX_TRIPS = 65536;
75/** Total dispatches one run may encode, across all loops. */
76const MAX_DISPATCHES = 2_000_000;
77/** Loop-variable uniform slot stride (minUniformBufferOffsetAlignment). */
78const LV_STRIDE = 256;
80export interface Slot {
81 buffer: GPUBuffer;
82 count: number;
83}
85/** What a printed/displayed value reads from. */
86export type ValueRef =
87 | { kind: 'literal'; value: number }
88 | { kind: 'buffer'; slot: Slot; count: number }
89 | { kind: 'host'; cName: string };
91/** One piece of an fprintf: fixed text or a formatted value. */
92export type EmitPart =
93 | { kind: 'text'; text: string }
94 | { kind: 'value'; ref: ValueRef; spec: string };
96export 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 };
128export interface ScriptPlan {
129 ops: Op[];
130 /** Human-readable op sequence — what the script actually compiled to. */
131 describe(): string[];
132 destroy(): void;
133}
135interface PlannerVar {
136 slot: Slot;
137 shape: number[];
138}
140export 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[] = [];
158 let seedCounter = 1;
159 let tempCounter = 0;
160 let tocCounter = 0;
161 let dispatchBudget = MAX_DISPATCHES;
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 }
180 const resolve = (cName: string): string => {
181 let c = cName;
182 while (aliases.has(c)) c = aliases.get(c)!;
183 return c;
184 };
186 const maxBytes = device.limits.maxStorageBufferBindingSize;
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 };
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 };
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 };
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 }
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 }
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;
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[] = [];
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 });
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 }
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 }
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 }
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 }
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 }
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 }
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 }
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);
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;
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 }
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 }
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;
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 }
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;
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 }
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 }
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 }
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 }
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' };
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);
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;
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 }
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 }
823 // ── Host-op planning ──────────────────────────────────────────────────
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 };
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 }
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 }
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 }
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 }
965 // ── Statement walk ────────────────────────────────────────────────────
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 }
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 }
994 const isTemp = (name: string): boolean => name.startsWith('_mtoc2_') || name.startsWith('%tmp');
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 }
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 }
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 }
1058 await planValue(stmt.cName, stmt.name, stmt.ty, stmt.expr, stmt.span);
1059 maybeEcho(stmt);
1060 }
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 }
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 }
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 }
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 }
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);
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)`);
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 }
1209 for (const stmt of compiled.stmts) {
1210 await planStmt(stmt);
1211 }
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}