1// Multi-dispatch step: each stage its own dispatch at full GPU width.
2// Representative step = 2 stencil stages + 4 transforms + 2 page solves,
3// recorded 512 steps per command buffer.
4import { installNodeWebGpu, requestDevice } from './gpu.mjs';
6const NRT = 26, NTH = 60, N = NRT * NTH;
7const WG = Number(process.env.WG ?? 128);
8const NWG = Math.ceil(N / WG);
9const STEPS = 512;
11const DYNOFF = !!process.env.DYNOFF;
12const GATE = !!process.env.GATE;
13const HDR = `
14const NRT = ${NRT}u; const NTH = ${NTH}u; const N = ${N}u;
15@group(0) @binding(0) var<storage, read> ops: array<f32>;
16@group(0) @binding(1) var<storage, read_write> st: array<f32>;
17${DYNOFF ? '@group(0) @binding(2) var<uniform> ctl: vec4<u32>;' : ''}
18`;
20const KERNELS = {
21 stencil: `${HDR}
22@compute @workgroup_size(${WG})
23fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
24 let i = gid.x;
25 if (i >= N) { return; }
26 ${GATE ? 'if (st[1] == 0.0) { return; }' : ''}
27 let l = i % NRT; let j = i / NRT;
28 let jm = (j + NTH - 1u) % NTH; let jp = (j + 1u) % NTH;
29 var v: f32 = 0.0;
30 if (l >= 1u && l <= NRT - 2u) {
31 v = (st[N + (l - 1u) + j * NRT] - 2.0 * st[N + l + j * NRT] + st[N + (l + 1u) + j * NRT]) * 0.5
32 + (st[N + l + jm * NRT] - st[N + l + jp * NRT]) * 0.25;
33 }
34 st[2u * N + i] = v;
35}`,
36 transform: `${HDR}
37@compute @workgroup_size(${WG})
38fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
39 let i = gid.x;
40 if (i >= N) { return; }
41 ${GATE ? 'if (st[1] == 0.0) { return; }' : ''}
42 let l = i % NRT; let j = i / NRT;
43 var sum: f32 = 0.0;
44 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
45 sum = sum + st[N + l + q * NRT] * ops[q + j * NTH];
46 }
47 st[2u * N + i] = sum + f32(${DYNOFF ? 'ctl.x' : '0u'});
48}`,
49 page: `${HDR}
50@compute @workgroup_size(${WG})
51fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
52 let i = gid.x;
53 if (i >= N) { return; }
54 ${GATE ? 'if (st[1] == 0.0) { return; }' : ''}
55 let l = i % NRT; let j = i / NRT;
56 let base = j * NRT * NRT;
57 var sum: f32 = 0.0;
58 for (var p: u32 = 0u; p < NRT; p = p + 1u) {
59 sum = sum + ops[base + l + p * NRT] * st[N + p + j * NRT];
60 }
61 st[2u * N + i] = sum;
62}`,
63};
65async function main() {
66 await installNodeWebGpu();
67 const { device, info } = await requestDevice();
68 console.log(`device: ${info}, wg size ${WG} x ${NWG} workgroups`);
70 const BIG = Number(process.env.BIG ?? 1);
71 const ops = device.createBuffer({ size: 4 * NRT * NRT * NTH * BIG, usage: GPUBufferUsage.STORAGE });
72 const stb = device.createBuffer({ size: 4 * N * 4 * BIG, usage: GPUBufferUsage.STORAGE });
73 const uni = device.createBuffer({ size: STEPS * 256, usage: GPUBufferUsage.UNIFORM });
75 const pipes = {};
76 for (const [name, code] of Object.entries(KERNELS)) {
77 const module = device.createShaderModule({ code });
78 const errs = (await module.getCompilationInfo()).messages.filter((m) => m.type === 'error');
79 if (errs.length) throw new Error(name + ': ' + errs.map((m) => m.message).join('\n'));
80 const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
81 const entries = [
82 { binding: 0, resource: { buffer: ops } },
83 { binding: 1, resource: { buffer: stb } },
84 ];
85 if (DYNOFF && name === 'transform') entries.push({ binding: 2, resource: { buffer: uni, offset: 0, size: 16 } });
86 const bg = device.createBindGroup({
87 layout: pipeline.getBindGroupLayout(0),
88 entries,
89 });
90 pipes[name] = { pipeline, bg };
91 }
93 let stepSeq = ['stencil', 'transform', 'page', 'transform', 'stencil', 'transform', 'page', 'transform'];
94 if (process.env.SEQ1) stepSeq = ['transform'];
95 if (process.env.SEQN) stepSeq = Array(Number(process.env.SEQN)).fill('transform');
96 const record = () => {
97 const enc = device.createCommandEncoder();
98 const pass = enc.beginComputePass();
99 for (let s = 0; s < STEPS; s++) {
100 for (const name of stepSeq) {
101 const p = pipes[name];
102 pass.setPipeline(p.pipeline);
103 pass.setBindGroup(0, p.bg);
104 pass.dispatchWorkgroups(NWG);
105 }
106 }
107 pass.end();
108 return enc.finish();
109 };
111 device.queue.writeBuffer(stb, 4, new Float32Array([1]));
112 // warmup
113 device.queue.submit([record()]);
114 await device.queue.onSubmittedWorkDone();
116 let tRec = 0;
117 const t0 = performance.now();
118 const reps = Number(process.env.REPS ?? 4);
119 for (let r = 0; r < reps; r++) {
120 const tr = performance.now();
121 const cb = record();
122 tRec += performance.now() - tr;
123 device.queue.submit([cb]);
124 }
125 await device.queue.onSubmittedWorkDone();
126 const ms = (performance.now() - t0) / reps;
127 console.log(`8-dispatch step: ${(1000 * ms / STEPS).toFixed(2)} us/step total (cpu recording ${(1000 * tRec / reps / STEPS).toFixed(2)} us/step)`);
128}
130main().then(() => process.exit(0)).catch((e) => { console.error(e); process.exit(1); });