10729f3WebGPU time loop for the cymbal2 shell model, with numbl setup and output stagesJeremy Magland 1// vec4 variants with the radial dimension padded to 28 (7 vec4 lanes).
2// Transform: out[l,j] = sum_q X[l,q] M[q,j]; X stored 28 x 60, M 60 x 60.
3// Page: out[l,j] = sum_p A[l,p,j] X[p,j]; A pages 28 x 28.
4import { installNodeWebGpu, requestDevice } from './gpu.mjs';
6const PADL = 28, NTH = 60, NP = PADL * NTH; // 1680
7const NPV = NP / 4;
8const WS = 256;
9const STEPS = 2048;
11function wrap(body) {
12 return `
13const PADL = ${PADL}u; const NTH = ${NTH}u; const NP = ${NP}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> opsv: array<vec4<f32>>;
16@group(0) @binding(2) var<storage, read_write> stv: array<vec4<f32>>;
17fn ld(i: u32) -> f32 { return stv[i / 4u][i % 4u]; }
18@compute @workgroup_size(${WS})
19fn main(@builtin(local_invocation_index) lid: u32) {
20 for (var s: u32 = 0u; s < STEPS; s = s + 1u) {
21${body}
22 workgroupBarrier();
23 storageBarrier();
24 }
25 if (lid == 0u) { stv[0].x = stv[0].y + 1.0; }
26}
27`;
28}
30const CASES = {
31 // 60 threads, 7 vec4 accumulators, X read as vec4, M scalar broadcast
32 colRegsVec4: wrap(`
33 if (lid < NTH) {
34 let j = lid;
35 var acc: array<vec4<f32>, 7>;
36 for (var v: u32 = 0u; v < 7u; v = v + 1u) { acc[v] = vec4<f32>(0.0); }
37 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
38 let mq = ops[q + j * NTH];
39 let xb = (NP + q * PADL) / 4u;
40 for (var v: u32 = 0u; v < 7u; v = v + 1u) {
41 acc[v] = acc[v] + stv[xb + v] * mq;
42 }
43 }
44 let ob = (2u * NP + j * PADL) / 4u;
45 for (var v: u32 = 0u; v < 7u; v = v + 1u) { stv[ob + v] = acc[v]; }
46 }`),
48 // 120 threads: j column plus upper/lower half (4 vec4s h=0, 3 vec4s h=1)
49 colRegsVec4x120: wrap(`
50 if (lid < 120u) {
51 let j = lid % NTH; let h = lid / NTH;
52 let v0 = h * 4u;
53 let nv = 4u - h; // 4 or 3
54 var acc: array<vec4<f32>, 4>;
55 for (var v: u32 = 0u; v < 4u; v = v + 1u) { acc[v] = vec4<f32>(0.0); }
56 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
57 let mq = ops[q + j * NTH];
58 let xb = (NP + q * PADL) / 4u + v0;
59 for (var v: u32 = 0u; v < 4u; v = v + 1u) {
60 if (v < nv) { acc[v] = acc[v] + stv[xb + v] * mq; }
61 }
62 }
63 let ob = (2u * NP + j * PADL) / 4u + v0;
64 for (var v: u32 = 0u; v < 4u; v = v + 1u) {
65 if (v < nv) { stv[ob + v] = acc[v]; }
66 }
67 }`),
69 // page solve: 60 threads, A columns read as vec4 (28-padded), X broadcast
70 pageVec4: wrap(`
71 if (lid < NTH) {
72 let j = lid;
73 let base = j * PADL * PADL;
74 var acc: array<vec4<f32>, 7>;
75 for (var v: u32 = 0u; v < 7u; v = v + 1u) { acc[v] = vec4<f32>(0.0); }
76 for (var p: u32 = 0u; p < PADL; p = p + 1u) {
77 let xq = ld(NP + p + j * PADL);
78 let ab = (base + p * PADL) / 4u;
79 for (var v: u32 = 0u; v < 7u; v = v + 1u) {
80 acc[v] = acc[v] + opsv[ab + v] * xq;
81 }
82 }
83 let ob = (2u * NP + j * PADL) / 4u;
84 for (var v: u32 = 0u; v < 7u; v = v + 1u) { stv[ob + v] = acc[v]; }
85 }`),
87 // triple page solve (the shell update): same but three A blocks, three X
88 page3Vec4: wrap(`
89 if (lid < NTH) {
90 let j = lid;
91 let base = j * PADL * PADL;
92 var acc: array<vec4<f32>, 7>;
93 for (var v: u32 = 0u; v < 7u; v = v + 1u) { acc[v] = vec4<f32>(0.0); }
94 for (var p: u32 = 0u; p < PADL; p = p + 1u) {
95 let x1 = ld(NP + p + j * PADL);
96 let x2 = ld(2u * NP + p + j * PADL);
97 let x3 = ld(3u * NP + p + j * PADL);
98 let ab = (base + p * PADL) / 4u;
99 let s2 = ${PADL * PADL * NTH / 4}u;
100 for (var v: u32 = 0u; v < 7u; v = v + 1u) {
101 acc[v] = acc[v] + opsv[ab + v] * x1 + opsv[ab + v + s2] * x2 + opsv[ab + v + 2u * s2] * x3;
102 }
103 }
104 let ob = (j * PADL) / 4u;
105 for (var v: u32 = 0u; v < 7u; v = v + 1u) { stv[ob + v] = acc[v]; }
106 }`),
107};
109async function time(device, code, label) {
110 const module = device.createShaderModule({ code });
111 const info = await module.getCompilationInfo();
112 const errs = info.messages.filter((m) => m.type === 'error');
113 if (errs.length) throw new Error(label + ':\n' + errs.map((m) => `${m.lineNum}: ${m.message}`).join('\n'));
114 const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
115 const ops = device.createBuffer({ size: 4 * PADL * PADL * NTH * 3, usage: GPUBufferUsage.STORAGE });
116 const stb = device.createBuffer({ size: 4 * NP * 4, usage: GPUBufferUsage.STORAGE });
117 const entries = [];
118 if (/\bops\[/.test(code)) entries.push({ binding: 0, resource: { buffer: ops } });
119 if (/\bopsv\[/.test(code)) entries.push({ binding: 1, resource: { buffer: ops } });
120 entries.push({ binding: 2, resource: { buffer: stb } });
121 const bg = device.createBindGroup({
122 layout: pipeline.getBindGroupLayout(0),
123 entries,
124 });
125 const run = async (reps) => {
126 const enc = device.createCommandEncoder();
127 for (let r = 0; r < reps; r++) {
128 const pass = enc.beginComputePass();
129 pass.setPipeline(pipeline); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(1); pass.end();
130 }
131 device.queue.submit([enc.finish()]);
132 await device.queue.onSubmittedWorkDone();
133 };
134 await run(1);
135 const t0 = performance.now();
136 await run(4);
137 const ms = (performance.now() - t0) / 4;
138 console.log(`${label}: ${(1000 * ms / STEPS).toFixed(2)} us per stage-step`);
139 ops.destroy(); stb.destroy();
140}
142async function main() {
143 await installNodeWebGpu();
144 const { device, info } = await requestDevice();
145 console.log(`device: ${info}`);
146 for (const [label, code] of Object.entries(CASES)) await time(device, code, label);
147}
149main().then(() => process.exit(0)).catch((e) => { console.error(e); process.exit(1); });