Compute reduction

// ═══ typeshade example — compute reduction kernel ═══
//
// A self-contained compute pass: each invocation reduces a fixed-size WINDOW of an input
// storage array into one output element (a segmented sum). Exercises two storage buffers
// (read + read_write), a uniform count, a `@workgroup_size` compute entry, and the
// `reduce()` loop-fold combinator. Emits WGSL + the std430 storage layouts.
//
// NOT WebGL2-renderable (`renderable: false`): GLSL ES 3.00 has no compute / SSBO, so the
// GLSL backend fails-closed on this module. The site shows its WGSL + Reflection only —
// this is the WebGPU/compute face of the same DSL.

import {
  fn,
  module,
  f32,
  u32,
  reduce,
  f32T,
  vec3uT,
  vec4uT,
  voidT,
  If,
  Return,
  storageBuffer,
  resource,
  builtin,
} from '../src/index.js';
import type { ShaderExample } from './_shared.js';

// Each invocation sums WINDOW consecutive input elements.
const WINDOW = u32(8);

const inputB = storageBuffer('input', f32T, { group: 0, binding: 0, access: 'read' });
const outputB = storageBuffer('output', f32T, { group: 0, binding: 1, access: 'read_write' });
// .x = number of output elements (one reduced window each).
const params = resource('params', vec4uT, { group: 0, binding: 2 });

const reduceKernel = fn(
  'reduce_windows',
  { gid: builtin('global_invocation_id', vec3uT) },
  voidT,
  ({ gid }) => {
    const idx = gid.x;
    If(idx.ge(params.node.x), () => {
      Return();
    });

    const base = idx.mul(WINDOW);
    // Fold WINDOW elements: acc starts at 0, loop j in [0, WINDOW), accumulate.
    const sum = reduce(
      f32(0),
      u32(0),
      (j) => j.lt(WINDOW),
      (acc, j) => acc.add(inputB.at(base.add(j))),
      u32(1),
    );

    outputB.at(idx).assign(sum);
  },
  { stage: 'compute', workgroupSize: 64 },
);

const reductionModule = module({
  bindings: [inputB.binding, outputB.binding, params.binding],
  funcs: [reduceKernel],
});

export const computeReduction: ShaderExample = {
  id: 'compute-reduction',
  title: 'Compute reduction',
  blurb:
    'A @workgroup_size compute kernel folding a window of an input storage buffer into one output element with reduce(). WebGPU-only — GLSL ES 3.00 has no compute, so this shows WGSL + reflection.',
  category: 'compute',
  file: 'compute-reduction.ts',
  module: reductionModule,
  renderable: false,
};

@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
@group(0) @binding(2) var<uniform> params: vec4<u32>;

@compute @workgroup_size(64)
fn reduce_windows(@builtin(global_invocation_id) gid: vec3<u32>) {
  let _licm0 = (gid.x * 8u);
  if ((gid.x >= params.x)) {
    return;
  }
  var _v0: f32 = 0.0;
  for (var _v1: u32 = 0u; (_v1 < 8u); _v1 = (_v1 + 1u)) {
    _v0 = (_v0 + input[(_licm0 + _v1)]);
  }
  output[gid.x] = _v0;
}