10729f3WebGPU time loop for the cymbal2 shell model, with numbl setup and output stagesJeremy Magland 1// Which matmul formulation is fast on this device, for X (26x60) * M (60x60),
2// one workgroup? Candidates: per-output serial reduction (baseline), shared-
3// memory staged input, column-per-thread with register accumulators, and the
4// combination. Also the 26x26 page solve (pagemtimes) both ways.
5import { installNodeWebGpu, requestDevice } from './gpu.mjs';
7const NRT = 26, NTH = 60, N = NRT * NTH;
8const WS = Number(process.env.WS ?? 256);
9const STEPS = 2048;
11function wrap(body, decl = '') {
12 return `
13const NRT = ${NRT}u; const NTH = ${NTH}u; const N = ${N}u; const WS = ${WS}u; const STEPS = ${STEPS}u;
14@group(0) @binding(0) var<storage, read> ops: array<f32>;
15@group(0) @binding(1) var<storage, read_write> st: array<f32>;
16${decl}
17@compute @workgroup_size(${WS})
18fn main(@builtin(local_invocation_index) lid: u32) {
19 for (var s: u32 = 0u; s < STEPS; s = s + 1u) {
20${body}
21 workgroupBarrier();
22 storageBarrier();
23 }
24 if (lid == 0u) { st[0] = st[1] + 1.0; }
25}
26`;
27}
29const CASES = {
30 baseline: wrap(`
31 for (var i: u32 = lid; i < N; i = i + WS) {
32 let l = i % NRT; let j = i / NRT;
33 var sum: f32 = 0.0;
34 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
35 sum = sum + st[N + l + q * NRT] * ops[q + j * NTH];
36 }
37 st[2u * N + i] = sum;
38 }`),
40 sharedX: wrap(`
41 for (var i: u32 = lid; i < N; i = i + WS) { xsh[i] = st[N + i]; }
42 workgroupBarrier();
43 for (var i: u32 = lid; i < N; i = i + WS) {
44 let l = i % NRT; let j = i / NRT;
45 var sum: f32 = 0.0;
46 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
47 sum = sum + xsh[l + q * NRT] * ops[q + j * NTH];
48 }
49 st[2u * N + i] = sum;
50 }`, 'var<workgroup> xsh: array<f32, 1560>;'),
52 colRegs: wrap(`
53 if (lid < NTH) {
54 let j = lid;
55 var acc: array<f32, ${NRT}>;
56 for (var l: u32 = 0u; l < NRT; l = l + 1u) { acc[l] = 0.0; }
57 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
58 let mq = ops[q + j * NTH];
59 for (var l: u32 = 0u; l < NRT; l = l + 1u) {
60 acc[l] = acc[l] + st[N + l + q * NRT] * mq;
61 }
62 }
63 for (var l: u32 = 0u; l < NRT; l = l + 1u) { st[2u * N + l + j * NRT] = acc[l]; }
64 }`),
66 colRegsShared: wrap(`
67 for (var i: u32 = lid; i < N; i = i + WS) { xsh[i] = st[N + i]; }
68 workgroupBarrier();
69 if (lid < NTH) {
70 let j = lid;
71 var acc: array<f32, ${NRT}>;
72 for (var l: u32 = 0u; l < NRT; l = l + 1u) { acc[l] = 0.0; }
73 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
74 let mq = ops[q + j * NTH];
75 for (var l: u32 = 0u; l < NRT; l = l + 1u) {
76 acc[l] = acc[l] + xsh[l + q * NRT] * mq;
77 }
78 }
79 for (var l: u32 = 0u; l < NRT; l = l + 1u) { st[2u * N + l + j * NRT] = acc[l]; }
80 }`, 'var<workgroup> xsh: array<f32, 1560>;'),
82 colBlockShared: wrap(`
83 for (var i: u32 = lid; i < N; i = i + WS) { xsh[i] = st[N + i]; }
84 workgroupBarrier();
85 if (lid < 240u) {
86 let j = lid % NTH; let b = lid / NTH; // 4 blocks of 7 rows (last short)
87 let l0 = b * 7u;
88 let nl = min(7u, NRT - l0);
89 var acc: array<f32, 7>;
90 for (var l: u32 = 0u; l < 7u; l = l + 1u) { acc[l] = 0.0; }
91 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
92 let mq = ops[q + j * NTH];
93 for (var l: u32 = 0u; l < nl; l = l + 1u) {
94 acc[l] = acc[l] + xsh[l0 + l + q * NRT] * mq;
95 }
96 }
97 for (var l: u32 = 0u; l < nl; l = l + 1u) { st[2u * N + l0 + l + j * NRT] = acc[l]; }
98 }`, 'var<workgroup> xsh: array<f32, 1560>;'),
100 pageBaseline: wrap(`
101 for (var i: u32 = lid; i < N; i = i + WS) {
102 let l = i % NRT; let j = i / NRT;
103 let base = j * NRT * NRT;
104 var sum: f32 = 0.0;
105 for (var p: u32 = 0u; p < NRT; p = p + 1u) {
106 sum = sum + ops[base + l + p * NRT] * st[N + p + j * NRT];
107 }
108 st[2u * N + i] = sum;
109 }`),
111 pageColRegsShared: wrap(`
112 for (var i: u32 = lid; i < N; i = i + WS) { xsh[i] = st[N + i]; }
113 workgroupBarrier();
114 if (lid < 240u) {
115 let j = lid % NTH; let b = lid / NTH;
116 let l0 = b * 7u;
117 let nl = min(7u, NRT - l0);
118 let base = j * NRT * NRT;
119 var acc: array<f32, 7>;
120 for (var l: u32 = 0u; l < 7u; l = l + 1u) { acc[l] = 0.0; }
121 for (var p: u32 = 0u; p < NRT; p = p + 1u) {
122 let xq = xsh[p + j * NRT];
123 for (var l: u32 = 0u; l < nl; l = l + 1u) {
124 acc[l] = acc[l] + ops[base + l0 + l + p * NRT] * xq;
125 }
126 }
127 for (var l: u32 = 0u; l < nl; l = l + 1u) { st[2u * N + l0 + l + j * NRT] = acc[l]; }
128 }`, 'var<workgroup> xsh: array<f32, 1560>;'),
129};
131async function time(device, code, label) {
132 const module = device.createShaderModule({ code });
133 const info = await module.getCompilationInfo();
134 const errs = info.messages.filter((m) => m.type === 'error');
135 if (errs.length) throw new Error(label + ':\n' + errs.map((m) => `${m.lineNum}: ${m.message}`).join('\n'));
136 const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
137 const ops = device.createBuffer({ size: 4 * NRT * NRT * NTH, usage: GPUBufferUsage.STORAGE });
138 const stb = device.createBuffer({ size: 4 * N * 4, usage: GPUBufferUsage.STORAGE });
139 const bg = device.createBindGroup({
140 layout: pipeline.getBindGroupLayout(0),
141 entries: [
142 { binding: 0, resource: { buffer: ops } },
143 { binding: 1, resource: { buffer: stb } },
144 ],
145 });
146 const run = async (reps) => {
147 const enc = device.createCommandEncoder();
148 for (let r = 0; r < reps; r++) {
149 const pass = enc.beginComputePass();
150 pass.setPipeline(pipeline); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(1); pass.end();
151 }
152 device.queue.submit([enc.finish()]);
153 await device.queue.onSubmittedWorkDone();
154 };
155 await run(1);
156 const t0 = performance.now();
157 await run(4);
158 const ms = (performance.now() - t0) / 4;
159 console.log(`${label}: ${(1000 * ms / STEPS).toFixed(2)} us per stage-step`);
160 ops.destroy(); stb.destroy();
161}
163async function main() {
164 await installNodeWebGpu();
165 const { device, info } = await requestDevice();
166 console.log(`device: ${info}, WS=${WS}`);
167 for (const [label, code] of Object.entries(CASES)) await time(device, code, label);
168}
170if (!process.env.MORE) { main().then(() => process.exit(0)).catch((e) => { console.error(e); process.exit(1); }); }
172// --- appended variants ---
173export const MORE = {
174 transposedBaseline: wrap(`
175 for (var i: u32 = lid; i < N; i = i + WS) {
176 let l = i % NRT; let j = i / NRT;
177 var sum: f32 = 0.0;
178 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
179 sum = sum + st[N + q + l * NTH] * ops[q + j * NTH];
180 }
181 st[2u * N + i] = sum;
182 }`),
183 colRegs13x120: wrap(`
184 if (lid < 120u) {
185 let j = lid % NTH; let b = lid / NTH;
186 let l0 = b * 13u;
187 var acc: array<f32, 13>;
188 for (var l: u32 = 0u; l < 13u; l = l + 1u) { acc[l] = 0.0; }
189 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
190 let mq = ops[q + j * NTH];
191 for (var l: u32 = 0u; l < 13u; l = l + 1u) {
192 acc[l] = acc[l] + st[N + l0 + l + q * NRT] * mq;
193 }
194 }
195 for (var l: u32 = 0u; l < 13u; l = l + 1u) { st[2u * N + l0 + l + j * NRT] = acc[l]; }
196 }`),
197 page3colRegs: wrap(`
198 if (lid < 120u) {
199 let j = lid % NTH; let b = lid / NTH;
200 let l0 = b * 13u;
201 let base = j * NRT * NRT;
202 var acc: array<f32, 13>;
203 for (var l: u32 = 0u; l < 13u; l = l + 1u) { acc[l] = 0.0; }
204 for (var p: u32 = 0u; p < NRT; p = p + 1u) {
205 let x1 = st[N + p + j * NRT];
206 let x2 = st[2u * N + p + j * NRT];
207 let x3 = st[3u * N + p + j * NRT];
208 for (var l: u32 = 0u; l < 13u; l = l + 1u) {
209 acc[l] = acc[l] + ops[base + l0 + l + p * NRT] * x1
210 + ops[base + l0 + l + p * NRT + 676u] * x2
211 + ops[base + l0 + l + p * NRT + 1352u] * x3;
212 }
213 }
214 for (var l: u32 = 0u; l < 13u; l = l + 1u) { st[l0 + l + j * NRT] = acc[l]; }
215 }`),
216};
218async function main2() {
219 await installNodeWebGpu();
220 const { device, info } = await requestDevice();
221 console.log(`device: ${info}, WS=${WS}`);
222 for (const [label, code] of Object.entries(MORE)) await time(device, code, label);
223 // dispatch overhead: many tiny dispatches in one submit
224 {
225 const code = `
226@group(0) @binding(0) var<storage, read_write> st: array<f32>;
227@compute @workgroup_size(64)
228fn main(@builtin(local_invocation_index) lid: u32) { if (lid == 0u) { st[0] = st[0] + 1.0; } }
229`;
230 const module = device.createShaderModule({ code });
231 const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
232 const buf = device.createBuffer({ size: 16, usage: GPUBufferUsage.STORAGE });
233 const bg = device.createBindGroup({ layout: pipeline.getBindGroupLayout(0), entries: [{ binding: 0, resource: { buffer: buf } }] });
234 const enc = device.createCommandEncoder();
235 const pass = enc.beginComputePass();
236 pass.setPipeline(pipeline); pass.setBindGroup(0, bg);
237 for (let i = 0; i < 8192; i++) pass.dispatchWorkgroups(1);
238 pass.end();
239 const cb = enc.finish();
240 const t0 = performance.now();
241 device.queue.submit([cb]);
242 await device.queue.onSubmittedWorkDone();
243 console.log(`dispatch overhead: ${((performance.now() - t0) * 1000 / 8192).toFixed(2)} us/dispatch`);
244 }
245}
246if (process.env.MORE) { main2().then(() => process.exit(0)).catch((e) => { console.error(e); process.exit(1); }); }