// Microbenchmark: what do workgroup/storage barriers and the stage loops cost
// on this device, per simulated step, for one workgroup?
import { installNodeWebGpu, requestDevice } from './gpu.mjs';

const N = 1560;
const WS = Number(process.env.WS ?? 256);
const STEPS = 2048;

function kernel(body, nbar) {
  return `
const N = ${N}u; const WS = ${WS}u; const STEPS = ${STEPS}u;
@group(0) @binding(0) 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) {
    ${Array.from({ length: nbar }, () => body).join('\n')}
  }
  if (lid == 0u) { st[0] = f32(STEPS); }
}
`;
}

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(errs.map((m) => m.message).join('\n'));
  const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });
  const buf = device.createBuffer({ size: 4 * N * 20, usage: GPUBufferUsage.STORAGE });
  const bg = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [{ binding: 0, resource: { buffer: buf } }],
  });
  // warmup
  let enc = device.createCommandEncoder();
  let pass = enc.beginComputePass(); pass.setPipeline(pipeline); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(1); pass.end();
  device.queue.submit([enc.finish()]);
  await device.queue.onSubmittedWorkDone();
  const t0 = performance.now();
  const reps = 4;
  enc = device.createCommandEncoder();
  for (let r = 0; r < reps; r++) {
    pass = enc.beginComputePass(); pass.setPipeline(pipeline); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(1); pass.end();
  }
  device.queue.submit([enc.finish()]);
  await device.queue.onSubmittedWorkDone();
  const ms = (performance.now() - t0) / reps;
  console.log(`${label}: ${(1000 * ms / STEPS).toFixed(2)} us/step`);
  buf.destroy();
}

async function main() {
  await installNodeWebGpu();
  const { device, info } = await requestDevice();
  console.log(`device: ${info}, WS=${WS}`);

  await time(device, kernel('', 0), 'empty loop');
  await time(device, kernel('workgroupBarrier();', 14), '14 workgroupBarrier');
  await time(device, kernel('storageBarrier();', 14), '14 storageBarrier');
  await time(device, kernel(`
    for (var i: u32 = lid; i < N; i = i + WS) {
      st[N + i] = st[i] * 1.0001 + 0.5;
    }
    storageBarrier();`, 10), '10 x (elementwise N + storageBarrier)');
  await time(device, kernel(`
    for (var i: u32 = lid; i < N; i = i + WS) {
      var sum: f32 = 0.0;
      let l = i % 26u; let j = i / 26u;
      for (var q: u32 = 0u; q < 60u; q = q + 1u) {
        sum = sum + st[l + q * 26u] * st[2u * N + q + j * 60u];
      }
      st[N + i] = sum;
    }
    storageBarrier();`, 4), '4 x (60-MAC matmul N + storageBarrier)');
}

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