// Fully unrolled register-accumulator variants: no local arrays anywhere,
// so nothing can spill to scratch memory; accumulators are named scalars.
import { installNodeWebGpu, requestDevice } from './gpu.mjs';

const NRT = 26, NTH = 60, N = NRT * NTH;
const WS = 256;
const STEPS = 2048;

function wrap(body) {
  return `
const NRT = ${NRT}u; const NTH = ${NTH}u; const N = ${N}u; const WS = ${WS}u; const STEPS = ${STEPS}u;
@group(0) @binding(0) var<storage, read> ops: array<f32>;
@group(0) @binding(2) var<storage, read_write> st: array<f32>;
@compute @workgroup_size(${WS})
fn main(@builtin(local_invocation_index) lid: u32) {
  for (var s: u32 = 0u; s < STEPS; s = s + 1u) {
${body}
    workgroupBarrier();
    storageBarrier();
  }
  if (lid == 0u) { st[0] = st[1] + 1.0; }
}
`;
}

// transform: 60 threads, 26 scalar accumulators, unrolled inner block
function transformUnrolled() {
  const acc = Array.from({ length: NRT }, (_, l) => `var a${l}: f32 = 0.0;`).join(' ');
  const fmas = Array.from({ length: NRT }, (_, l) => `a${l} = a${l} + st[xb + ${l}u] * mq;`).join('\n        ');
  const stores = Array.from({ length: NRT }, (_, l) => `st[ob + ${l}u] = a${l};`).join(' ');
  return wrap(`
    if (lid < NTH) {
      let j = lid;
      ${acc}
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        let mq = ops[q + j * NTH];
        let xb = N + q * NRT;
        ${fmas}
      }
      let ob = 2u * N + j * NRT;
      ${stores}
    }`);
}

// transform with two columns per thread (120 threads), 26 accs each... too many regs;
// instead: 13 accumulators x 120 threads (j, half), unrolled
function transformUnrolledHalf() {
  const acc = Array.from({ length: 13 }, (_, l) => `var a${l}: f32 = 0.0;`).join(' ');
  const fmas = Array.from({ length: 13 }, (_, l) => `a${l} = a${l} + st[xb + ${l}u] * mq;`).join('\n        ');
  const stores = Array.from({ length: 13 }, (_, l) => `st[ob + ${l}u] = a${l};`).join(' ');
  return wrap(`
    if (lid < 120u) {
      let j = lid % NTH; let h = lid / NTH;
      let l0 = h * 13u;
      ${acc}
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        let mq = ops[q + j * NTH];
        let xb = N + q * NRT + l0;
        ${fmas}
      }
      let ob = 2u * N + j * NRT + l0;
      ${stores}
    }`);
}

// page solve unrolled the same way: 60 threads, 26 accs, A read coalesced
function pageUnrolled() {
  const acc = Array.from({ length: NRT }, (_, l) => `var a${l}: f32 = 0.0;`).join(' ');
  const fmas = Array.from({ length: NRT }, (_, l) => `a${l} = a${l} + ops[ab + ${l}u] * xq;`).join('\n        ');
  const stores = Array.from({ length: NRT }, (_, l) => `st[ob + ${l}u] = a${l};`).join(' ');
  return wrap(`
    if (lid < NTH) {
      let j = lid;
      let base = j * NRT * NRT;
      ${acc}
      for (var p: u32 = 0u; p < NRT; p = p + 1u) {
        let xq = st[N + p + j * NRT];
        let ab = base + p * NRT;
        ${fmas}
      }
      let ob = 2u * N + j * NRT;
      ${stores}
    }`);
}

const CASES = {
  transformUnrolled: transformUnrolled(),
  transformUnrolledHalf: transformUnrolledHalf(),
  pageUnrolled: pageUnrolled(),
};

async function time(device, code, label) {
  const module = device.createShaderModule({ code });
  const info = await module.getCompilationInfo();
  const errs = info.messages.filter((m) => m.type === 'error');
  if (errs.length) throw new Error(label + ':\n' + errs.map((m) => `${m.lineNum}: ${m.message}`).join('\n'));
  const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
  const ops = device.createBuffer({ size: 4 * NRT * NRT * NTH, usage: GPUBufferUsage.STORAGE });
  const stb = device.createBuffer({ size: 4 * N * 4, usage: GPUBufferUsage.STORAGE });
  const bg = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [
      { binding: 0, resource: { buffer: ops } },
      { binding: 2, resource: { buffer: stb } },
    ],
  });
  const run = async (reps) => {
    const enc = device.createCommandEncoder();
    for (let r = 0; r < reps; r++) {
      const pass = enc.beginComputePass();
      pass.setPipeline(pipeline); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(1); pass.end();
    }
    device.queue.submit([enc.finish()]);
    await device.queue.onSubmittedWorkDone();
  };
  await run(1);
  const t0 = performance.now();
  await run(4);
  const ms = (performance.now() - t0) / 4;
  console.log(`${label}: ${(1000 * ms / STEPS).toFixed(2)} us per stage-step`);
  ops.destroy(); stb.destroy();
}

async function main() {
  await installNodeWebGpu();
  const { device, info } = await requestDevice();
  console.log(`device: ${info}`);
  for (const [label, code] of Object.entries(CASES)) await time(device, code, label);
}

main().then(() => process.exit(0)).catch((e) => { console.error(e); process.exit(1); });