// Which matmul formulation is fast on this device, for X (26x60) * M (60x60),
// one workgroup? Candidates: per-output serial reduction (baseline), shared-
// memory staged input, column-per-thread with register accumulators, and the
// combination. Also the 26x26 page solve (pagemtimes) both ways.
import { installNodeWebGpu, requestDevice } from './gpu.mjs';

const NRT = 26, NTH = 60, N = NRT * NTH;
const WS = Number(process.env.WS ?? 256);
const STEPS = 2048;

function wrap(body, decl = '') {
  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(1) var<storage, read_write> st: array<f32>;
${decl}
@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; }
}
`;
}

const CASES = {
  baseline: wrap(`
    for (var i: u32 = lid; i < N; i = i + WS) {
      let l = i % NRT; let j = i / NRT;
      var sum: f32 = 0.0;
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        sum = sum + st[N + l + q * NRT] * ops[q + j * NTH];
      }
      st[2u * N + i] = sum;
    }`),

  sharedX: wrap(`
    for (var i: u32 = lid; i < N; i = i + WS) { xsh[i] = st[N + i]; }
    workgroupBarrier();
    for (var i: u32 = lid; i < N; i = i + WS) {
      let l = i % NRT; let j = i / NRT;
      var sum: f32 = 0.0;
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        sum = sum + xsh[l + q * NRT] * ops[q + j * NTH];
      }
      st[2u * N + i] = sum;
    }`, 'var<workgroup> xsh: array<f32, 1560>;'),

  colRegs: wrap(`
    if (lid < NTH) {
      let j = lid;
      var acc: array<f32, ${NRT}>;
      for (var l: u32 = 0u; l < NRT; l = l + 1u) { acc[l] = 0.0; }
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        let mq = ops[q + j * NTH];
        for (var l: u32 = 0u; l < NRT; l = l + 1u) {
          acc[l] = acc[l] + st[N + l + q * NRT] * mq;
        }
      }
      for (var l: u32 = 0u; l < NRT; l = l + 1u) { st[2u * N + l + j * NRT] = acc[l]; }
    }`),

  colRegsShared: wrap(`
    for (var i: u32 = lid; i < N; i = i + WS) { xsh[i] = st[N + i]; }
    workgroupBarrier();
    if (lid < NTH) {
      let j = lid;
      var acc: array<f32, ${NRT}>;
      for (var l: u32 = 0u; l < NRT; l = l + 1u) { acc[l] = 0.0; }
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        let mq = ops[q + j * NTH];
        for (var l: u32 = 0u; l < NRT; l = l + 1u) {
          acc[l] = acc[l] + xsh[l + q * NRT] * mq;
        }
      }
      for (var l: u32 = 0u; l < NRT; l = l + 1u) { st[2u * N + l + j * NRT] = acc[l]; }
    }`, 'var<workgroup> xsh: array<f32, 1560>;'),

  colBlockShared: wrap(`
    for (var i: u32 = lid; i < N; i = i + WS) { xsh[i] = st[N + i]; }
    workgroupBarrier();
    if (lid < 240u) {
      let j = lid % NTH; let b = lid / NTH;   // 4 blocks of 7 rows (last short)
      let l0 = b * 7u;
      let nl = min(7u, NRT - l0);
      var acc: array<f32, 7>;
      for (var l: u32 = 0u; l < 7u; l = l + 1u) { acc[l] = 0.0; }
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        let mq = ops[q + j * NTH];
        for (var l: u32 = 0u; l < nl; l = l + 1u) {
          acc[l] = acc[l] + xsh[l0 + l + q * NRT] * mq;
        }
      }
      for (var l: u32 = 0u; l < nl; l = l + 1u) { st[2u * N + l0 + l + j * NRT] = acc[l]; }
    }`, 'var<workgroup> xsh: array<f32, 1560>;'),

  pageBaseline: wrap(`
    for (var i: u32 = lid; i < N; i = i + WS) {
      let l = i % NRT; let j = i / NRT;
      let base = j * NRT * NRT;
      var sum: f32 = 0.0;
      for (var p: u32 = 0u; p < NRT; p = p + 1u) {
        sum = sum + ops[base + l + p * NRT] * st[N + p + j * NRT];
      }
      st[2u * N + i] = sum;
    }`),

  pageColRegsShared: wrap(`
    for (var i: u32 = lid; i < N; i = i + WS) { xsh[i] = st[N + i]; }
    workgroupBarrier();
    if (lid < 240u) {
      let j = lid % NTH; let b = lid / NTH;
      let l0 = b * 7u;
      let nl = min(7u, NRT - l0);
      let base = j * NRT * NRT;
      var acc: array<f32, 7>;
      for (var l: u32 = 0u; l < 7u; l = l + 1u) { acc[l] = 0.0; }
      for (var p: u32 = 0u; p < NRT; p = p + 1u) {
        let xq = xsh[p + j * NRT];
        for (var l: u32 = 0u; l < nl; l = l + 1u) {
          acc[l] = acc[l] + ops[base + l0 + l + p * NRT] * xq;
        }
      }
      for (var l: u32 = 0u; l < nl; l = l + 1u) { st[2u * N + l0 + l + j * NRT] = acc[l]; }
    }`, 'var<workgroup> xsh: array<f32, 1560>;'),
};

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: 1, 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}, WS=${WS}`);
  for (const [label, code] of Object.entries(CASES)) await time(device, code, label);
}

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

// --- appended variants ---
export const MORE = {
  transposedBaseline: wrap(`
    for (var i: u32 = lid; i < N; i = i + WS) {
      let l = i % NRT; let j = i / NRT;
      var sum: f32 = 0.0;
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        sum = sum + st[N + q + l * NTH] * ops[q + j * NTH];
      }
      st[2u * N + i] = sum;
    }`),
  colRegs13x120: wrap(`
    if (lid < 120u) {
      let j = lid % NTH; let b = lid / NTH;
      let l0 = b * 13u;
      var acc: array<f32, 13>;
      for (var l: u32 = 0u; l < 13u; l = l + 1u) { acc[l] = 0.0; }
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        let mq = ops[q + j * NTH];
        for (var l: u32 = 0u; l < 13u; l = l + 1u) {
          acc[l] = acc[l] + st[N + l0 + l + q * NRT] * mq;
        }
      }
      for (var l: u32 = 0u; l < 13u; l = l + 1u) { st[2u * N + l0 + l + j * NRT] = acc[l]; }
    }`),
  page3colRegs: wrap(`
    if (lid < 120u) {
      let j = lid % NTH; let b = lid / NTH;
      let l0 = b * 13u;
      let base = j * NRT * NRT;
      var acc: array<f32, 13>;
      for (var l: u32 = 0u; l < 13u; l = l + 1u) { acc[l] = 0.0; }
      for (var p: u32 = 0u; p < NRT; p = p + 1u) {
        let x1 = st[N + p + j * NRT];
        let x2 = st[2u * N + p + j * NRT];
        let x3 = st[3u * N + p + j * NRT];
        for (var l: u32 = 0u; l < 13u; l = l + 1u) {
          acc[l] = acc[l] + ops[base + l0 + l + p * NRT] * x1
                          + ops[base + l0 + l + p * NRT + 676u] * x2
                          + ops[base + l0 + l + p * NRT + 1352u] * x3;
        }
      }
      for (var l: u32 = 0u; l < 13u; l = l + 1u) { st[l0 + l + j * NRT] = acc[l]; }
    }`),
};

async function main2() {
  await installNodeWebGpu();
  const { device, info } = await requestDevice();
  console.log(`device: ${info}, WS=${WS}`);
  for (const [label, code] of Object.entries(MORE)) await time(device, code, label);
  // dispatch overhead: many tiny dispatches in one submit
  {
    const code = `
@group(0) @binding(0) var<storage, read_write> st: array<f32>;
@compute @workgroup_size(64)
fn main(@builtin(local_invocation_index) lid: u32) { if (lid == 0u) { st[0] = st[0] + 1.0; } }
`;
    const module = device.createShaderModule({ code });
    const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
    const buf = device.createBuffer({ size: 16, usage: GPUBufferUsage.STORAGE });
    const bg = device.createBindGroup({ layout: pipeline.getBindGroupLayout(0), entries: [{ binding: 0, resource: { buffer: buf } }] });
    const enc = device.createCommandEncoder();
    const pass = enc.beginComputePass();
    pass.setPipeline(pipeline); pass.setBindGroup(0, bg);
    for (let i = 0; i < 8192; i++) pass.dispatchWorkgroups(1);
    pass.end();
    const cb = enc.finish();
    const t0 = performance.now();
    device.queue.submit([cb]);
    await device.queue.onSubmittedWorkDone();
    console.log(`dispatch overhead: ${((performance.now() - t0) * 1000 / 8192).toFixed(2)} us/dispatch`);
  }
}
if (process.env.MORE) { main2().then(() => process.exit(0)).catch((e) => { console.error(e); process.exit(1); }); }