1import type { BackendInfo, ComputeInfo, RunRequest, RunResult } from "./types";
2import { executeBundle, formatError } from "./execute";
3import { LIBRARY_FILES } from "../generate/library";
4import { runFdLoopGpu, type GpuHandle } from "../gpu/fdLoop";
5import { fdLoopRecordToPlain, fdLoopSetupFromPlain } from "../gpu/fdStage";
7/** Supplies the (cached) WebGPU device of the host, or rejects when there is none. */
8export type GpuProvider = () => Promise<GpuHandle>;
10/**
11 * Run a request, on the GPU when it asks for that and the host has one. The
12 * plain path is executeBundle. The staged path runs the setup script in numbl,
13 * the loop on the GPU and the output script in numbl, and reports where the
14 * loop ran in `compute`; when the GPU was asked for but cannot be used, the
15 * whole bundle runs in numbl and `compute.note` says why.
16 */
17export async function executeRequest(
18 req: RunRequest,
19 backend: BackendInfo,
20 onOutput?: (text: string) => void,
21 gpu?: GpuProvider,
22): Promise<RunResult> {
23 const stages = req.gpu;
24 if (!stages) return executeBundle(req, backend, onOutput);
25 const fallback = (note: string): RunResult => ({
26 ...executeBundle(req, backend, onOutput),
27 compute: { loop: "numbl", note },
28 });
30 if (!gpu) return fallback("WebGPU is not available here, so the time loop ran in numbl");
31 const shipped = LIBRARY_FILES.find((f) => f.name === stages.loopFile)?.source;
32 const current = req.bundle.files.find((f) => f.name === stages.loopFile)?.source;
33 if (shipped !== undefined && current !== shipped)
34 return fallback(
35 `${stages.loopFile} was edited, so the time loop ran in numbl, where the edit is live`,
36 );
37 let handle: GpuHandle;
38 try {
39 handle = await gpu();
40 } catch (e) {
41 return fallback(`no WebGPU device (${formatError(e)}), so the time loop ran in numbl`);
42 }
44 const t0 = performance.now();
45 const output: string[] = [];
46 const say = (text: string) => {
47 output.push(text);
48 onOutput?.(text);
49 };
50 const libs = req.bundle.files.filter((f) => f.name !== req.bundle.entry);
51 const fail = (error: string, compute?: ComputeInfo): RunResult => ({
52 ok: false,
53 error,
54 output,
55 wallMs: performance.now() - t0,
56 backend,
57 compute,
58 });
60 const setupRun = executeBundle(
61 {
62 bundle: {
63 files: [{ name: "gpu_setup.m", source: stages.setup }, ...libs],
64 entry: "gpu_setup.m",
65 resultVariable: "S",
66 },
67 },
68 backend,
69 say,
70 );
71 if (!setupRun.ok) return fail(setupRun.error);
72 let setup;
73 try {
74 setup = fdLoopSetupFromPlain(setupRun.value);
75 } catch (e) {
76 return fail(`the setup stage returned an unexpected shape: ${formatError(e)}`);
77 }
79 let lastPct = 0;
80 let rec;
81 try {
82 rec = await runFdLoopGpu(handle.device, setup, (done, total) => {
83 const pct = Math.floor((10 * done) / total);
84 if (pct > lastPct && done < total) {
85 lastPct = pct;
86 say(`gpu: ${done.toLocaleString()} of ${total.toLocaleString()} steps\n`);
87 }
88 });
89 } catch (e) {
90 return fail(`WebGPU loop: ${formatError(e)}`, { loop: "webgpu", adapter: handle.adapter });
91 }
93 const R = fdLoopRecordToPlain(rec, setup, handle.adapter);
94 const outRun = executeBundle(
95 {
96 bundle: {
97 files: [{ name: "gpu_output.m", source: stages.output }, ...libs],
98 entry: "gpu_output.m",
99 resultVariable: req.bundle.resultVariable,
100 },
101 inputs: { ...(req.inputs ?? {}), R },
102 },
103 backend,
104 say,
105 );
106 const compute: ComputeInfo = {
107 loop: "webgpu",
108 adapter: handle.adapter,
109 gpuMs: rec.gpuMs,
110 numblMs: setupRun.wallMs + outRun.wallMs,
111 };
112 if (!outRun.ok) return fail(outRun.error, compute);
113 return {
114 ok: true,
115 value: outRun.value,
116 output,
117 wallMs: performance.now() - t0,
118 backend,
119 compute,
120 };
121}