/ concept-collection / math-webgpu-sandbox
Sign in
concept-collection / math-webgpu-sandbox
math-webgpu-sandbox / src / mgpu / plan.ts
1234 lines · 42.8 KBCodeBlameHistory
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;
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 };
4ded6ffShow the generated WGSL kernels in the compiled-plan paneJeremy Magland 128/** One generated compute shader, for the "show me the kernels" pane. */
129export interface KernelSource {
130 /** The label of the statement that first created it. */
131 label: string;
132 code: string;
136 ops: Op[];
137 /** Human-readable op sequence — what the script actually compiled to. */
138 describe(): string[];
4ded6ffShow the generated WGSL kernels in the compiled-plan paneJeremy Magland 139 /** Every distinct WGSL kernel, in creation order (reused kernels appear
140 * once, under their first label). */
141 kernels: KernelSource[];
145interface PlannerVar {
146 slot: Slot;
147 shape: number[];
150export async function planScript(
151 device: GPUDevice,
152 compiled: CompiledScript,
153): Promise<ScriptPlan> {
154 const owned: GPUBuffer[] = [];
155 /** cName -> buffer-backed variable (tensors and runtime scalars). */
156 const vars = new Map<string, PlannerVar>();
157 /** cName -> cName it is a view of. */
158 const aliases = new Map<string, string>();
159 /** cName -> char value (format strings). */
160 const chars = new Map<string, string>();
161 /** cNames whose value the executor knows on the host (tic/toc results). */
162 const hostScalars = new Set<string>();
163 /** Loop-variable uniform buffers, by cVar. */
164 const loopUniforms = new Map<string, GPUBuffer>();
165 const pipelines = new Map<string, GPUComputePipeline>();
4ded6ffShow the generated WGSL kernels in the compiled-plan paneJeremy Magland 166 const kernelSources: KernelSource[] = [];
e01ddf1MATLAB-syntax scripts on WebGPU: fused kernels, tic/toc timing, CPU comparisonJeremy Magland 167 const describeLines: string[] = [];
169 let seedCounter = 1;
170 let tempCounter = 0;
171 let tocCounter = 0;
172 let dispatchBudget = MAX_DISPATCHES;
174 // How many times each cName is assigned, anywhere. Decides when a variable
175 // may be a view of another's buffer, and when an exact scalar still needs a
176 // real buffer (a later assignment makes its uses flow-dependent).
177 const assignCounts = new Map<string, number>();
178 {
179 const walkCounts = (stmts: IRStmt[]): void => {
180 for (const s of stmts) {
181 if (s.kind === 'Assign') {
182 assignCounts.set(s.cName, (assignCounts.get(s.cName) ?? 0) + 1);
183 } else if (s.kind === 'For') {
184 walkCounts(s.body);
185 }
186 }
187 };
188 walkCounts(compiled.stmts);
189 }
191 const resolve = (cName: string): string => {
192 let c = cName;
193 while (aliases.has(c)) c = aliases.get(c)!;
194 return c;
195 };
197 const maxBytes = device.limits.maxStorageBufferBindingSize;
199 const makeSlot = (label: string, count: number): Slot => {
200 const bytes = Math.max(4, 4 * count);
201 if (bytes > maxBytes) {
202 throw new UnsupportedOnGpu(
203 `'${label}' needs ${(bytes / 1e6).toFixed(0)} MB, over this device's ` +
204 `storage-buffer limit of ${(maxBytes / 1e6).toFixed(0)} MB`,
205 );
206 }
207 const buffer = device.createBuffer({
208 label: `mgpu-${label}`,
209 size: bytes,
210 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
211 });
212 owned.push(buffer);
213 return { buffer, count };
214 };
216 /** The variable's slot, allocating on first assignment. */
217 const slotFor = (cName: string, name: string, ty: NumericType, span: Span): Slot => {
218 const key = resolve(cName);
219 const count = numel(ty);
220 const existing = vars.get(key);
221 if (existing) {
222 if (existing.slot.count !== count) {
223 throw new UnsupportedOnGpu(
224 `'${name}' changes size between assignments (${existing.slot.count} ` +
225 `-> ${count} elements); the sandbox fixes each variable's storage once`,
226 span,
227 );
228 }
229 existing.shape = ty.shape ?? [count, 1];
230 return existing.slot;
231 }
232 const slot = makeSlot(name, count);
233 vars.set(key, { slot, shape: ty.shape ?? [count, 1] });
234 return slot;
235 };
237 const readSlot = (cName: string, name: string, span: Span): Slot => {
238 const v = vars.get(resolve(cName));
239 if (!v) {
240 throw new UnsupportedOnGpu(`'${name}' is read before it has a value`, span);
241 }
242 return v.slot;
243 };
245 async function pipeline(code: string, label: string, layout: GPUBindGroupLayout): Promise<GPUComputePipeline> {
246 const hit = pipelines.get(code);
247 if (hit) return hit;
248 device.pushErrorScope('validation');
249 const module = device.createShaderModule({ code, label });
250 const info = await module.getCompilationInfo();
251 const errors = info.messages.filter((m) => m.type === 'error');
252 if (errors.length) {
253 throw new UnsupportedOnGpu(
254 `generated WGSL failed to compile for '${label}':\n` +
255 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n') +
256 `\n--- shader ---\n${code}`,
257 );
258 }
259 const p = await device.createComputePipelineAsync({
260 layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }),
261 compute: { module, entryPoint: 'main' },
262 label,
263 });
264 const err = await device.popErrorScope();
265 if (err) throw new UnsupportedOnGpu(`pipeline '${label}': ${err.message}`);
266 pipelines.set(code, p);
4ded6ffShow the generated WGSL kernels in the compiled-plan paneJeremy Magland 267 kernelSources.push({ label, code });
269 }
271 /** Bind group layout: out at 0, `inputs` read-only buffers, then `loops`
272 * dynamic-offset uniforms. Explicit so unused bindings still match. */
273 function kernelLayout(inputs: number, loops: number): GPUBindGroupLayout {
274 const entries: GPUBindGroupLayoutEntry[] = [
275 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
276 ];
277 for (let i = 0; i < inputs; i++) {
278 entries.push({
279 binding: i + 1,
280 visibility: GPUShaderStage.COMPUTE,
281 buffer: { type: 'read-only-storage' },
282 });
283 }
284 for (let k = 0; k < loops; k++) {
285 entries.push({
286 binding: inputs + 1 + k,
287 visibility: GPUShaderStage.COMPUTE,
288 buffer: { type: 'uniform', hasDynamicOffset: true },
289 });
290 }
291 return device.createBindGroupLayout({ entries });
292 }
294 /** Ops are appended to the innermost list; loops nest via this stack. */
295 const opStack: Op[][] = [[]];
296 const ops = (): Op[] => opStack[opStack.length - 1];
297 const enclosingLoops: string[] = [];
298 const inLoop = (): boolean => enclosingLoops.length > 0;
300 const spendDispatches = (n: number, span: Span): void => {
301 // Inside loops the body is re-encoded per iteration; multiply out.
302 let mult = 1;
303 for (const f of loopTrips) mult *= f;
304 dispatchBudget -= n * mult;
305 if (dispatchBudget < 0) {
306 throw new UnsupportedOnGpu(
307 `this script would encode more than ${MAX_DISPATCHES.toLocaleString()} ` +
308 `GPU dispatches; shrink loop counts`,
309 span,
310 );
311 }
312 };
313 const loopTrips: number[] = [];
315 /** KernelInputs seeded with fresh maps; buffers registered on demand. */
316 const freshIo = (): KernelInputs => ({
317 buffers: new Map(),
318 loopVars: new Map(),
319 nextSeed: () => seedCounter++,
320 });
322 /** Register every buffer-backed read in `expr` into `io.buffers`. */
323 function collectBuffers(expr: IRExpr, io: KernelInputs, span: Span): void {
324 const walk = (x: IRExpr): void => {
325 // A constant-folded subtree emits as a literal; nothing under it is read.
326 if (exactValue(x) !== undefined) return;
327 switch (x.kind) {
328 case 'Var': {
329 if (enclosingLoops.includes(x.cName)) {
330 if (!io.loopVars.has(x.cName)) io.loopVars.set(x.cName, io.loopVars.size);
331 return;
332 }
333 if (!isNumeric(x.ty)) {
334 throw new UnsupportedOnGpu(`'${x.name}' is not numeric`, x.span);
335 }
336 if (typeof x.ty.exact === 'number') return; // folds to a literal
337 const key = resolve(x.cName);
338 if (!io.buffers.has(key)) io.buffers.set(key, io.buffers.size);
339 // Rewrite so the emitter sees the resolved name.
340 x.cName = key;
341 return;
342 }
343 case 'Binary':
344 walk(x.left);
345 walk(x.right);
346 return;
347 case 'Unary':
348 walk(x.operand);
349 return;
350 case 'IndexSlice':
351 walk(x.base);
352 return;
353 case 'MakeRange':
354 walk(x.start);
355 walk(x.step);
356 return;
357 case 'Call':
358 if (!['zeros', 'ones', 'eye', 'rand', 'randn'].includes(x.name)) {
359 x.args.forEach(walk);
360 }
361 return;
362 default:
363 return;
364 }
365 };
366 walk(expr);
367 void span;
368 }
370 /** Bind-group entries for a kernel built from `io`, writing `target`. */
371 function kernelBindGroup(
372 io: KernelInputs,
373 bufferOrder: string[],
374 loopOrder: string[],
375 target: GPUBuffer,
376 layout: GPUBindGroupLayout,
377 span: Span,
378 ): GPUBindGroup {
379 const entries: GPUBindGroupEntry[] = [{ binding: 0, resource: { buffer: target } }];
380 bufferOrder.forEach((cName, i) => {
381 const v = vars.get(resolve(cName));
382 if (!v) throw new UnsupportedOnGpu(`a value read here has no buffer`, span);
383 entries.push({ binding: i + 1, resource: { buffer: v.slot.buffer } });
384 });
385 loopOrder.forEach((cVar, k) => {
386 const u = loopUniforms.get(cVar);
387 if (!u) throw new UnsupportedOnGpu(`internal: no uniform for loop '${cVar}'`, span);
388 entries.push({
389 binding: bufferOrder.length + 1 + k,
390 resource: { buffer: u, offset: 0, size: 8 },
391 });
392 });
393 return device.createBindGroup({ layout, entries });
394 }
396 /** Materialize an arbitrary tensor expression into a slot, planning
397 * whatever ops that takes. A plain Var is returned as-is. */
398 async function materialize(expr: IRExpr): Promise<{ cName: string; slot: Slot }> {
399 if (expr.kind === 'Var' && isTensor(expr.ty) && typeof (expr.ty as NumericType).exact !== 'object') {
400 return { cName: resolve(expr.cName), slot: readSlot(expr.cName, expr.name, expr.span) };
401 }
402 if (expr.kind === 'IndexSlice') {
403 const base = fullColonBase(expr);
404 if (base && base.kind === 'Var') {
405 return { cName: resolve(base.cName), slot: readSlot(base.cName, base.name, base.span) };
406 }
407 }
408 if (!isNumeric(expr.ty)) {
409 throw new UnsupportedOnGpu(`expression is not numeric`, expr.span);
410 }
411 const cName = `%tmp${tempCounter++}`;
412 await planValue(cName, cName, expr.ty, expr, expr.span);
413 return { cName, slot: vars.get(resolve(cName))!.slot };
414 }
416 /** Replace non-elementwise subtrees (GEMM, matrix transpose, reductions)
417 * with materialized temps, so what remains is one fused kernel. */
418 async function hoistNonElementwise(expr: IRExpr): Promise<IRExpr> {
419 const hoist = async (x: IRExpr): Promise<IRExpr> => {
420 if (isHoistRoot(x)) {
421 const { cName, slot } = await materialize(x);
422 void slot;
423 return {
424 kind: 'Var',
425 name: cName,
426 cName,
427 ty: x.ty,
428 span: x.span,
429 };
430 }
431 switch (x.kind) {
432 case 'Binary':
433 x.left = await hoist(x.left);
434 x.right = await hoist(x.right);
435 return x;
436 case 'Unary':
437 x.operand = await hoist(x.operand);
438 return x;
439 case 'Call':
440 for (let i = 0; i < x.args.length; i++) x.args[i] = await hoist(x.args[i]);
441 return x;
442 case 'IndexSlice':
443 x.base = await hoist(x.base);
444 return x;
445 default:
446 return x;
447 }
448 };
449 // The root itself was already routed by planValue; only hoist children.
450 switch (expr.kind) {
451 case 'Binary':
452 expr.left = await hoist(expr.left);
453 expr.right = await hoist(expr.right);
454 return expr;
455 case 'Unary':
456 expr.operand = await hoist(expr.operand);
457 return expr;
458 case 'Call':
459 for (let i = 0; i < expr.args.length; i++) expr.args[i] = await hoist(expr.args[i]);
460 return expr;
461 case 'IndexSlice':
462 expr.base = await hoist(expr.base);
463 return expr;
464 default:
465 return expr;
466 }
467 }
469 function isHoistRoot(x: IRExpr): boolean {
470 if (x.kind === 'Binary' && x.builtin === 'mtimes' && isTensor(x.left.ty) && isTensor(x.right.ty)) {
471 return true;
472 }
473 if (x.kind === 'Unary' && x.builtin === 'transpose' && isTensor(x.operand.ty) && !isVectorish(x.operand.ty)) {
474 return true;
475 }
476 if (x.kind === 'Call' && REDUCTIONS.has(x.name)) {
477 // Two-arg max/min is elementwise, not a reduction.
478 if ((x.name === 'max' || x.name === 'min') && x.args.length === 2) return false;
479 return x.args.some((a) => isTensor(a.ty));
480 }
481 return false;
482 }
484 /** Plan `dest = expr`, routing to the right kind of op. */
485 async function planValue(
486 cName: string,
487 name: string,
488 ty: NumericType,
489 expr: IRExpr,
490 span: Span,
491 ): Promise<void> {
492 // View-or-copy forms: X(:), reshape, vector transpose, plain `B = A`.
493 const viewOf = viewSource(expr);
494 if (viewOf) {
495 const srcSlot = readSlot(viewOf.cName, viewOf.name, span);
496 const srcKey = resolve(viewOf.cName);
497 if (numel(ty) !== srcSlot.count) {
498 throw new UnsupportedOnGpu(`'${name}' and '${viewOf.name}' differ in size`, span);
499 }
500 const destAssigns = assignCounts.get(cName) ?? 1;
501 const srcAssigns = assignCounts.get(viewOf.cName) ?? 1;
502 if (destAssigns === 1 && srcAssigns === 1 && !inLoop() && !vars.has(cName)) {
503 aliases.set(cName, srcKey);
504 // Track the new orientation under the alias's own name via vars of
505 // the base: display uses the Assign's ty directly, so nothing to do.
506 describeLines.push(`view ${name} -> ${viewOf.name}`);
507 return;
508 }
509 const dest = slotFor(cName, name, ty, span);
510 if (dest.buffer !== srcSlot.buffer) {
511 ops().push({
512 kind: 'copy',
513 from: srcSlot.buffer,
514 to: dest.buffer,
515 bytes: 4 * srcSlot.count,
516 label: `${name} = ${viewOf.name}`,
517 });
518 describeLines.push(`copy ${name} = ${viewOf.name}`);
519 }
520 return;
521 }
523 // GEMM: A * B with both sides tensors.
524 if (expr.kind === 'Binary' && expr.builtin === 'mtimes' && isTensor(expr.left.ty) && isTensor(expr.right.ty)) {
525 return planGemm(cName, name, ty, expr, span);
526 }
527 // Matrix transpose.
528 if (expr.kind === 'Unary' && expr.builtin === 'transpose' && isTensor(expr.operand.ty) && !isVectorish(expr.operand.ty)) {
529 return planTranspose(cName, name, ty, expr, span);
530 }
531 // Reductions.
532 if (expr.kind === 'Call' && REDUCTIONS.has(expr.name) &&
533 !((expr.name === 'max' || expr.name === 'min') && expr.args.length === 2) &&
534 expr.args.some((a) => isTensor(a.ty))) {
535 return planReduce(cName, name, ty, expr as Call, span);
536 }
538 // Fused elementwise kernel (with any non-elementwise subtrees hoisted).
539 const hoisted = await hoistNonElementwise(expr);
540 const io = freshIo();
541 collectBuffers(hoisted, io, span);
542 const label = `${name} = <${numel(ty)} elem>`;
543 const kernel = buildKernel({ name, cName, ty, expr: hoisted, span }, io, [...enclosingLoops], label);
544 const layout = kernelLayout(kernel.buffers.length, kernel.loops.length);
545 const pipe = await pipeline(kernel.code, label, layout);
547 const dest = slotFor(cName, name, ty, span);
548 // WebGPU forbids aliasing a writable binding with a readable one, so an
549 // in-place update (`u = u + 1`) writes scratch and copies back.
550 const aliased = kernel.buffers.some((c) => vars.get(resolve(c))?.slot.buffer === dest.buffer);
551 const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;
553 const bindGroup = kernelBindGroup(io, kernel.buffers, kernel.loops, target.buffer, layout, span);
554 spendDispatches(1, span);
555 ops().push({
556 kind: 'kernel',
557 pipeline: pipe,
558 bindGroup,
559 dispatch: dispatchFor(kernel.count),
560 loops: kernel.loops,
561 label,
562 copyBack: aliased
563 ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
564 : undefined,
565 });
566 describeLines.push(`kernel ${label}${aliased ? ' (+copy-back)' : ''}`);
567 }
569 /** `B = A`, `X(:)`, `reshape(A, ...)`, vector `x'` — pure views/copies. */
570 function viewSource(expr: IRExpr): { cName: string; name: string } | null {
571 if (expr.kind === 'Var' && isTensor(expr.ty)) {
572 return { cName: expr.cName, name: expr.name };
573 }
574 if (expr.kind === 'IndexSlice') {
575 const base = fullColonBase(expr);
576 if (base && base.kind === 'Var' && isTensor(base.ty)) {
577 return { cName: base.cName, name: base.name };
578 }
579 return null;
580 }
581 if (expr.kind === 'Call' && expr.name === 'reshape' && expr.args.length >= 1 &&
582 expr.args[0].kind === 'Var' && isTensor(expr.args[0].ty)) {
583 return { cName: expr.args[0].cName, name: expr.args[0].name };
584 }
585 if (expr.kind === 'Unary' && expr.builtin === 'transpose' &&
586 isVectorish(expr.operand.ty) && expr.operand.kind === 'Var' && isTensor(expr.operand.ty)) {
587 return { cName: expr.operand.cName, name: expr.operand.name };
588 }
589 return null;
590 }
592 async function planGemm(
593 cName: string,
594 name: string,
595 ty: NumericType,
596 expr: IRExpr & { kind: 'Binary' },
597 span: Span,
598 ): Promise<void> {
599 const a = await materialize(expr.left);
600 const b = await materialize(expr.right);
601 const [m, k] = shapeOf(expr.left.ty, span);
602 const [, n] = shapeOf(expr.right.ty, span);
603 const dest = slotFor(cName, name, ty, span);
604 const aliased = dest.buffer === a.slot.buffer || dest.buffer === b.slot.buffer;
605 const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;
607 const label = `${name} = ${m}x${k} * ${k}x${n}`;
608 const layout = kernelLayout(2, 0);
609 const pipe = await pipeline(gemmKernel(m, k, n), label, layout);
610 const bindGroup = device.createBindGroup({
611 layout,
612 entries: [
613 { binding: 0, resource: { buffer: target.buffer } },
614 { binding: 1, resource: { buffer: a.slot.buffer } },
615 { binding: 2, resource: { buffer: b.slot.buffer } },
616 ],
617 });
618 spendDispatches(1, span);
619 ops().push({
620 kind: 'kernel',
621 pipeline: pipe,
622 bindGroup,
623 dispatch: gemmDispatch(m, n),
624 loops: [],
625 label,
626 copyBack: aliased
627 ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
628 : undefined,
629 });
630 describeLines.push(`gemm ${label}`);
631 }
633 async function planTranspose(
634 cName: string,
635 name: string,
636 ty: NumericType,
637 expr: IRExpr & { kind: 'Unary' },
638 span: Span,
639 ): Promise<void> {
640 const src = await materialize(expr.operand);
641 const [m, n] = shapeOf(expr.operand.ty, span);
642 const dest = slotFor(cName, name, ty, span);
643 const aliased = dest.buffer === src.slot.buffer;
644 const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;
646 const label = `${name} = (${m}x${n})'`;
647 const layout = kernelLayout(1, 0);
648 const pipe = await pipeline(transposeKernel(m, n), label, layout);
649 const bindGroup = device.createBindGroup({
650 layout,
651 entries: [
652 { binding: 0, resource: { buffer: target.buffer } },
653 { binding: 1, resource: { buffer: src.slot.buffer } },
654 ],
655 });
656 spendDispatches(1, span);
657 ops().push({
658 kind: 'kernel',
659 pipeline: pipe,
660 bindGroup,
661 dispatch: transposeDispatch(m, n),
662 loops: [],
663 label,
664 copyBack: aliased
665 ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
666 : undefined,
667 });
668 describeLines.push(`transp ${label}`);
669 }
671 async function planReduce(
672 cName: string,
673 name: string,
674 ty: NumericType,
675 call: Call,
676 span: Span,
677 ): Promise<void> {
678 const fn = call.name;
679 if (fn !== 'dot' && call.args.length !== 1) {
680 throw new UnsupportedOnGpu(
681 `'${fn}' supports only the one-argument form here (no dim/'all' ` +
682 `arguments — use ${fn}(X(:)) for the whole array)`,
683 span,
684 );
685 }
687 // Build the loader expression: the (fused) element the reduction eats.
688 let loaderExpr: IRExpr;
689 let map: MapKind = 'id';
690 if (fn === 'dot') {
691 if (call.args.length !== 2) {
692 throw new UnsupportedOnGpu(`'dot' takes two vectors`, span);
693 }
694 loaderExpr = {
695 kind: 'Binary',
696 builtin: 'times',
697 left: call.args[0],
698 right: call.args[1],
699 ty: call.args[0].ty,
700 span,
701 };
702 } else {
703 loaderExpr = call.args[0];
704 if (fn === 'norm') {
705 if (!isVectorish(loaderExpr.ty)) {
706 throw new UnsupportedOnGpu(
707 `'norm' of a matrix is the spectral norm, which the sandbox does ` +
708 `not compute; norm(v) for vectors only`,
709 span,
710 );
711 }
712 map = 'sq';
713 }
714 }
715 // See through X(:) so the loader reads the base buffer directly.
716 if (loaderExpr.kind === 'IndexSlice') {
717 const base = fullColonBase(loaderExpr);
718 if (base) loaderExpr = { ...base, ty: loaderExpr.ty } as IRExpr;
719 }
721 const inputTy = loaderExpr.ty as NumericType;
722 const inCount = numel(inputTy);
723 const outCount = numel(ty);
724 const full = outCount === 1;
725 if (!full) {
726 const [m, n] = shapeOf(inputTy, span);
727 const [om, on] = shapeOf(ty, span);
728 if (om !== 1 || on !== n) {
729 throw new UnsupportedOnGpu(
730 `'${fn}' along that dimension is not supported — transpose first, ` +
731 `or reduce the whole array with ${fn}(X(:))`,
732 span,
733 );
734 }
735 void m;
736 }
738 const combine: Combine =
739 fn === 'prod' ? 'mul' : fn === 'max' ? 'max' : fn === 'min' ? 'min' : 'add';
740 const epilogue: Epilogue =
741 fn === 'mean'
742 ? { kind: 'scale', by: 1 / (full ? inCount : shapeOf(inputTy, span)[0]) }
743 : fn === 'norm'
744 ? { kind: 'sqrt' }
745 : { kind: 'none' };
747 // Hoist nested non-elementwise pieces, then emit the fused loader.
748 loaderExpr = await hoistNonElementwise(
749 loaderExpr.kind === 'Binary' || loaderExpr.kind === 'Unary' || loaderExpr.kind === 'Call' || loaderExpr.kind === 'IndexSlice'
750 ? loaderExpr
751 : loaderExpr,
752 );
753 if (isHoistRoot(loaderExpr)) {
754 const { cName: mc } = await materialize(loaderExpr);
755 loaderExpr = { kind: 'Var', name: mc, cName: mc, ty: loaderExpr.ty, span };
756 }
757 const io = freshIo();
758 collectBuffers(loaderExpr, io, span);
759 const loader = emitLoader(loaderExpr, io, [...enclosingLoops]);
760 const { decls, buffers, loops } = bindingDecls(io);
762 const dest = slotFor(cName, name, ty, span);
763 const aliased = buffers.some((c) => vars.get(resolve(c))?.slot.buffer === dest.buffer);
764 const target = aliased ? makeSlot(`${name}-scratch`, dest.count) : dest;
766 if (full) {
767 const numWg = reducePartials(inCount);
768 const partials = makeSlot(`${name}-partials`, numWg);
769 const label1 = `${name} = ${fn}(<${inCount} elem>) pass1`;
770 const layout1 = kernelLayout(buffers.length, loops.length);
771 const pipe1 = await pipeline(
772 reduceFullPass1(decls, loader, inCount, numWg, combine, map),
773 label1,
774 layout1,
775 );
776 const bg1 = kernelBindGroup(io, buffers, loops, partials.buffer, layout1, span);
777 const label2 = `${name} = ${fn}(...) pass2`;
778 const layout2 = kernelLayout(1, 0);
779 const pipe2 = await pipeline(reduceFullPass2(numWg, combine, epilogue), label2, layout2);
780 const bg2 = device.createBindGroup({
781 layout: layout2,
782 entries: [
783 { binding: 0, resource: { buffer: target.buffer } },
784 { binding: 1, resource: { buffer: partials.buffer } },
785 ],
786 });
787 spendDispatches(2, span);
788 ops().push({
789 kind: 'kernel', pipeline: pipe1, bindGroup: bg1,
790 dispatch: [numWg, 1], loops, label: label1,
791 });
792 ops().push({
793 kind: 'kernel', pipeline: pipe2, bindGroup: bg2,
794 dispatch: [1, 1], loops: [], label: label2,
795 copyBack: aliased
796 ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
797 : undefined,
798 });
799 describeLines.push(`reduce ${name} = ${fn}(<${inCount} elem>)`);
800 } else {
801 const [m, n] = shapeOf(inputTy, span);
802 if (n > 65535) {
803 throw new UnsupportedOnGpu(`'${fn}' over ${n} columns exceeds the dispatch limit`, span);
804 }
805 const label = `${name} = ${fn}(${m}x${n} by columns)`;
806 const layout = kernelLayout(buffers.length, loops.length);
807 const pipe = await pipeline(
808 reduceColumns(decls, loader, m, combine, map, epilogue),
809 label,
810 layout,
811 );
812 const bg = kernelBindGroup(io, buffers, loops, target.buffer, layout, span);
813 spendDispatches(1, span);
814 ops().push({
815 kind: 'kernel', pipeline: pipe, bindGroup: bg,
816 dispatch: [n, 1], loops, label,
817 copyBack: aliased
818 ? { from: target.buffer, to: dest.buffer, bytes: 4 * dest.count }
819 : undefined,
820 });
821 describeLines.push(`reduce ${label}`);
822 }
823 }
825 function shapeOf(t: Type, span: Span): [number, number] {
826 if (!isNumeric(t) || !t.shape) {
827 throw new UnsupportedOnGpu(`shape is not known at compile time`, span);
828 }
829 if (t.shape.length !== 2) {
830 throw new UnsupportedOnGpu(`only 2-D arrays are supported (got ${t.shape.length}-D)`, span);
831 }
832 return [t.shape[0], t.shape[1]];
833 }
835 // ── Host-op planning ──────────────────────────────────────────────────
837 const noHostOpsInLoops = (what: string, span: Span): void => {
838 if (inLoop()) {
839 throw new UnsupportedOnGpu(
840 `'${what}' inside a for loop is not supported: the loop body is ` +
841 `compiled once and replayed on the GPU, so per-iteration host I/O ` +
842 `has nowhere to run. Hoist it out of the loop (or time the whole loop).`,
843 span,
844 );
845 }
846 };
848 /** A ValueRef for a scalar-valued expression the host wants to print. The
849 * inline pass folds argument temps into the call, so this routinely sees
850 * whole expressions — they get computed into a 1-element buffer. */
851 async function scalarRef(e: IRExpr, span: Span): Promise<ValueRef> {
852 if (isNumeric(e.ty) && isTensor(e.ty)) {
853 throw new UnsupportedOnGpu(
854 `printing an array here is not supported (MATLAB would recycle the ` +
855 `format); print a scalar, or use disp`,
856 span,
857 );
858 }
859 const exact = exactValue(e);
860 if (exact !== undefined) return { kind: 'literal', value: exact };
861 if (e.kind === 'Var') {
862 if (hostScalars.has(e.cName)) return { kind: 'host', cName: e.cName };
863 const slot = readSlot(e.cName, e.name, span);
864 return { kind: 'buffer', slot, count: 1 };
865 }
866 if (!isNumeric(e.ty)) {
867 throw new UnsupportedOnGpu(`only numeric values can be printed`, span);
868 }
869 const cName = `%tmp${tempCounter++}`;
870 await planValue(cName, cName, e.ty, e, span);
871 return { kind: 'buffer', slot: vars.get(resolve(cName))!.slot, count: 1 };
872 }
874 async function planDisp(call: Call, span: Span): Promise<void> {
875 noHostOpsInLoops('disp', span);
876 if (call.args.length !== 1) {
877 throw new UnsupportedOnGpu(`'disp' takes one argument`, span);
878 }
879 const a = call.args[0];
880 if (a.kind === 'StringLit') {
881 ops().push({ kind: 'emit', parts: [{ kind: 'text', text: a.value + '\n' }] });
882 return;
883 }
884 if (a.kind === 'Var' && chars.has(a.cName)) {
885 ops().push({ kind: 'emit', parts: [{ kind: 'text', text: chars.get(a.cName)! + '\n' }] });
886 return;
887 }
888 if (!isNumeric(a.ty)) throw new UnsupportedOnGpu(`'disp' argument is not numeric`, span);
889 if (isTensor(a.ty)) {
890 if (a.kind !== 'Var') {
891 throw new UnsupportedOnGpu(`'disp' of an expression — give it a name first`, span);
892 }
893 const slot = readSlot(a.cName, a.name, span);
894 ops().push({
895 kind: 'display',
896 label: null,
897 ref: { kind: 'buffer', slot, count: slot.count },
898 shape: a.ty.shape ?? [slot.count, 1],
899 });
900 return;
901 }
902 ops().push({ kind: 'display', label: null, ref: await scalarRef(a, span), shape: [1, 1] });
903 }
905 async function planFprintf(call: Call, span: Span): Promise<void> {
906 noHostOpsInLoops('fprintf', span);
907 let args = call.args;
908 // fprintf(1, fmt, ...) — MATLAB's stdout file id.
909 if (args.length >= 2 && args[0].kind === 'NumLit' && args[0].value === 1) {
910 args = args.slice(1);
911 }
912 if (args.length === 0) throw new UnsupportedOnGpu(`'fprintf' needs a format string`, span);
913 const fmtArg = args[0];
914 const fmt =
915 fmtArg.kind === 'StringLit'
916 ? fmtArg.value
917 : fmtArg.kind === 'Var' && chars.has(fmtArg.cName)
918 ? chars.get(fmtArg.cName)!
919 : null;
920 if (fmt === null) {
921 throw new UnsupportedOnGpu(`'fprintf' format must be a literal string`, span);
922 }
923 const parts = parseFormat(fmt, span);
924 const specs = parts.filter((p) => p.kind === 'spec');
925 const values = args.slice(1);
926 if (specs.length !== values.length) {
927 throw new UnsupportedOnGpu(
928 `'fprintf' format has ${specs.length} conversion(s) but ${values.length} ` +
929 `value(s); the sandbox does not recycle the format over arrays`,
930 span,
931 );
932 }
933 let vi = 0;
934 const emitParts: EmitPart[] = [];
935 for (const p of parts) {
936 if (p.kind === 'text') emitParts.push({ kind: 'text', text: p.text });
937 else emitParts.push({ kind: 'value', ref: await scalarRef(values[vi++], span), spec: p.spec });
938 }
939 ops().push({ kind: 'emit', parts: emitParts });
940 }
942 /** printf-format split: text runs (escapes decoded) and % conversions. */
943 function parseFormat(
944 fmt: string,
945 span: Span,
946 ): ({ kind: 'text'; text: string } | { kind: 'spec'; spec: string })[] {
947 const out: ({ kind: 'text'; text: string } | { kind: 'spec'; spec: string })[] = [];
948 let text = '';
949 for (let i = 0; i < fmt.length; i++) {
950 const c = fmt[i];
951 if (c === '\\') {
952 const n = fmt[i + 1];
953 if (n === 'n') { text += '\n'; i++; }
954 else if (n === 't') { text += '\t'; i++; }
955 else if (n === '\\') { text += '\\'; i++; }
956 else text += c;
957 } else if (c === '%') {
958 if (fmt[i + 1] === '%') { text += '%'; i++; continue; }
959 const m = /^%[-+ 0#]*\d*(?:\.\d+)?[diufeEgGs]/.exec(fmt.slice(i));
960 if (!m) {
961 throw new UnsupportedOnGpu(
962 `'fprintf': unsupported conversion at "${fmt.slice(i, i + 6)}"`,
963 span,
964 );
965 }
966 if (text) { out.push({ kind: 'text', text }); text = ''; }
967 out.push({ kind: 'spec', spec: m[0] });
968 i += m[0].length - 1;
969 } else {
970 text += c;
971 }
972 }
973 if (text) out.push({ kind: 'text', text });
974 return out;
975 }
977 // ── Statement walk ────────────────────────────────────────────────────
979 async function planStmt(stmt: IRStmt): Promise<void> {
980 switch (stmt.kind) {
981 case 'Assign':
982 return planAssign(stmt);
983 case 'ExprStmt':
984 return planExprStmt(stmt);
985 case 'For':
986 return planFor(stmt);
987 default:
988 throw new UnsupportedOnGpu(
989 `'${stmtName(stmt.kind)}' is not supported in the sandbox`,
990 stmt.span,
991 );
992 }
993 }
995 function stmtName(kind: string): string {
996 return (
997 {
998 If: 'if', While: 'while', Break: 'break', Continue: 'continue',
999 IndexStore: 'indexed assignment', IndexSliceStore: 'indexed assignment',
1000 MultiAssignCall: 'multiple assignment', MemberStore: 'struct assignment',
1001 CellIndexStore: 'cell assignment',
1002 }[kind] ?? kind
1003 );
1006 const isTemp = (name: string): boolean => name.startsWith('_mtoc2_') || name.startsWith('%tmp');
1008 async function planAssign(stmt: Assign): Promise<void> {
1009 // Char values (format strings) ride along on the host. Their type kind
1010 // is numbl's 'Char'/'String', not Numeric.
1011 if (stmt.expr.kind === 'StringLit') {
1012 chars.set(stmt.cName, stmt.expr.value);
1013 return;
1015 if (!isNumeric(stmt.ty)) {
1016 throw new UnsupportedOnGpu(
1017 `'${stmt.name}' is not a numeric value (only numbers, strings for ` +
1018 `printing, and numeric arrays exist in the sandbox)`,
1019 stmt.span,
1020 );
1023 // tic/toc as values.
1024 if (stmt.expr.kind === 'Call' && (stmt.expr.name === 'tic' || stmt.expr.name === 'toc')) {
1025 noHostOpsInLoops(stmt.expr.name, stmt.span);
1026 const slot = slotFor(stmt.cName, stmt.name, stmt.ty, stmt.span);
1027 hostScalars.add(stmt.cName);
1028 if (stmt.expr.name === 'tic') {
1029 ops().push({ kind: 'tic', assignTo: { cName: stmt.cName, slot } });
1030 } else {
1031 const since = tocSince(stmt.expr, stmt.span);
1032 ops().push({
1033 kind: 'toc', print: false, sinceCName: since,
1034 assignTo: { cName: stmt.cName, slot }, seq: ++tocCounter,
1035 });
1037 maybeEcho(stmt);
1038 return;
1041 // Exact scalar: folds into every consumer. It only needs real storage if
1042 // the variable is reassigned elsewhere (later uses then read the buffer).
1043 if (typeof stmt.ty.exact === 'number') {
1044 if ((assignCounts.get(stmt.cName) ?? 1) > 1) {
1045 const slot = slotFor(stmt.cName, stmt.name, stmt.ty, stmt.span);
1046 ops().push({
1047 kind: 'write', slot,
1048 data: new Float32Array([stmt.ty.exact]),
1049 label: `${stmt.name} = ${stmt.ty.exact}`,
1050 });
1052 maybeEcho(stmt);
1053 return;
1055 // Exact tensor (a literal like [1 2 3]): upload the data.
1056 if (stmt.ty.exact instanceof Float64Array) {
1057 const slot = slotFor(stmt.cName, stmt.name, stmt.ty, stmt.span);
1058 ops().push({
1059 kind: 'write', slot,
1060 data: Float32Array.from(stmt.ty.exact),
1061 label: `${stmt.name} = <literal ${slot.count} elem>`,
1062 });
1063 maybeEcho(stmt);
1064 return;
1066 if (stmt.ty.exact !== undefined) {
1067 throw new UnsupportedOnGpu(`complex values are not supported (f32 backend)`, stmt.span);
1070 await planValue(stmt.cName, stmt.name, stmt.ty, stmt.expr, stmt.span);
1071 maybeEcho(stmt);
1074 /** MATLAB-style echo for statements without a trailing semicolon. */
1075 function maybeEcho(stmt: Assign): void {
1076 if (isTemp(stmt.name)) return;
1077 if (!compiled.isEchoed(stmt.span.start)) return;
1078 noHostOpsInLoops(`echo of '${stmt.name}' (add a semicolon)`, stmt.span);
1079 if (!isNumeric(stmt.ty)) return;
1080 const exact = typeof stmt.ty.exact === 'number' ? stmt.ty.exact : undefined;
1081 const shape = stmt.ty.shape ?? [1, 1];
1082 if (exact !== undefined) {
1083 ops().push({
1084 kind: 'display', label: stmt.name,
1085 ref: { kind: 'literal', value: exact }, shape,
1086 });
1087 return;
1089 if (hostScalars.has(stmt.cName)) {
1090 ops().push({
1091 kind: 'display', label: stmt.name,
1092 ref: { kind: 'host', cName: stmt.cName }, shape,
1093 });
1094 return;
1096 const v = vars.get(resolve(stmt.cName));
1097 if (!v) return;
1098 ops().push({
1099 kind: 'display', label: stmt.name,
1100 ref: { kind: 'buffer', slot: v.slot, count: v.slot.count }, shape,
1101 });
1104 function tocSince(call: IRExpr & { kind: 'Call' }, span: Span): string | undefined {
1105 if (call.args.length === 0) return undefined;
1106 const a = call.args[0];
1107 if (a.kind === 'Var' && hostScalars.has(a.cName)) return a.cName;
1108 throw new UnsupportedOnGpu(
1109 `'toc(t)' needs a value produced by 't = tic'`,
1110 span,
1111 );
1114 async function planExprStmt(stmt: ExprStmt): Promise<void> {
1115 const e = stmt.expr;
1116 if (e.kind === 'Call') {
1117 switch (e.name) {
1118 case 'tic':
1119 noHostOpsInLoops('tic', stmt.span);
1120 ops().push({ kind: 'tic' });
1121 return;
1122 case 'toc':
1123 case 'toc_print': // numbl's lowering of a bare `toc` statement
1124 noHostOpsInLoops('toc', stmt.span);
1125 ops().push({
1126 kind: 'toc', print: true,
1127 sinceCName: tocSince(e, stmt.span), seq: ++tocCounter,
1128 });
1129 return;
1130 case 'disp':
1131 return planDisp(e, stmt.span);
1132 case 'fprintf':
1133 return planFprintf(e, stmt.span);
1134 case 'rng':
1135 throw new UnsupportedOnGpu(
1136 `'rng' is not supported: the sandbox's rand/randn streams are ` +
1137 `deterministic per run already`,
1138 stmt.span,
1139 );
1140 default:
1141 break;
1144 // A bare expression with a value: numbl assigns display-relevant results
1145 // to `ans` as an Assign, so a leftover ExprStmt is side-effect-free.
1146 if (e.kind === 'Call') {
1147 throw new UnsupportedOnGpu(`'${e.name}' is not supported in the sandbox`, stmt.span);
1151 async function planFor(stmt: For): Promise<void> {
1152 const from = exactValue(stmt.start);
1153 const to = exactValue(stmt.end);
1154 if (from === undefined || to === undefined) {
1155 throw new UnsupportedOnGpu(
1156 `a for loop's bounds must be known when the script is compiled ` +
1157 `(assign them from literals)`,
1158 stmt.span,
1159 );
1161 const trips = Math.floor((to - from) / stmt.step + 1e-9) + 1;
1162 if (trips <= 0) return; // never executes
1163 if (trips > MAX_TRIPS) {
1164 throw new UnsupportedOnGpu(
1165 `'for ${stmt.varName}' runs ${trips} iterations, over the sandbox ` +
1166 `limit of ${MAX_TRIPS}`,
1167 stmt.span,
1168 );
1171 // One 256-byte uniform slot per iteration: { v: f32, it: u32 }.
1172 const uniform = device.createBuffer({
1173 label: `mgpu-loop-${stmt.varName}`,
1174 size: trips * LV_STRIDE,
1175 usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
1176 });
1177 owned.push(uniform);
1178 const data = new ArrayBuffer(trips * LV_STRIDE);
1179 const f32 = new Float32Array(data);
1180 const u32 = new Uint32Array(data);
1181 for (let i = 0; i < trips; i++) {
1182 f32[(i * LV_STRIDE) / 4] = from + i * stmt.step;
1183 u32[(i * LV_STRIDE) / 4 + 1] = i;
1185 device.queue.writeBuffer(uniform, 0, data);
1186 loopUniforms.set(stmt.cVar, uniform);
1188 const body: Op[] = [];
1189 opStack.push(body);
1190 enclosingLoops.push(stmt.cVar);
1191 loopTrips.push(trips);
1192 try {
1193 for (const s of stmt.body) await planStmt(s);
1194 } finally {
1195 loopTrips.pop();
1196 enclosingLoops.pop();
1197 opStack.pop();
1199 ops().push({
1200 kind: 'loop', cVar: stmt.cVar, trips, uniform, body,
1201 label: `for ${stmt.varName} = ${from}:${stmt.step}:${to}`,
1202 });
1203 describeLines.push(`loop for ${stmt.varName} (${trips} iterations, body above)`);
1205 // MATLAB leaves the loop variable holding its final value; make that
1206 // readable afterwards.
1207 const finalValue = from + (trips - 1) * stmt.step;
1208 const slot = slotFor(
1209 stmt.cVar, stmt.varName,
1210 { kind: 'Numeric', elem: 'double', isComplex: false,
1211 dims: [{ kind: 'exact', value: 1 }, { kind: 'exact', value: 1 }],
1212 shape: [1, 1], sign: 'unknown' },
1213 stmt.span,
1214 );
1215 ops().push({
1216 kind: 'write', slot, data: new Float32Array([finalValue]),
1217 label: `${stmt.varName} = ${finalValue} (final)`,
1218 });
1221 for (const stmt of compiled.stmts) {
1222 await planStmt(stmt);
1225 return {
1226 ops: opStack[0],
1227 describe: () => describeLines,
4ded6ffShow the generated WGSL kernels in the compiled-plan paneJeremy Magland 1228 kernels: kernelSources,
1230 for (const b of owned) b.destroy();
1231 owned.length = 0;
1232 },
1233 };
moveopenescclose