1// Microbenchmark: what do workgroup/storage barriers and the stage loops cost
2// on this device, per simulated step, for one workgroup?
3import { installNodeWebGpu, requestDevice } from './gpu.mjs';
5const N = 1560;
6const WS = Number(process.env.WS ?? 256);
7const STEPS = 2048;
9function kernel(body, nbar) {
10 return `
11const N = ${N}u; const WS = ${WS}u; const STEPS = ${STEPS}u;
12@group(0) @binding(0) var<storage, read_write> st: array<f32>;
13@compute @workgroup_size(${WS})
14fn main(@builtin(local_invocation_index) lid: u32) {
15 for (var s: u32 = 0u; s < STEPS; s = s + 1u) {
16 ${Array.from({ length: nbar }, () => body).join('\n')}
17 }
18 if (lid == 0u) { st[0] = f32(STEPS); }
19}
20`;
21}
23async function time(device, code, label) {
24 const module = device.createShaderModule({ code });
25 const info = await module.getCompilationInfo();
26 const errs = info.messages.filter((m) => m.type === 'error');
27 if (errs.length) throw new Error(errs.map((m) => m.message).join('\n'));
28 const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
29 const buf = device.createBuffer({ size: 4 * N * 20, usage: GPUBufferUsage.STORAGE });
30 const bg = device.createBindGroup({
31 layout: pipeline.getBindGroupLayout(0),
32 entries: [{ binding: 0, resource: { buffer: buf } }],
33 });
34 // warmup
35 let enc = device.createCommandEncoder();
36 let pass = enc.beginComputePass(); pass.setPipeline(pipeline); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(1); pass.end();
37 device.queue.submit([enc.finish()]);
38 await device.queue.onSubmittedWorkDone();
39 const t0 = performance.now();
40 const reps = 4;
41 enc = device.createCommandEncoder();
42 for (let r = 0; r < reps; r++) {
43 pass = enc.beginComputePass(); pass.setPipeline(pipeline); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(1); pass.end();
44 }
45 device.queue.submit([enc.finish()]);
46 await device.queue.onSubmittedWorkDone();
47 const ms = (performance.now() - t0) / reps;
48 console.log(`${label}: ${(1000 * ms / STEPS).toFixed(2)} us/step`);
49 buf.destroy();
50}
52async function main() {
53 await installNodeWebGpu();
54 const { device, info } = await requestDevice();
55 console.log(`device: ${info}, WS=${WS}`);
57 await time(device, kernel('', 0), 'empty loop');
58 await time(device, kernel('workgroupBarrier();', 14), '14 workgroupBarrier');
59 await time(device, kernel('storageBarrier();', 14), '14 storageBarrier');
60 await time(device, kernel(`
61 for (var i: u32 = lid; i < N; i = i + WS) {
62 st[N + i] = st[i] * 1.0001 + 0.5;
63 }
64 storageBarrier();`, 10), '10 x (elementwise N + storageBarrier)');
65 await time(device, kernel(`
66 for (var i: u32 = lid; i < N; i = i + WS) {
67 var sum: f32 = 0.0;
68 let l = i % 26u; let j = i / 26u;
69 for (var q: u32 = 0u; q < 60u; q = q + 1u) {
70 sum = sum + st[l + q * 26u] * st[2u * N + q + j * 60u];
71 }
72 st[N + i] = sum;
73 }
74 storageBarrier();`, 4), '4 x (60-MAC matmul N + storageBarrier)');
75}
77main().then(() => process.exit(0)).catch((e) => { console.error(e); process.exit(1); });