[jowens]

A test of the load-balancing programming model of Osama, Porumbescu and Owens (PPoPP 2023) on WebGPU: a load-balanced reduction of each tile (segment) of a CSR-style array, with the schedule separate from the work.

• SegmentedReduce (src/segmentedreduce.mjs) holds the work, which the caller gives as JS functions that return WGSL: atom: (atom, tile) => expression, emit: (tile, value) => statement, and optionally tileStart: (t) => expression in place of an offsets buffer. These are Gridwise's stand-in for the paper's C++ lambdas.
• Schedules (src/schedules.mjs) are classes that write the kernels and call the work through binop, tileStartOf, atomValue and emitTile. SegmentedReduce knows no schedule by more than its name, and a test defines one outside the library.
  • ThreadMapped: one invocation per tile.
  • WorkgroupMapped: one workgroup per tile (the paper's block-mapped).
  • MergePath (Merrill and Garland): tiles + atoms path steps split evenly. Split tiles are combined by a segmented scan within the workgroup, and between workgroups by a fixup that is itself a SegmentedReduce: its atoms are the workgroups' carries, its tiles are the workgroups. It runs merge-path, recursively until one workgroup suffices, so a tile spanning any number of workgroups costs nothing extra. Every tile is emitted once, with its whole value, with no float atomics.
• SpMV (src/spmv.mjs): y = Ax (CSR, f32) as a SegmentedReduce.
• Tests (segred, 32): u32 add, i32 max and SpMV against the CPU; all three schedules and an outside one; 1, 7 and 16 items per invocation; three levels of fixup; kernels repeated in place (trials); uniform, power-law, skewed, one-tile and all-empty shapes; sizes changing between executes.
• misc/benchmark_spmv.mjs: the schedules on uniform, R-MAT and dense-row matrices or Matrix Market files, checked against the CPU.
• docs/segmented-reduce.md: the model, the schedules, writing a schedule, a worked example of the fixup, measurements on an M1 Max, a V100 and an A100 (merge-path at 2^22: 2× thread-mapped on uniform, 8–15× on R-MAT, 1,500–2,800× with one dense row, at the same GFLOP/s on all three).

https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ

---

Generated by Claude Code