1// WebGPU host for the megakernel: device setup (Dawn in Node via the
2// optional `webgpu` package; in a browser this maps onto navigator.gpu),
3// buffer upload, the dispatch queue, and readback.
5export async function installNodeWebGpu() {
6 const mod = await import('webgpu');
7 Object.assign(globalThis, mod.globals);
8 const dawnFlags = process.env.DAWN_FLAGS?.split(';').filter(Boolean) ?? [];
9 Object.defineProperty(globalThis, 'navigator', {
10 value: { gpu: mod.create(dawnFlags) },
11 configurable: true,
12 writable: true,
13 });
14}
16export async function requestDevice() {
17 const adapter = await navigator.gpu.requestAdapter();
18 if (!adapter) throw new Error('no WebGPU adapter');
19 const L = adapter.limits;
20 const device = await adapter.requestDevice({
21 requiredLimits: {
22 maxComputeInvocationsPerWorkgroup: L.maxComputeInvocationsPerWorkgroup,
23 maxComputeWorkgroupSizeX: L.maxComputeWorkgroupSizeX,
24 maxComputeWorkgroupStorageSize: L.maxComputeWorkgroupStorageSize,
25 },
26 });
27 device.addEventListener?.('uncapturederror', (e) => {
28 console.error('WebGPU uncaptured error:', e.error?.message ?? e);
29 });
30 const info = adapter.info
31 ? `${adapter.info.vendor ?? '?'} ${adapter.info.architecture ?? ''} ${adapter.info.description ?? ''}`.trim()
32 : 'unknown adapter';
33 return { device, info };
34}
36// Multi-dispatch backend: 8 pipelines, one bind group with a dynamic-offset
37// uniform carrying the within-chunk step index; the host writes the chunk
38// base into state cell 0 before each chunk and keeps two chunks in flight.
39export async function runMultiDispatch(device, L, stages, onProgress) {
40 const nsteps = L.nsteps;
41 const bgl = device.createBindGroupLayout({
42 entries: [
43 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
44 { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
45 { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
46 { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform', hasDynamicOffset: true, minBindingSize: 16 } },
47 ],
48 });
49 const layout = device.createPipelineLayout({ bindGroupLayouts: [bgl] });
51 const pipelines = [];
52 const simpleBGs = [];
53 for (const name of stages.order) {
54 const module = device.createShaderModule({ code: stages.sources[name] });
55 const errs = (await module.getCompilationInfo()).messages.filter((m) => m.type === 'error');
56 if (errs.length) {
57 throw new Error(`stage ${name} WGSL errors:\n` +
58 errs.map((m) => ` ${m.lineNum}:${m.linePos} ${m.message}`).join('\n'));
59 }
60 pipelines.push(device.createComputePipeline({ layout, compute: { module, entryPoint: 'main' } }));
61 }
63 const chunk = Number(process.env.CHUNK ?? 1024);
64 const mkBuf = (size, usage) => device.createBuffer({ size, usage });
65 const opsBuf = mkBuf(L.opsData.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
66 device.queue.writeBuffer(opsBuf, 0, L.opsData);
67 const stBuf = mkBuf(L.stLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST);
68 const outBuf = mkBuf(2 * nsteps * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
69 const uniBuf = mkBuf(chunk * 256, GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
70 {
71 const uni = new Uint32Array(chunk * 64);
72 for (let s = 0; s < chunk; s++) uni[s * 64] = s;
73 device.queue.writeBuffer(uniBuf, 0, uni);
74 }
75 const readOut = mkBuf(2 * nsteps * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
76 const readSt = mkBuf(L.stLen * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
78 const bindGroup = device.createBindGroup({
79 layout: bgl,
80 entries: [
81 { binding: 0, resource: { buffer: opsBuf } },
82 { binding: 1, resource: { buffer: stBuf } },
83 { binding: 2, resource: { buffer: outBuf } },
84 { binding: 3, resource: { buffer: uniBuf, offset: 0, size: 16 } },
85 ],
86 });
88 const t0 = performance.now();
89 const pending = [];
90 let kbase = 0;
91 let recMs = 0;
92 const inflight = Number(process.env.INFLIGHT ?? 2);
93 while (kbase < nsteps) {
94 const n = Math.min(chunk, nsteps - kbase);
95 device.queue.writeBuffer(stBuf, 0, new Float32Array([kbase]));
96 const tr = performance.now();
97 const enc = device.createCommandEncoder();
98 const pass = enc.beginComputePass();
100 for (let s = 0; s < n; s++) {
101 pass.setBindGroup(0, bindGroup, [s * 256]);
102 for (const p of pipelines) {
103 pass.setPipeline(p);
104 pass.dispatchWorkgroups(stages.nwg);
105 }
106 }
107 pass.end();
108 const cb = enc.finish();
109 recMs += performance.now() - tr;
110 device.queue.submit([cb]);
111 pending.push(device.queue.onSubmittedWorkDone());
112 if (pending.length >= inflight) await pending.shift();
113 kbase += n;
114 if (onProgress) onProgress(kbase, nsteps, performance.now() - t0);
115 }
116 while (pending.length) await pending.shift();
117 const gpuMs = performance.now() - t0;
118 console.log(`\n [recording ${(1000 * recMs / nsteps).toFixed(1)} us/step of ${(1000 * gpuMs / nsteps).toFixed(1)} wall]`);
120 const enc = device.createCommandEncoder();
121 enc.copyBufferToBuffer(outBuf, 0, readOut, 0, 2 * nsteps * 4);
122 enc.copyBufferToBuffer(stBuf, 0, readSt, 0, L.stLen * 4);
123 device.queue.submit([enc.finish()]);
124 await readOut.mapAsync(GPUMapMode.READ);
125 await readSt.mapAsync(GPUMapMode.READ);
126 const outData = new Float32Array(readOut.getMappedRange().slice(0));
127 const stData = new Float32Array(readSt.getMappedRange().slice(0));
128 readOut.unmap();
129 readSt.unmap();
130 for (const b of [opsBuf, stBuf, outBuf, uniBuf, readOut, readSt]) b.destroy();
132 return {
133 acc: outData.subarray(0, nsteps),
134 fhist: outData.subarray(nsteps, 2 * nsteps),
135 state: stData,
136 gpuMs,
137 };
138}
141// Persistent-kernel backend: identical dispatches, each integrating CHUNK
142// steps with global spin barriers between stages. sync[] holds the barrier
143// atomics; state cell 0 (kbase) advances GPU-side at the end of each chunk.
144export async function runPersistent(device, L, pk, onProgress) {
145 const nsteps = L.nsteps;
146 const module = device.createShaderModule({ code: pk.code });
147 const errs = (await module.getCompilationInfo()).messages.filter((m) => m.type === 'error');
148 if (errs.length) {
149 throw new Error('persistent kernel WGSL errors:\n' +
150 errs.map((m) => ` ${m.lineNum}:${m.linePos} ${m.message}`).join('\n'));
151 }
152 const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
154 const mkBuf = (size, usage) => device.createBuffer({ size, usage });
155 const opsBuf = mkBuf(L.opsData.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
156 device.queue.writeBuffer(opsBuf, 0, L.opsData);
157 const stBuf = mkBuf(L.stLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
158 const outBuf = mkBuf(2 * nsteps * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
159 const syncBuf = mkBuf(16, GPUBufferUsage.STORAGE);
160 const readOut = mkBuf(2 * nsteps * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
161 const readSt = mkBuf(L.stLen * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
163 const bindGroup = device.createBindGroup({
164 layout: pipeline.getBindGroupLayout(0),
165 entries: [
166 { binding: 0, resource: { buffer: opsBuf } },
167 { binding: 1, resource: { buffer: stBuf } },
168 { binding: 2, resource: { buffer: outBuf } },
169 { binding: 3, resource: { buffer: syncBuf } },
170 ],
171 });
173 const totalDispatches = Math.ceil(nsteps / pk.chunk);
174 const perSubmit = 8;
175 const t0 = performance.now();
176 let done = 0;
177 const pending = [];
178 while (done < totalDispatches) {
179 const n = Math.min(perSubmit, totalDispatches - done);
180 const enc = device.createCommandEncoder();
181 for (let i = 0; i < n; i++) {
182 const pass = enc.beginComputePass();
183 pass.setPipeline(pipeline);
184 pass.setBindGroup(0, bindGroup);
185 pass.dispatchWorkgroups(pk.nwg);
186 pass.end();
187 }
188 device.queue.submit([enc.finish()]);
189 pending.push(device.queue.onSubmittedWorkDone());
190 if (pending.length >= 2) await pending.shift();
191 done += n;
192 if (onProgress) onProgress(Math.min(done * pk.chunk, nsteps), nsteps, performance.now() - t0);
193 }
194 while (pending.length) await pending.shift();
195 const gpuMs = performance.now() - t0;
197 const enc = device.createCommandEncoder();
198 enc.copyBufferToBuffer(outBuf, 0, readOut, 0, 2 * nsteps * 4);
199 enc.copyBufferToBuffer(stBuf, 0, readSt, 0, L.stLen * 4);
200 device.queue.submit([enc.finish()]);
201 await readOut.mapAsync(GPUMapMode.READ);
202 await readSt.mapAsync(GPUMapMode.READ);
203 const outData = new Float32Array(readOut.getMappedRange().slice(0));
204 const stData = new Float32Array(readSt.getMappedRange().slice(0));
205 readOut.unmap();
206 readSt.unmap();
207 for (const b of [opsBuf, stBuf, outBuf, syncBuf, readOut, readSt]) b.destroy();
209 return {
210 acc: outData.subarray(0, nsteps),
211 fhist: outData.subarray(nsteps, 2 * nsteps),
212 state: stData,
213 gpuMs,
214 };
215}
217export async function runMegakernel(device, kernel, nsteps, onProgress) {
218 const module = device.createShaderModule({ code: kernel.code });
219 const compInfo = await module.getCompilationInfo();
220 const errors = compInfo.messages.filter((m) => m.type === 'error');
221 if (errors.length) {
222 throw new Error(
223 'WGSL compile errors:\n' +
224 errors.map((m) => ` ${m.lineNum}:${m.linePos} ${m.message}`).join('\n'),
225 );
226 }
228 const pipeline = device.createComputePipeline({
229 layout: 'auto',
230 compute: { module, entryPoint: 'main' },
231 });
233 const mkBuf = (size, usage) => device.createBuffer({ size, usage });
234 const opsBuf = mkBuf(kernel.opsData.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
235 device.queue.writeBuffer(opsBuf, 0, kernel.opsData);
236 const stBuf = mkBuf(kernel.stLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
237 const outBuf = mkBuf(2 * nsteps * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
238 const readOut = mkBuf(2 * nsteps * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
239 const readSt = mkBuf(kernel.stLen * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
241 const bindGroup = device.createBindGroup({
242 layout: pipeline.getBindGroupLayout(0),
243 entries: [
244 { binding: 0, resource: { buffer: opsBuf } },
245 { binding: 1, resource: { buffer: stBuf } },
246 { binding: 2, resource: { buffer: outBuf } },
247 ],
248 });
250 const totalDispatches = Math.ceil(nsteps / kernel.chunk);
251 const perSubmit = 16;
252 const t0 = performance.now();
253 let done = 0;
254 while (done < totalDispatches) {
255 const n = Math.min(perSubmit, totalDispatches - done);
256 const enc = device.createCommandEncoder();
257 for (let i = 0; i < n; i++) {
258 const pass = enc.beginComputePass();
259 pass.setPipeline(pipeline);
260 pass.setBindGroup(0, bindGroup);
261 pass.dispatchWorkgroups(1);
262 pass.end();
263 }
264 device.queue.submit([enc.finish()]);
265 await device.queue.onSubmittedWorkDone();
266 done += n;
267 if (onProgress) onProgress(Math.min(done * kernel.chunk, nsteps), nsteps, performance.now() - t0);
268 }
269 const gpuMs = performance.now() - t0;
271 const enc = device.createCommandEncoder();
272 enc.copyBufferToBuffer(outBuf, 0, readOut, 0, 2 * nsteps * 4);
273 enc.copyBufferToBuffer(stBuf, 0, readSt, 0, kernel.stLen * 4);
274 device.queue.submit([enc.finish()]);
275 await readOut.mapAsync(GPUMapMode.READ);
276 await readSt.mapAsync(GPUMapMode.READ);
277 const outData = new Float32Array(readOut.getMappedRange().slice(0));
278 const stData = new Float32Array(readSt.getMappedRange().slice(0));
279 readOut.unmap();
280 readSt.unmap();
281 for (const b of [opsBuf, stBuf, outBuf, readOut, readSt]) b.destroy();
283 return {
284 acc: outData.subarray(0, nsteps),
285 fhist: outData.subarray(nsteps, 2 * nsteps),
286 state: stData,
287 gpuMs,
288 };
289}