1// Fully unrolled register-accumulator variants: no local arrays anywhere,
2// so nothing can spill to scratch memory; accumulators are named scalars.
3import { installNodeWebGpu, requestDevice } from './gpu.mjs';
5const NRT = 26, NTH = 60, N = NRT * NTH;
6const WS = 256;
7const STEPS = 2048;
9function wrap(body) {
10 return `
11const NRT = ${NRT}u; const NTH = ${NTH}u; const N = ${N}u; const WS = ${WS}u; const STEPS = ${STEPS}u;
12@group(0) @binding(0) var<storage, read> ops: array<f32>;
13@group(0) @binding(2) var<storage, read_write> st: array<f32>;
14@compute @workgroup_size(${WS})
15fn main(@builtin(local_invocation_index) lid: u32) {
16 for (var s: u32 = 0u; s < STEPS; s = s + 1u) {
17${body}
18 workgroupBarrier();
19 storageBarrier();
20 }
21 if (lid == 0u) { st[0] = st[1] + 1.0; }
22}
23`;
24}
26// transform: 60 threads, 26 scalar accumulators, unrolled inner block
27function transformUnrolled() {
28 const acc = Array.from({ length: NRT }, (_, l) => `var a${l}: f32 = 0.0;`).join(' ');
29 const fmas = Array.from({ length: NRT }, (_, l) => `a${l} = a${l} + st[xb + ${l}u] * mq;`).join('\n ');
30 const stores = Array.from({ length: NRT }, (_, l) => `st[ob + ${l}u] = a${l};`).join(' ');
31 return wrap(`
32 if (lid < NTH) {
33 let j = lid;
34 ${acc}
35 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
36 let mq = ops[q + j * NTH];
37 let xb = N + q * NRT;
38 ${fmas}
39 }
40 let ob = 2u * N + j * NRT;
41 ${stores}
42 }`);
43}
45// transform with two columns per thread (120 threads), 26 accs each... too many regs;
46// instead: 13 accumulators x 120 threads (j, half), unrolled
47function transformUnrolledHalf() {
48 const acc = Array.from({ length: 13 }, (_, l) => `var a${l}: f32 = 0.0;`).join(' ');
49 const fmas = Array.from({ length: 13 }, (_, l) => `a${l} = a${l} + st[xb + ${l}u] * mq;`).join('\n ');
50 const stores = Array.from({ length: 13 }, (_, l) => `st[ob + ${l}u] = a${l};`).join(' ');
51 return wrap(`
52 if (lid < 120u) {
53 let j = lid % NTH; let h = lid / NTH;
54 let l0 = h * 13u;
55 ${acc}
56 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
57 let mq = ops[q + j * NTH];
58 let xb = N + q * NRT + l0;
59 ${fmas}
60 }
61 let ob = 2u * N + j * NRT + l0;
62 ${stores}
63 }`);
64}
66// page solve unrolled the same way: 60 threads, 26 accs, A read coalesced
67function pageUnrolled() {
68 const acc = Array.from({ length: NRT }, (_, l) => `var a${l}: f32 = 0.0;`).join(' ');
69 const fmas = Array.from({ length: NRT }, (_, l) => `a${l} = a${l} + ops[ab + ${l}u] * xq;`).join('\n ');
70 const stores = Array.from({ length: NRT }, (_, l) => `st[ob + ${l}u] = a${l};`).join(' ');
71 return wrap(`
72 if (lid < NTH) {
73 let j = lid;
74 let base = j * NRT * NRT;
75 ${acc}
76 for (var p: u32 = 0u; p < NRT; p = p + 1u) {
77 let xq = st[N + p + j * NRT];
78 let ab = base + p * NRT;
79 ${fmas}
80 }
81 let ob = 2u * N + j * NRT;
82 ${stores}
83 }`);
84}
86const CASES = {
87 transformUnrolled: transformUnrolled(),
88 transformUnrolledHalf: transformUnrolledHalf(),
89 pageUnrolled: pageUnrolled(),
90};
92async function time(device, code, label) {
93 const module = device.createShaderModule({ code });
94 const info = await module.getCompilationInfo();
95 const errs = info.messages.filter((m) => m.type === 'error');
96 if (errs.length) throw new Error(label + ':\n' + errs.map((m) => `${m.lineNum}: ${m.message}`).join('\n'));
97 const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
98 const ops = device.createBuffer({ size: 4 * NRT * NRT * NTH, usage: GPUBufferUsage.STORAGE });
99 const stb = device.createBuffer({ size: 4 * N * 4, usage: GPUBufferUsage.STORAGE });
100 const bg = device.createBindGroup({
101 layout: pipeline.getBindGroupLayout(0),
102 entries: [
103 { binding: 0, resource: { buffer: ops } },
104 { binding: 2, resource: { buffer: stb } },
105 ],
106 });
107 const run = async (reps) => {
108 const enc = device.createCommandEncoder();
109 for (let r = 0; r < reps; r++) {
110 const pass = enc.beginComputePass();
111 pass.setPipeline(pipeline); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(1); pass.end();
112 }
113 device.queue.submit([enc.finish()]);
114 await device.queue.onSubmittedWorkDone();
115 };
116 await run(1);
117 const t0 = performance.now();
118 await run(4);
119 const ms = (performance.now() - t0) / 4;
120 console.log(`${label}: ${(1000 * ms / STEPS).toFixed(2)} us per stage-step`);
121 ops.destroy(); stb.destroy();
122}
124async function main() {
125 await installNodeWebGpu();
126 const { device, info } = await requestDevice();
127 console.log(`device: ${info}`);
128 for (const [label, code] of Object.entries(CASES)) await time(device, code, label);
129}
131main().then(() => process.exit(0)).catch((e) => { console.error(e); process.exit(1); });