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 */
14import { WORKGROUP_SIZE } from './wgsl.ts';
15import type { EmitPart, Op, ScriptPlan, Slot, ValueRef } from './plan.ts';
17void WORKGROUP_SIZE;
19export interface TimingSegment {
20 /** 1-based tic..toc pair index. */
21 seq: number;
22 seconds: number;
23}
25export 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}
34const LV_STRIDE = 256;
36export 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 };
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 };
76 /** Values of tic/toc-produced variables, in seconds. */
77 const hostVals = new Map<string, number>();
78 let ticStartMs: number | null = null;
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 };
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 };
108 const writeScalar = (slot: Slot, value: number): void => {
109 device.queue.writeBuffer(slot.buffer, 0, new Float32Array([value]) as Float32Array<ArrayBuffer>);
110 };
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 }
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 }
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 }
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;
245 return { output, segments, totalSeconds, error };
246}
248/** MATLAB `format short`-flavored scalar rendering. */
249export 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}
258/** One printf-style conversion. */
259function 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}