From aad73ceb87e62cd5a49716ca00e50951c0cda9d9 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 17:39:39 -0800 Subject: [PATCH 01/10] SegmentedReduce (thread-mapped and merge-path schedules) and SpMV A load-balanced reduction of each tile of a CSR-style array, after Osama, Porumbescu and Owens (PPoPP 2023): the schedule is separate from the work, which the caller writes as JS functions that return WGSL (atom, emit), Gridwise's stand-in for the paper's C++ lambdas. Merge-path combines split tiles with a segmented scan in the workgroup and a fixup kernel between workgroups, so each tile is emitted once with its whole value (no float atomics needed). SpMV is a SegmentedReduce. Tests (segred) and misc/benchmark_spmv.mjs (uniform, R-MAT, dense-row and Matrix Market). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- docs/docs.md | 3 + docs/segmented-reduce.md | 95 +++++++++++++ examples/regression.mjs | 124 +++++++++++++++++ index.d.ts | 67 ++++++++++ index.mjs | 2 + misc/benchmark_spmv.mjs | 224 +++++++++++++++++++++++++++++++ src/segmentedreduce.mjs | 278 +++++++++++++++++++++++++++++++++++++++ src/spmv.mjs | 56 ++++++++ 8 files changed, 849 insertions(+) create mode 100644 docs/segmented-reduce.md create mode 100644 misc/benchmark_spmv.mjs create mode 100644 src/segmentedreduce.mjs create mode 100644 src/spmv.mjs diff --git a/docs/docs.md b/docs/docs.md index 7f0f851..0e4390f 100644 --- a/docs/docs.md +++ b/docs/docs.md @@ -17,6 +17,9 @@ Complete documentation for Gridwise's OneSweepSort implementation. Covers key-on ### [Histogram]({{ "/histogram/" | relative_url }}) Counts of samples per bin, modeled on CUB's DeviceHistogram: equal-width or level-bounded bins (bin functions `EvenBins`, `RangeBins`, `RadixBins`), the workgroup-histogram helper shared with permutations, and options. +### [Segmented Reduce and SpMV]({{ "/segmented-reduce/" | relative_url }}) +`SegmentedReduce`, a load-balanced reduction of each segment of a CSR-style array with the schedule (thread-mapped or merge-path) separate from the work, which the caller writes as JS functions that return WGSL; and `SpMV`, built on it. + ### [Binary Operations]({{ "/binop/" | relative_url }}) Guide to binary operations (monoids) used in Gridwise's scan and reduce. Documents standard operators (Add, Min, Max, Multiply) and shows how to specialize them by datatype or implement custom binary operators by extending the `BinOp` class. diff --git a/docs/segmented-reduce.md b/docs/segmented-reduce.md new file mode 100644 index 0000000..349ea11 --- /dev/null +++ b/docs/segmented-reduce.md @@ -0,0 +1,95 @@ +--- +layout: page +title: "Gridwise Segmented Reduce and SpMV" +category: gridwise +permalink: /segmented-reduce/ +excerpt: "SegmentedReduce: a load-balanced reduction of each segment of a CSR-style array, with the schedule (thread-mapped or merge-path) separate from the work, which the caller writes as JS functions that return WGSL; SpMV is a SegmentedReduce." +order: 20 +--- + +`SegmentedReduce` reduces each segment of a CSR-style array, with the load-balancing schedule kept separate from the work. `SpMV` (y = Ax, CSR, f32) is a `SegmentedReduce` that supplies the work. + +It follows the programming model of Osama, Porumbescu and Owens, [A Programming Model for GPU Load Balancing](https://escholarship.org/uc/item/9nq090zg) (PPoPP 2023): + +- **atoms** are units of work of equal cost, such as a matrix's nonzeros; +- **tiles** are groups of atoms of any size, such as its rows, given by an offsets array: tile t holds atoms `[offsets[t], offsets[t + 1])`; +- a **schedule** decides which invocation handles which atoms, independent of what the work is. + +For every tile, `SegmentedReduce` combines its atoms' values with a binop and emits the result. + +```js +import { SegmentedReduce, BinOpAdd } from "gridwise"; + +const spmv = new SegmentedReduce({ + device, + schedule: "mergepath", + binop: new BinOpAdd({ datatype: "f32" }), + tiles: rows, + atoms: nnz, + buffers: [ + { name: "indices", type: "array" }, + { name: "values", type: "array" }, + { name: "x", type: "array" }, + { name: "y", type: "array", access: "read_write" }, + ], + atom: (a) => `values[${a}] * x[indices[${a}]]`, + emit: (t, v) => `y[${t}] = ${v};`, +}); +await spmv.execute({ offsets, indices, values, x, y }); +``` + +That is all of `SpMV` (`src/spmv.mjs`), which also offers `rows` and `nnz` as names for `tiles` and `atoms`. + +## The work as JS functions that return WGSL + +The paper's C++ interface hands the schedule the work as lambdas, which the compiler inlines. WGSL has no lambdas, templates or function pointers, so Gridwise uses the nearest thing JavaScript offers: functions that return WGSL. + +- **`atom: (atom, tile) => expression`** gives an atom's value. +- **`emit: (tile, value) => statement`** consumes a tile's result. Without it, results go to `outputBuffer`. + +The schedule calls each function with the names of its own WGSL variables and places the result in its kernel. A schedule never knows what the work is, and the work never knows the schedule. Changing schedules is one option, as in the paper. + +This is code generation by string, and it is deliberately so: JavaScript composes it freely today, and the patterns it needs (functions as parameters, generic element types) are what a future WGSL might support directly. Generation happens once per pipeline, so it costs nothing per execute. + +**Buffers** the functions use are declared in `buffers` (`{ name, type, access }`, access `"read"` unless given) and passed to `execute()` by those names, with `offsets` (tiles + 1 u32s). + +**Options:** + +- `binop` (required): any `BinOp`; it may be non-commutative, since atoms are combined in order; +- `tiles`, `atoms` (required): the number of tiles, and `offsets[tiles]`; both may change between executes; +- `atom` (required), `emit`, `buffers`: above; +- `schedule`: `"mergepath"` (default) or `"thread"`; +- `itemsPerThread`: merge-path steps per invocation (default 7). + +## Schedules + +**`thread`** gives each invocation one tile at a time, striding over the tiles. It is balanced only when tiles are about the same size. One long tile holds up the whole kernel, because a single invocation walks all of it. + +**`mergepath`** (Merrill and Garland, SC 2016) walks the merge of two lists: the tiles' ends and the atom indices. Each step of the path either consumes an atom or ends a tile, so there are tiles + atoms steps, and every invocation takes `itemsPerThread` of them, whatever the tile sizes. A binary search over the offsets finds where each invocation's steps start. + +Tiles split between invocations are put back together in two stages: + +1. **Within a workgroup.** Each invocation emits the tiles it both starts and ends, and carries a partial value for the tile it is in at the end. A segmented inclusive scan of the carries, keyed by tile, gives each invocation what earlier invocations carried into the tile it started in. +2. **Between workgroups.** A tile that began in an earlier workgroup is not emitted by the workgroup that ends it. That workgroup records its value instead, and a second kernel (`fixup`) adds what the earlier workgroups carried in that tile, walking back over them, and emits it. + +So every tile is emitted exactly once, with its whole value, and `emit` can be any WGSL: it is never asked to combine partial results. + +### Fit to WebGPU + +- **Barriers in uniform control flow.** WGSL requires `workgroupBarrier()` where control flow is provably uniform across the workgroup. Merge-path's per-invocation loops have no barriers; its only barriers are in the workgroup scan, whose loop bounds are constants. A block-mapped schedule (a workgroup per tile, scanning atom counts per tile) would need its loop bounds made uniform with `workgroupUniformLoad`. +- **No float atomics.** CUDA merge-path SpMV can finish split rows with `atomicAdd` on floats. WebGPU atomics are 32-bit integers only, so the fixup kernel combines split tiles instead, which also keeps the result deterministic. +- **No subgroups needed.** Both schedules use only workgroup memory and barriers, so they run the same on every device, including those without subgroups. + +## Testing and measuring + +The regression tests (`segred`) check `SegmentedReduce` (u32 add, i32 max) and `SpMV` against the CPU on several shapes: uniform tiles, a power law, a few tiles spanning many workgroups, one tile, all tiles empty. They cover both schedules and 1, 7 and 16 items per invocation, and change `tiles` and `atoms` between executes. + +`misc/benchmark_spmv.mjs` times both schedules on generated matrices (`uniform`, `rmat` scale-free graphs, `denserow` with one row holding half the nonzeros) and on Matrix Market files (`--mtx=`, e.g. from the [SuiteSparse collection](https://sparse.tamu.edu)), checks every result, and reports GFLOP/s and bandwidth against a copy kernel's. + +## Possible improvements not made + +- **Coalesced loads.** Each invocation reads its atoms consecutively, so a subgroup's loads are strided by `itemsPerThread`. CUB stages a tile's atoms through workgroup memory first. +- **One search per workgroup.** Each invocation searches the global offsets. CUB searches once per workgroup, then within a tile of row ends in workgroup memory. +- **A parallel fixup.** A tile spanning w workgroups makes one invocation walk back over w carries. A segmented scan of the carries would bound that for very long tiles. +- **Atoms from the GPU.** The caller passes `atoms` (offsets[tiles]) to size the grid. An indirect dispatch could read it on the GPU instead, for offsets produced by an earlier kernel. +- **More schedules.** Workgroup-mapped and subgroup-mapped schedules, as in the paper. diff --git a/examples/regression.mjs b/examples/regression.mjs index 47adf13..5fb5186 100644 --- a/examples/regression.mjs +++ b/examples/regression.mjs @@ -21,6 +21,8 @@ import { Multisplit } from "../src/multisplit.mjs"; import { ShunPerm } from "../src/shunperm.mjs"; import { SuffixArray, suffixArrayReference, checkSuffixArray } from "../src/suffixarray.mjs"; import { Histogram } from "../src/histogram.mjs"; +import { SegmentedReduce, segmentedReduceReference } from "../src/segmentedreduce.mjs"; +import { SpMV, spmvReference } from "../src/spmv.mjs"; import { EvenBins, RangeBins, RadixBins } from "../src/bins.mjs"; export async function main(navigator) { @@ -1110,6 +1112,103 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} return errors; } + /* SegmentedReduce and SpMV inputs: tile (row) lengths by shape */ + const srHash = (i) => (Math.imul(i ^ 0x2545f491, 0x2c1b3c6d) ^ (i >>> 7)) >>> 0; + const srShapes = { + uniform: (tiles) => Array(tiles).fill(8), + "one long tile": () => [70000], + "all empty": (tiles) => Array(tiles).fill(0), + /* a few long tiles (spanning many workgroups) among short and empty ones */ + skewed: (tiles) => Array.from({ length: tiles }, (_, i) => (i % 1009 === 3 ? 20000 + (i % 7) : srHash(i) % 5)), + /* power law: length ~ 1 / u */ + "power law": (tiles) => Array.from({ length: tiles }, (_, i) => Math.min(50000, Math.floor(1 / ((srHash(i) + 1) / 2 ** 32)) - 1)), + }; + function srOffsets(lengths) { + const offsets = new Uint32Array(lengths.length + 1); + lengths.forEach((l, i) => (offsets[i + 1] = offsets[i] + l)); + return offsets; + } + const srBuffer = (array) => { + const buffer = device.createBuffer({ size: Math.max(16, array.byteLength), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC }); + device.queue.writeBuffer(buffer, 0, array); + return buffer; + }; + async function srRead(buffer, TypedArray, length) { + const mappable = device.createBuffer({ size: Math.max(16, length * 4), usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }); + const encoder = device.createCommandEncoder(); + encoder.copyBufferToBuffer(buffer, 0, mappable, 0, Math.max(16, length * 4)); + device.queue.submit([encoder.finish()]); + await mappable.mapAsync(GPUMapMode.READ); + const data = new TypedArray(mappable.getMappedRange().slice(0, length * 4)); + mappable.destroy(); + return data; + } + + /* SegmentedReduce: each tile's reduction of data[atom] (u32 add or i32 + * max), equal to the CPU's; one primitive over several shapes in turn + * (tiles and atoms change between executes) */ + async function runSegmentedReduceTest({ shapes, tiles, binop, options = {} }) { + const T = binop.datatype; + const TypedArray = T === "i32" ? Int32Array : Uint32Array; + const primitive = new SegmentedReduce({ + device, + binop, + tiles: 1, + atoms: 0, + buffers: [{ name: "data", type: `array<${T}>` }], + atom: (a) => `data[${a}]`, + ...options, + }); + let errors = ""; + for (const shape of shapes) { + const lengths = srShapes[shape](tiles); + const offsets = srOffsets(lengths); + const atoms = offsets[lengths.length]; + const data = TypedArray.from({ length: atoms }, (_, a) => (T === "i32" ? (srHash(a) | 0) >> 8 : srHash(a) & 0xffff)); + const buffers = { offsets: srBuffer(offsets), data: srBuffer(data), outputBuffer: srBuffer(new TypedArray(lengths.length)) }; + Object.assign(primitive, { tiles: lengths.length, atoms }); + await primitive.execute({ ...buffers }); + const out = await srRead(buffers.outputBuffer, TypedArray, lengths.length); + const expected = segmentedReduceReference({ offsets, tiles: lengths.length, binop, atomValue: (a) => data[a] }); + const t = expected.findIndex((v, i) => out[i] !== (T === "i32" ? v | 0 : v >>> 0)); + if (t >= 0) errors += `\n${shape}: tile ${t} (${lengths[t]} atoms) = ${out[t]}, expected ${expected[t]}`; + for (const b of Object.values(buffers)) b.destroy(); + } + primitive.destroy(); + return errors; + } + + /* SpMV: y = A x within f32 rounding of the CPU's f64 sums */ + async function runSpMVTest({ shape, rows, cols, options = {} }) { + const lengths = srShapes[shape](rows); + const offsets = srOffsets(lengths); + const nnz = offsets[lengths.length]; + const unit = (i) => (srHash(i) / 2 ** 31) - 1; + const indices = Uint32Array.from({ length: nnz }, (_, k) => srHash(k + 7) % cols); + const values = Float32Array.from({ length: nnz }, (_, k) => unit(k)); + const x = Float32Array.from({ length: cols }, (_, c) => unit(c + 99991)); + const buffers = { + offsets: srBuffer(offsets), + indices: srBuffer(indices), + values: srBuffer(values), + x: srBuffer(x), + y: srBuffer(new Float32Array(lengths.length)), + }; + const primitive = new SpMV({ device, rows: lengths.length, nnz, ...options }); + await primitive.execute({ ...buffers }); + const y = await srRead(buffers.y, Float32Array, lengths.length); + const expected = spmvReference({ offsets, indices, values, x }); + let errors = ""; + for (let r = 0; r < lengths.length && !errors; r++) { + let magnitude = 0; + for (let k = offsets[r]; k < offsets[r + 1]; k++) magnitude += Math.abs(values[k] * x[indices[k]]); + if (!(Math.abs(y[r] - expected[r]) <= 1e-5 * magnitude + 1e-30)) errors = `\nrow ${r} (${lengths[r]} nonzeros): ${y[r]}, expected ${expected[r]}`; + } + primitive.destroy(); + for (const b of Object.values(buffers)) b.destroy(); + return errors; + } + class RandomBitsTest extends BasePrimitive { constructor(args) { super(args); @@ -1458,6 +1557,31 @@ ${randomBitsWGSL(this.rng)} runSuffixArrayTest({ input: "english", n: 20000, options: { speculate: false }, expectSpeculation: false }) ); + { + const allShapes = ["uniform", "skewed", "one long tile", "all empty", "power law"]; + for (const [label, schedule, extra] of [ + ["thread", "thread", {}], + ["mergepath", "mergepath", {}], + ["mergepath, 1 item per thread", "mergepath", { itemsPerThread: 1 }], + ["mergepath, 16 items per thread", "mergepath", { itemsPerThread: 16 }], + ]) { + await runTest(`segred SegmentedReduce u32 add, ${label}`, () => + runSegmentedReduceTest({ shapes: allShapes, tiles: 30011, binop: new BinOpAdd({ datatype: "u32" }), options: { schedule, ...extra } }) + ); + } + for (const schedule of ["thread", "mergepath"]) { + await runTest(`segred SegmentedReduce i32 max, ${schedule}`, () => + runSegmentedReduceTest({ shapes: ["skewed", "all empty", "power law"], tiles: 5003, binop: new BinOpMax({ datatype: "i32" }), options: { schedule } }) + ); + await runTest(`segred SegmentedReduce one tile, ${schedule}`, () => + runSegmentedReduceTest({ shapes: ["uniform", "all empty"], tiles: 1, binop: new BinOpAdd({ datatype: "u32" }), options: { schedule } }) + ); + for (const shape of ["uniform", "power law", "skewed"]) { + await runTest(`segred SpMV ${shape}, ${schedule}`, () => runSpMVTest({ shape, rows: 20000, cols: 5000, options: { schedule } })); + } + } + } + for (const [label, bins, n, gen, repeat] of histogramTests) { await runTest(`hist ${label} n=${n}`, () => runHistogramTest({ bins, n, gen, repeat })); } diff --git a/index.d.ts b/index.d.ts index 0b56ee2..8c59da1 100644 --- a/index.d.ts +++ b/index.d.ts @@ -482,6 +482,73 @@ export function suffixArrayTraffic(args: { /** the phase a SuffixArray kernel label belongs to */ export function suffixArrayPhase(label: string): string; +/** For each tile (segment) of a CSR-style offsets array, the binop over its + * atoms' values, with a load-balancing schedule separate from the work + * (Osama et al., PPoPP 2023). docs/segmented-reduce.md */ +export interface SegmentedReduceOptions { + device: GPUDevice; + binop: BinOp; + /** number of tiles; offsets holds tiles + 1 entries (may change between executes) */ + tiles: number; + /** offsets[tiles] (may change between executes) */ + atoms: number; + /** WGSL expression of atom `atom` (in tile `tile`)'s value, given the names of the WGSL variables */ + atom: (atom: string, tile: string) => string; + /** WGSL statement that consumes tile `tile`'s result `value` (default: outputBuffer[tile] = value;) */ + emit?: (tile: string, value: string) => string; + /** the buffers atom and emit use, bound by name in execute() (access "read" unless given) */ + buffers?: { name: string; type: string; access?: "read" | "read_write" }[]; + /** "mergepath" (default) or "thread" */ + schedule?: "mergepath" | "thread"; + /** mergepath: path steps per invocation (default 7) */ + itemsPerThread?: number; + adapter?: GPUAdapter; + [key: string]: any; +} +export class SegmentedReduce extends BasePrimitive { + tiles: number; + atoms: number; + schedule: "mergepath" | "thread"; + constructor(args: SegmentedReduceOptions); + /** offsets (tiles + 1 u32s), the buffers named in options.buffers, and + * outputBuffer when there is no emit */ + execute(args: { offsets: GPUBuffer; [buffer: string]: any }): Promise; +} +/** the reduction on the CPU, for checking */ +export function segmentedReduceReference(args: { + offsets: Uint32Array; + tiles: number; + binop: BinOp; + atomValue: (atom: number, tile: number) => number; +}): number[]; + +/** y = A x, A in CSR form (f32): a SegmentedReduce over A's rows. docs/segmented-reduce.md */ +export class SpMV extends SegmentedReduce { + rows: number; + nnz: number; + constructor(args: { + device: GPUDevice; + rows: number; + nnz: number; + schedule?: "mergepath" | "thread"; + itemsPerThread?: number; + adapter?: GPUAdapter; + [key: string]: any; + }); + /** offsets: rows + 1 u32s; indices: nnz u32 columns; values: nnz f32s; x: f32 per column; y: rows f32s (written) */ + execute(args: { + offsets: GPUBuffer; + indices: GPUBuffer; + values: GPUBuffer; + x: GPUBuffer; + y: GPUBuffer; + enableGPUTiming?: boolean; + [key: string]: any; + }): Promise; +} +/** y = A x on the CPU (f64 sums) */ +export function spmvReference(args: { offsets: Uint32Array; indices: Uint32Array; values: Float32Array; x: Float32Array }): Float64Array; + export function tuneMultisplitPerm(args: { device: GPUDevice; n: number; diff --git a/index.mjs b/index.mjs index e5585fd..f1afb94 100644 --- a/index.mjs +++ b/index.mjs @@ -25,6 +25,8 @@ export { suffixArrayTraffic, suffixArrayPhase, } from "./src/suffixarray.mjs"; +export { SegmentedReduce, segmentedReduceReference } from "./src/segmentedreduce.mjs"; +export { SpMV, spmvReference } from "./src/spmv.mjs"; export { randomBitsWGSL, randomBitsJS, diff --git a/misc/benchmark_spmv.mjs b/misc/benchmark_spmv.mjs new file mode 100644 index 0000000..49fe711 --- /dev/null +++ b/misc/benchmark_spmv.mjs @@ -0,0 +1,224 @@ +/* Benchmark SpMV (SegmentedReduce) schedules in Node (Dawn): thread-mapped + * against merge-path, on balanced and unbalanced sparse matrices. + * + * node misc/benchmark_spmv.mjs [matrices] [log2 sizes] [--mtx=file.mtx,...] [--items=K,...] + * node misc/benchmark_spmv.mjs uniform,rmat,denserow 18,20,22 + * node misc/benchmark_spmv.mjs --mtx=soc-LiveJournal1.mtx,road_usa.mtx + * + * Generated matrices, 2^s rows (and columns), about 16 nonzeros per row: + * uniform every row 16 nonzeros, random columns: balanced + * rmat an R-MAT graph (a, b, c, d = 0.57, 0.19, 0.19, 0.05; 16 2^s + * edges; vertex ids shuffled): scale-free, a few rows with + * tens of thousands of nonzeros, many empty + * denserow 8 nonzeros per row, plus one row holding half of all + * nonzeros: the worst case for one invocation per row + * --mtx: Matrix Market files (e.g. from the SuiteSparse collection, + * sparse.tamu.edu): coordinate real, integer or pattern (values 1), + * general or symmetric. --items: merge-path steps per invocation to try + * (default 7). + * + * Per matrix: rows, nonzeros, mean and longest row; then per schedule the + * GPU time (median after a 1 s warm-up), GFLOP/s (2 per nonzero), and the + * bandwidth on the bytes SpMV must move (values, column indices and one x + * per nonzero; offsets and y per row: 12 nnz + 8 rows), against a copy + * kernel's. Every result is checked against a CPU SpMV. See + * docs/segmented-reduce.md. */ +import { readFileSync } from "node:fs"; +import { createNodeDevice } from "./node_device.mjs"; +import { SpMV, spmvReference } from "../src/spmv.mjs"; +import { copyBandwidth } from "./device_rates.mjs"; + +const positional = process.argv.slice(2).filter((a) => !a.startsWith("--")); +const option = (name) => process.argv.find((a) => a.startsWith(`--${name}=`))?.slice(name.length + 3); +const mtxFiles = option("mtx")?.split(",") ?? []; +const itemsList = (option("items") ?? "7").split(",").map(Number); +const generated = mtxFiles.length && !positional.length ? [] : (positional[0] ?? "uniform,rmat,denserow").split(","); +const logSizes = (positional[1] ?? "18,20,22").split(",").map(Number); + +const { adapter, device } = await createNodeDevice(); +console.log(`adapter: ${adapter.info.vendor} ${adapter.info.architecture} ${adapter.info.description ?? ""}`); +if (!device.features.has("timestamp-query")) throw new Error("needs GPU timestamps"); +const copyGBps = await copyBandwidth(device, Math.min(256 << 20, device.limits.maxBufferSize), { adapter }); +console.log(`copy bandwidth (copy kernel, read + write): ${copyGBps.toFixed(1)} GB/s\n`); + +/* ---- matrices */ +function xorshift(seed) { + let x = seed >>> 0 || 1; + return () => { + x ^= x << 13; + x ^= x >>> 17; + x ^= x << 5; + return x >>> 0; + }; +} + +/** CSR from coordinate lists, by a counting sort on rows */ +function csrFromCoo(rows, cols, rowOf, colOf, valueOf) { + const nnz = rowOf.length; + const offsets = new Uint32Array(rows + 1); + for (let k = 0; k < nnz; k++) offsets[rowOf[k] + 1]++; + for (let r = 0; r < rows; r++) offsets[r + 1] += offsets[r]; + const cursor = offsets.slice(0, rows); + const indices = new Uint32Array(nnz); + const values = new Float32Array(nnz); + for (let k = 0; k < nnz; k++) { + const p = cursor[rowOf[k]]++; + indices[p] = colOf[k]; + values[p] = valueOf ? valueOf[k] : 1; + } + return { rows, cols, offsets, indices, values }; +} + +const randomValues = (nnz, rand) => Float32Array.from({ length: nnz }, () => rand() / 2 ** 31 - 1); + +const generators = { + uniform(s) { + const rows = 2 ** s, d = 16, rand = xorshift(0x1234 + s); + const offsets = Uint32Array.from({ length: rows + 1 }, (_, r) => r * d); + const indices = Uint32Array.from({ length: rows * d }, () => rand() % rows); + return { rows, cols: rows, offsets, indices, values: randomValues(rows * d, rand) }; + }, + rmat(s) { + const n = 2 ** s, edges = 16 * n, rand = xorshift(0x5eed + s); + /* quadrant thresholds for a, a + b, a + b + c, as 32-bit integers */ + const [ta, tb, tc] = [0.57, 0.76, 0.95].map((p) => Math.floor(p * 2 ** 32)); + const perm = Uint32Array.from({ length: n }, (_, i) => i); + for (let i = n - 1; i > 0; i--) { + const j = rand() % (i + 1); + [perm[i], perm[j]] = [perm[j], perm[i]]; + } + const rowOf = new Uint32Array(edges), colOf = new Uint32Array(edges); + for (let e = 0; e < edges; e++) { + let r = 0, c = 0; + for (let bit = 0; bit < s; bit++) { + const u = rand(); + r = (r << 1) | (u >= tb ? 1 : 0); + c = (c << 1) | ((u >= ta && u < tb) || u >= tc ? 1 : 0); + } + rowOf[e] = perm[r >>> 0]; + colOf[e] = perm[c >>> 0]; + } + return csrFromCoo(n, n, rowOf, colOf, randomValues(edges, rand)); + }, + denserow(s) { + const rows = 2 ** s, d = 8, dense = d * rows, rand = xorshift(0xd00d + s); + const offsets = new Uint32Array(rows + 1); + const heavy = rows >> 1; /* the dense row, in the middle */ + for (let r = 0; r < rows; r++) offsets[r + 1] = offsets[r] + d + (r === heavy ? dense : 0); + const nnz = offsets[rows]; + const indices = Uint32Array.from({ length: nnz }, () => rand() % rows); + return { rows, cols: rows, offsets, indices, values: randomValues(nnz, rand) }; + }, +}; + +/** A Matrix Market coordinate file, parsed from its bytes */ +function readMatrixMarket(path) { + const bytes = readFileSync(path); + let p = 0; + const line = () => { + const start = p; + while (p < bytes.length && bytes[p] !== 10) p++; + return bytes.toString("latin1", start, p++); + }; + const header = line().toLowerCase().split(/\s+/); + if (header[0] !== "%%matrixmarket" || header[2] !== "coordinate") throw new Error(`${path}: not a Matrix Market coordinate file`); + const field = header[3], symmetric = header[4] !== "general"; + if (!["real", "integer", "pattern"].includes(field)) throw new Error(`${path}: field ${field} not supported`); + let size; + do size = line().trim(); while (size.startsWith("%") || size === ""); + const [rows, cols, entries] = size.split(/\s+/).map(Number); + const skip = () => { while (bytes[p] === 32 || bytes[p] === 9 || bytes[p] === 10 || bytes[p] === 13) p++; }; + const int = () => { skip(); let v = 0; while (bytes[p] >= 48 && bytes[p] <= 57) v = v * 10 + bytes[p++] - 48; return v; }; + const token = () => { skip(); const start = p; while (p < bytes.length && bytes[p] > 32) p++; return Number(bytes.toString("latin1", start, p)); }; + const capacity = symmetric ? 2 * entries : entries; + const rowOf = new Uint32Array(capacity), colOf = new Uint32Array(capacity), valueOf = new Float32Array(capacity); + let k = 0; + for (let e = 0; e < entries; e++) { + const r = int() - 1, c = int() - 1, v = field === "pattern" ? 1 : token(); + rowOf[k] = r; colOf[k] = c; valueOf[k++] = v; + if (symmetric && r !== c) { rowOf[k] = c; colOf[k] = r; valueOf[k++] = v; } + } + return csrFromCoo(rows, cols, rowOf.subarray(0, k), colOf.subarray(0, k), valueOf.subarray(0, k)); +} + +/* ---- runs */ +const usage = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST; +function upload(array) { + const buffer = device.createBuffer({ size: Math.max(16, Math.ceil(array.byteLength / 4) * 4), usage }); + device.queue.writeBuffer(buffer, 0, array); + return buffer; +} +async function download(buffer, length) { + const mappable = device.createBuffer({ size: length * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }); + const encoder = device.createCommandEncoder(); + encoder.copyBufferToBuffer(buffer, 0, mappable, 0, length * 4); + device.queue.submit([encoder.finish()]); + await mappable.mapAsync(GPUMapMode.READ); + const data = new Float32Array(mappable.getMappedRange().slice()); + mappable.destroy(); + return data; +} + +async function benchmark(name, m) { + const { rows, cols, offsets, indices, values } = m; + const nnz = offsets[rows]; + let longest = 0, empty = 0; + for (let r = 0; r < rows; r++) { + const len = offsets[r + 1] - offsets[r]; + longest = Math.max(longest, len); + if (!len) empty++; + } + console.log( + `${name}: ${rows} rows, ${cols} columns, ${nnz} nonzeros; ${(nnz / rows).toFixed(1)} per row, longest ${longest}, ${((100 * empty) / rows).toFixed(0)}% empty` + ); + const rand = xorshift(0xabcdef); + const x = randomValues(cols, rand); + const buffers = { offsets: upload(offsets), indices: upload(indices), values: upload(values), x: upload(x), y: upload(new Float32Array(rows)) }; + const expected = spmvReference({ offsets, indices, values, x }); + const bytes = 12 * nnz + 8 * rows; + const results = []; + const runs = [["thread", {}], ...itemsList.map((K) => [`mergepath${itemsList.length > 1 ? ` K=${K}` : ""}`, { itemsPerThread: K }])]; + for (const [label, extra] of runs) { + const schedule = label.startsWith("thread") ? "thread" : "mergepath"; + const spmv = new SpMV({ device, adapter, rows, nnz, schedule, ...extra }); + device.queue.writeBuffer(buffers.y, 0, new Float32Array(rows)); + const { ns, kernels } = await spmv.benchmark({ ...buffers, warmMs: 1000 }); + const y = await download(buffers.y, rows); + let bad = -1; + for (let r = 0; r < rows && bad < 0; r++) { + let magnitude = 0; + for (let k = offsets[r]; k < offsets[r + 1]; k++) magnitude += Math.abs(values[k] * x[indices[k]]); + if (!(Math.abs(y[r] - expected[r]) <= 1e-4 * magnitude + 1e-30)) bad = r; + } + spmv.destroy(); + const fixup = kernels.find((k) => k.label?.includes("fixup")); + results.push({ label, ns }); + console.log( + ` ${label.padEnd(14)} ${(ns / 1e6).toFixed(3).padStart(9)} ms ${((2 * nnz) / ns).toFixed(2).padStart(7)} GFLOP/s ` + + `${(bytes / ns).toFixed(1).padStart(7)} GB/s (${((100 * bytes) / ns / copyGBps).toFixed(0).padStart(3)}% of copy)` + + (fixup ? ` fixup ${(fixup.ns / 1e3).toFixed(1)} µs` : "") + + (bad >= 0 ? ` WRONG: row ${bad} = ${y[bad]}, expected ${expected[bad]}` : "") + ); + } + const thread = results[0].ns; + const best = Math.min(...results.slice(1).map((r) => r.ns)); + console.log(` merge-path / thread: ${(thread / best).toFixed(2)}x faster\n`); + for (const b of Object.values(buffers)) b.destroy(); +} + +for (const name of generated) { + if (!generators[name]) throw new Error(`unknown matrix ${name}; choose from ${Object.keys(generators)} or --mtx=`); + for (const s of logSizes) { + const start = performance.now(); + const m = generators[name](s); + console.log(`(generated in ${((performance.now() - start) / 1e3).toFixed(1)} s)`); + await benchmark(`${name} 2^${s}`, m); + } +} +for (const path of mtxFiles) { + const start = performance.now(); + const m = readMatrixMarket(path); + console.log(`(read in ${((performance.now() - start) / 1e3).toFixed(1)} s)`); + await benchmark(path.split("/").pop(), m); +} +process.exit(0); diff --git a/src/segmentedreduce.mjs b/src/segmentedreduce.mjs new file mode 100644 index 0000000..50c5656 --- /dev/null +++ b/src/segmentedreduce.mjs @@ -0,0 +1,278 @@ +/** + * SegmentedReduce: a load-balanced reduction of each segment of a CSR-style + * array, with the load-balancing schedule separate from the work. + * + * The work is described as in Osama, Porumbescu and Owens, "A Programming + * Model for GPU Load Balancing" (PPoPP 2023): atoms (units of work of equal + * cost, e.g. a matrix's nonzeros) grouped into tiles (segments of any length, + * e.g. its rows), given by an offsets array: tile t holds atoms + * [offsets[t], offsets[t + 1]). For every tile, SegmentedReduce combines the + * values of its atoms with a binop and emits the result: + * + * result[t] = binop over a in tile t of atom(a, t) (identity if t is empty) + * + * The caller says what an atom's value is and what to do with a tile's result + * as JS functions that return WGSL, called with the names of the WGSL + * variables that hold the atom and tile index: Gridwise's stand-in for the C++ + * lambdas the paper uses. A sparse matrix-vector product (y = Ax, CSR): + * + * const spmv = new SegmentedReduce({ + * device, schedule: "mergepath", binop: new BinOpAdd({ datatype: "f32" }), + * tiles: rows, atoms: nnz, + * buffers: [{ name: "indices", type: "array" }, { name: "values", type: "array" }, + * { name: "x", type: "array" }, { name: "y", type: "array", access: "read_write" }], + * atom: (a) => `values[${a}] * x[indices[${a}]]`, + * emit: (t, v) => `y[${t}] = ${v};`, + * }); + * await spmv.execute({ offsets, indices, values, x, y }); + * + * Without emit, results go to outputBuffer (array of the binop's datatype). + * Buffers named in `buffers` are bound by those names (access "read" unless + * given). The binop may be non-commutative: atoms are combined in order. + * + * Options: binop (required); tiles and atoms (required: the number of tiles, + * and offsets[tiles]; both may change between executes); atom (required), + * emit, buffers; schedule ("mergepath", default, or "thread"); + * itemsPerThread (mergepath: path steps per invocation, default 7). + * + * Schedules (the paper's names): + * thread one invocation per tile, striding over the tiles. Balanced + * only when tiles are about the same size: one long tile + * holds up the kernel. + * mergepath the merge path of the tile ends and the atom indices (Merrill + * and Garland, SC 2016): tiles + atoms steps, split evenly, so + * every invocation does itemsPerThread steps whatever the tile + * sizes. A tile split between invocations is combined by a + * segmented scan in the workgroup, and between workgroups by a + * second kernel (fixup) that walks back over the workgroups the + * tile spans: each tile is emitted exactly once, with its whole + * value, so emit may be any WGSL (it is never asked to combine). + */ +import { BasePrimitive, Kernel, AllocateBuffer, UniformStruct } from "./primitive.mjs"; + +const WORKGROUP_SIZE = 256; +const SCHEDULES = ["mergepath", "thread"]; + +export class SegmentedReduce extends BasePrimitive { + constructor(args) { + super(args); + if (!this.binop) throw new Error("SegmentedReduce: requires a binop (e.g. new BinOpAdd({ datatype: \"f32\" }))"); + if (typeof this.atom !== "function") throw new Error("SegmentedReduce: requires atom: (atom, tile) => WGSL expression"); + this.schedule ??= "mergepath"; + if (!SCHEDULES.includes(this.schedule)) throw new Error(`SegmentedReduce: schedule must be one of ${SCHEDULES.join(", ")}`); + this.itemsPerThread ??= 7; + if (!(Number.isInteger(this.itemsPerThread) && this.itemsPerThread >= 1)) throw new Error("SegmentedReduce: itemsPerThread must be a positive integer"); + this.datatype = this.binop.datatype; + this.userBuffers = (this.buffers ?? []).map((b) => ({ access: "read", ...b })); + if (!this.emit) { + this.emit = (t, v) => `outputBuffer[${t}] = ${v};`; + this.userBuffers.push({ name: "outputBuffer", access: "read_write", type: `array<${this.datatype}>` }); + } + this.knownBuffers = ["offsets", ...this.userBuffers.map((b) => b.name)]; + this.workgroupSize = WORKGROUP_SIZE; + this.params = new UniformStruct("SRParams", { tiles: "u32", atoms: "u32", workgroups: "u32" }); + if (!("label" in args)) this.label += this.makeParamString(["schedule"]); + } + + /** workgroups of the main kernel */ + workgroupCount() { + return this.schedule === "thread" + ? Math.ceil(this.tiles / WORKGROUP_SIZE) + : Math.ceil((this.tiles + this.atoms) / (WORKGROUP_SIZE * this.itemsPerThread)); + } + + /* the caller's pieces, as WGSL functions the schedules call */ + definitions() { + const T = this.datatype; + return /* wgsl */ ` +${this.binop.wgslfn} +const IDENTITY: ${T} = ${T}(${this.binop.identity}); +fn atomValue(atom: u32, tile: u32) -> ${T} { return ${this.atom("atom", "tile")}; } +fn emitTile(tile: u32, value: ${T}) { ${this.emit("tile", "value")} }`; + } + + threadWGSL = () => /* wgsl */ ` +${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} +${this.definitions()} +@compute @workgroup_size(${WORKGROUP_SIZE}) +fn thread_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${this.fnDeclarations.computeLinearizedGridParametersSplit()} + for (var tile = gid; tile < sr.tiles; tile += totalThreadCount) { + var acc = IDENTITY; + let end = offsets[tile + 1u]; + for (var atom = offsets[tile]; atom < end; atom++) { acc = binop(acc, atomValue(atom, tile)); } + emitTile(tile, acc); + } +}`; + + /* Merge path: step d of the path either consumes an atom (atom < the + * current tile's end) or ends the current tile. Invocation i takes steps + * [i K, (i + 1) K); a binary search finds where its path starts. */ + mergepathWGSL = () => { + const T = this.datatype; + const W = WORKGROUP_SIZE; + return /* wgsl */ ` +${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} +${this.definitions()} +const K = ${this.itemsPerThread}u; +const NONE = 0xffffffffu; + +/* the carried (unfinished) tile and value of each invocation, then their + * segmented inclusive scan */ +var scanTile: array; +var scanValue: array<${T}, ${W}>; +var wgStart: u32; /* the tile the workgroup's path starts in */ +var headDone: u32; /* did the workgroup end it? */ +var headValue: ${T}; + +@compute @workgroup_size(${W}) +fn mergepath_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${this.fnDeclarations.computeLinearizedGridParametersSplit()} + if (wgid >= sr.workgroups) { return; } + let lidx = builtinsNonuniform.lidx; + let total = sr.tiles + sr.atoms; + let diag = min(gid * K, total); + /* where this invocation's path starts: the most tiles ended (tile) such + * that tile ends come no later than the atoms before them */ + var lo = select(0u, diag - sr.atoms, diag > sr.atoms); + var hi = min(diag, sr.tiles); + while (lo < hi) { + let mid = (lo + hi) / 2u; + if (offsets[mid + 1u] <= diag - mid - 1u) { lo = mid + 1u; } else { hi = mid; } + } + var tile = lo; + var atom = diag - lo; + var tileEnd = select(sr.atoms, offsets[tile + 1u], tile < sr.tiles); + let stepEnd = min(diag + K, total); + var acc = IDENTITY; + var firstDone = false; /* ended the tile it started in? */ + var firstValue = IDENTITY; + for (var d = diag; d < stepEnd; d++) { + if (atom < tileEnd) { + acc = binop(acc, atomValue(atom, tile)); + atom++; + } else { + /* end of tile: the first tile ended may have started in an earlier + * invocation, so it waits for the scan; later ones are whole */ + if (!firstDone) { firstDone = true; firstValue = acc; } else { emitTile(tile, acc); } + acc = IDENTITY; + tile++; + tileEnd = select(sr.atoms, offsets[tile + 1u], tile < sr.tiles); + } + } + if (lidx == 0u) { + wgStart = lo; + headDone = 0u; + } + /* segmented inclusive scan of the carried values, keyed by tile: tiles are + * nondecreasing across invocations, so a tile's carries are contiguous */ + scanTile[lidx] = tile; + scanValue[lidx] = acc; + workgroupBarrier(); + for (var off = 1u; off < ${W}u; off <<= 1u) { + let joins = lidx >= off && scanTile[lidx - off] == scanTile[lidx]; + var earlier = IDENTITY; + if (joins) { earlier = scanValue[lidx - off]; } + workgroupBarrier(); + if (joins) { scanValue[lidx] = binop(earlier, scanValue[lidx]); } + workgroupBarrier(); + } + if (firstDone) { + /* the tile this invocation started in: add what earlier invocations of + * this workgroup carried in it */ + var value = firstValue; + if (lidx > 0u && scanTile[lidx - 1u] == lo) { value = binop(scanValue[lidx - 1u], value); } + if (wgid > 0u && lo == wgStart) { + /* the workgroup's first tile (one invocation ends it), possibly begun + * in earlier workgroups: the fixup kernel adds theirs and emits it */ + headDone = 1u; + headValue = value; + } else { + emitTile(lo, value); + } + } + workgroupBarrier(); + if (lidx == ${W - 1}u) { + /* what this workgroup carries out (to the tile unfinished at its end), + * and its first tile's value if it ended here */ + wgTile[wgid] = scanTile[${W - 1}u]; + wgValue[wgid] = scanValue[${W - 1}u]; + wgTile[sr.workgroups + wgid] = headDone; + wgValue[sr.workgroups + wgid] = headValue; + } +} + +/* Workgroup w (> 0) that ended the tile it started in (its head) adds what + * the earlier workgroups carried in that tile, walking back while they + * carried it, and emits it. */ +@compute @workgroup_size(${W}) +fn fixup_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${this.fnDeclarations.computeLinearizedGridParametersSplit()} + let w = gid; + if (w == 0u || w >= sr.workgroups) { return; } + if (wgTile[sr.workgroups + w] == 0u) { return; } + let tile = wgTile[w - 1u]; + var value = wgValue[sr.workgroups + w]; + var v = w; + while (v > 0u && wgTile[v - 1u] == tile) { + v--; + value = binop(wgValue[v], value); + } + emitTile(tile, value); +}`; + }; + + compute() { + if (!(Number.isInteger(this.tiles) && this.tiles >= 1)) throw new Error(`SegmentedReduce: tiles (${this.tiles}) must be a positive integer`); + if (!(Number.isInteger(this.atoms) && this.atoms >= 0)) throw new Error(`SegmentedReduce: atoms (${this.atoms}) must be a nonnegative integer`); + if (this.tiles + this.atoms >= 2 ** 32 - WORKGROUP_SIZE * this.itemsPerThread) throw new Error("SegmentedReduce: tiles + atoms must fit in 32 bits"); + const workgroups = this.workgroupCount(); + const T = this.datatype; + const buffers = [ + { name: "offsets", access: "read", type: "array" }, + ...this.userBuffers, + ...(this.schedule === "mergepath" + ? [ + { name: "wgTile", access: "read_write", type: "array" }, + { name: "wgValue", access: "read_write", type: `array<${T}>` }, + ] + : []), + { name: "sr", buffer: "srParams", uniform: this.params }, + ]; + const dispatch = (count) => () => this.getSimpleDispatchGeometry({ workgroupCount: count }); + const label = `SegmentedReduce (${this.schedule})`; + const actions = [this.params.allocate("srParams", { tiles: this.tiles, atoms: this.atoms, workgroups })]; + if (this.schedule === "thread") { + actions.push( + new Kernel({ kernel: this.threadWGSL, entryPoint: "thread_main", label, buffers, getDispatchGeometry: dispatch(workgroups) }) + ); + } else { + actions.push( + /* per workgroup: carried tile and value, then head flag and value */ + new AllocateBuffer({ label: "wgTile", size: Math.max(16, 8 * workgroups), clearBufferOnReuse: false }), + new AllocateBuffer({ label: "wgValue", size: Math.max(16, 8 * workgroups), clearBufferOnReuse: false }), + new Kernel({ kernel: this.mergepathWGSL, entryPoint: "mergepath_main", label: `${label} path`, buffers, getDispatchGeometry: dispatch(workgroups) }), + new Kernel({ + kernel: this.mergepathWGSL, + entryPoint: "fixup_main", + label: `${label} fixup`, + buffers, + getDispatchGeometry: dispatch(Math.ceil(workgroups / WORKGROUP_SIZE)), + }) + ); + } + return actions; + } +} + +/** The reduction on the CPU: result[t] = binop over tile t's atoms of + * atomValue(a, t), for checking. offsets: tiles + 1 entries. */ +export function segmentedReduceReference({ offsets, tiles, binop, atomValue }) { + const result = new Array(tiles); + for (let t = 0; t < tiles; t++) { + let acc = binop.identity; + for (let a = offsets[t]; a < offsets[t + 1]; a++) acc = binop.op(acc, atomValue(a, t)); + result[t] = acc; + } + return result; +} diff --git a/src/spmv.mjs b/src/spmv.mjs new file mode 100644 index 0000000..cd4a288 --- /dev/null +++ b/src/spmv.mjs @@ -0,0 +1,56 @@ +/** + * SpMV: y = A x for a sparse matrix A in CSR form (f32), a SegmentedReduce + * whose tiles are A's rows and whose atoms are its nonzeros. + * + * const spmv = new SpMV({ device, rows, nnz, schedule: "mergepath" }); + * await spmv.execute({ offsets, indices, values, x, y }); + * + * Buffers: offsets (rows + 1 u32: row r's nonzeros are + * [offsets[r], offsets[r + 1])), indices (nnz u32 column indices), values + * (nnz f32), x (f32, one per column), y (rows f32, written). + * Options: rows, nnz (both may change between executes), schedule + * ("mergepath", default, or "thread"), itemsPerThread; see SegmentedReduce. + */ +import { SegmentedReduce } from "./segmentedreduce.mjs"; +import { BinOpAdd } from "./binop.mjs"; + +export class SpMV extends SegmentedReduce { + constructor(args) { + super({ + binop: new BinOpAdd({ datatype: "f32" }), + buffers: [ + { name: "indices", type: "array" }, + { name: "values", type: "array" }, + { name: "x", type: "array" }, + { name: "y", type: "array", access: "read_write" }, + ], + atom: (a) => `values[${a}] * x[indices[${a}]]`, + emit: (r, v) => `y[${r}] = ${v};`, + ...args, + }); + } + get rows() { + return this.tiles; + } + set rows(r) { + this.tiles = r; + } + get nnz() { + return this.atoms; + } + set nnz(n) { + this.atoms = n; + } +} + +/** y = A x on the CPU (f64 sums), for checking. */ +export function spmvReference({ offsets, indices, values, x }) { + const rows = offsets.length - 1; + const y = new Float64Array(rows); + for (let r = 0; r < rows; r++) { + let sum = 0; + for (let k = offsets[r]; k < offsets[r + 1]; k++) sum += values[k] * x[indices[k]]; + y[r] = sum; + } + return y; +} From 8e9af8168e3b3486a48826bee5d2f92412d89add Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 18:08:15 -0800 Subject: [PATCH 02/10] SegmentedReduce: the fixup is a SegmentedReduce over the workgroup carries Combining the workgroups' carries is itself a segmented reduce: atoms are the carries, tiles are the workgroups (each holding the run of carries in the tile it ends, found by a binary search), and emit adds the run into the workgroup's head value. It runs merge-path too, recursively until one workgroup suffices, and binds only its own carry buffers; a last kernel emits the heads. Replaces the walk-back fixup, which took half the time on a matrix with one row of 2^25 nonzeros. Adds tileStart (a function in place of the offsets buffer) and wgsl options; itemsPerThread defaults to 4, the fastest on an M1 Max, a V100 and an A100. Test: three levels of fixup. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- docs/segmented-reduce.md | 23 ++++-- examples/regression.mjs | 18 ++++- misc/benchmark_spmv.mjs | 4 +- src/segmentedreduce.mjs | 156 ++++++++++++++++++++++++++++----------- 4 files changed, 148 insertions(+), 53 deletions(-) diff --git a/docs/segmented-reduce.md b/docs/segmented-reduce.md index 349ea11..01f6943 100644 --- a/docs/segmented-reduce.md +++ b/docs/segmented-reduce.md @@ -58,8 +58,10 @@ This is code generation by string, and it is deliberately so: JavaScript compose - `binop` (required): any `BinOp`; it may be non-commutative, since atoms are combined in order; - `tiles`, `atoms` (required): the number of tiles, and `offsets[tiles]`; both may change between executes; - `atom` (required), `emit`, `buffers`: above; +- `tileStart: (t) => expression`: tile t's first atom, for t from 0 to `tiles`, in place of an `offsets` buffer (the paper's iterators over a sparse format); +- `wgsl`: more WGSL declarations, such as functions that `atom`, `emit` or `tileStart` call; - `schedule`: `"mergepath"` (default) or `"thread"`; -- `itemsPerThread`: merge-path steps per invocation (default 7). +- `itemsPerThread`: merge-path steps per invocation (default 4, the fastest on an M1 Max, a V100 and an A100). ## Schedules @@ -70,19 +72,29 @@ This is code generation by string, and it is deliberately so: JavaScript compose Tiles split between invocations are put back together in two stages: 1. **Within a workgroup.** Each invocation emits the tiles it both starts and ends, and carries a partial value for the tile it is in at the end. A segmented inclusive scan of the carries, keyed by tile, gives each invocation what earlier invocations carried into the tile it started in. -2. **Between workgroups.** A tile that began in an earlier workgroup is not emitted by the workgroup that ends it. That workgroup records its value instead, and a second kernel (`fixup`) adds what the earlier workgroups carried in that tile, walking back over them, and emits it. +2. **Between workgroups.** A tile that began in an earlier workgroup is not emitted by the workgroup that ends it. That workgroup records the tile's value as its *head* instead, and every workgroup records what it carries out. -So every tile is emitted exactly once, with its whole value, and `emit` can be any WGSL: it is never asked to combine partial results. +Combining the workgroups' carries is itself a segmented reduce, so `SegmentedReduce` does it with a second `SegmentedReduce` (the *fixup*): + +- its atoms are the workgroups' carries; +- its tiles are the workgroups: tile w holds the run of carries in the tile that workgroup w ends (carried tiles are nondecreasing, so a binary search finds where the run starts), and none if w ends nothing it did not start; +- its `atom` reads a carry, its `tileStart` is that binary search, and its `emit` adds the run's sum in front of workgroup w's head value. + +The fixup uses the merge-path schedule too, so it is load-balanced however many workgroups a tile spans, and it has its own fixup, recursively, until one workgroup suffices. Each level shrinks the problem by a factor of 512 (with 4 items per invocation), so 2^28 atoms take three levels. A last kernel emits the heads through the caller's `emit`. + +So every tile is emitted exactly once, with its whole value, and `emit` can be any WGSL: it is never asked to combine partial results. The fixup levels bind only their own carry buffers, never the caller's, so the number of buffers a kernel binds does not grow with depth (WebGPU guarantees only 8 storage buffers per shader stage). + +The first version walked back over the carries with one invocation per finishing workgroup. On a matrix with one row holding half of 2^26 nonzeros, that walk crossed 33,000 workgroups in one invocation and took half the time (7.4 of 14.9 ms on an M1 Max). ### Fit to WebGPU - **Barriers in uniform control flow.** WGSL requires `workgroupBarrier()` where control flow is provably uniform across the workgroup. Merge-path's per-invocation loops have no barriers; its only barriers are in the workgroup scan, whose loop bounds are constants. A block-mapped schedule (a workgroup per tile, scanning atom counts per tile) would need its loop bounds made uniform with `workgroupUniformLoad`. -- **No float atomics.** CUDA merge-path SpMV can finish split rows with `atomicAdd` on floats. WebGPU atomics are 32-bit integers only, so the fixup kernel combines split tiles instead, which also keeps the result deterministic. +- **No float atomics.** CUDA merge-path SpMV can finish split rows with `atomicAdd` on floats. WebGPU atomics are 32-bit integers only, so the fixup combines split tiles instead, which also keeps the result deterministic. - **No subgroups needed.** Both schedules use only workgroup memory and barriers, so they run the same on every device, including those without subgroups. ## Testing and measuring -The regression tests (`segred`) check `SegmentedReduce` (u32 add, i32 max) and `SpMV` against the CPU on several shapes: uniform tiles, a power law, a few tiles spanning many workgroups, one tile, all tiles empty. They cover both schedules and 1, 7 and 16 items per invocation, and change `tiles` and `atoms` between executes. +The regression tests (`segred`) check `SegmentedReduce` (u32 add, i32 max) and `SpMV` against the CPU on several shapes: uniform tiles, a power law, a few tiles spanning many workgroups, one tile, all tiles empty. They cover both schedules, 1, 7 and 16 items per invocation, three levels of fixup (one tile of 5 million atoms, one item per invocation), and change `tiles` and `atoms` between executes. `misc/benchmark_spmv.mjs` times both schedules on generated matrices (`uniform`, `rmat` scale-free graphs, `denserow` with one row holding half the nonzeros) and on Matrix Market files (`--mtx=`, e.g. from the [SuiteSparse collection](https://sparse.tamu.edu)), checks every result, and reports GFLOP/s and bandwidth against a copy kernel's. @@ -90,6 +102,5 @@ The regression tests (`segred`) check `SegmentedReduce` (u32 add, i32 max) and ` - **Coalesced loads.** Each invocation reads its atoms consecutively, so a subgroup's loads are strided by `itemsPerThread`. CUB stages a tile's atoms through workgroup memory first. - **One search per workgroup.** Each invocation searches the global offsets. CUB searches once per workgroup, then within a tile of row ends in workgroup memory. -- **A parallel fixup.** A tile spanning w workgroups makes one invocation walk back over w carries. A segmented scan of the carries would bound that for very long tiles. - **Atoms from the GPU.** The caller passes `atoms` (offsets[tiles]) to size the grid. An indirect dispatch could read it on the GPU instead, for offsets produced by an earlier kernel. - **More schedules.** Workgroup-mapped and subgroup-mapped schedules, as in the paper. diff --git a/examples/regression.mjs b/examples/regression.mjs index 5fb5186..6396168 100644 --- a/examples/regression.mjs +++ b/examples/regression.mjs @@ -1117,6 +1117,8 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} const srShapes = { uniform: (tiles) => Array(tiles).fill(8), "one long tile": () => [70000], + /* with 1 item per invocation: 19500 workgroups, so three levels of fixup */ + "huge tile": () => [3, 5000000, 5], "all empty": (tiles) => Array(tiles).fill(0), /* a few long tiles (spanning many workgroups) among short and empty ones */ skewed: (tiles) => Array.from({ length: tiles }, (_, i) => (i % 1009 === 3 ? 20000 + (i % 7) : srHash(i) % 5)), @@ -1147,7 +1149,7 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} /* SegmentedReduce: each tile's reduction of data[atom] (u32 add or i32 * max), equal to the CPU's; one primitive over several shapes in turn * (tiles and atoms change between executes) */ - async function runSegmentedReduceTest({ shapes, tiles, binop, options = {} }) { + async function runSegmentedReduceTest({ shapes, tiles, binop, options = {}, fixupLevels }) { const T = binop.datatype; const TypedArray = T === "i32" ? Int32Array : Uint32Array; const primitive = new SegmentedReduce({ @@ -1172,6 +1174,11 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} const expected = segmentedReduceReference({ offsets, tiles: lengths.length, binop, atomValue: (a) => data[a] }); const t = expected.findIndex((v, i) => out[i] !== (T === "i32" ? v | 0 : v >>> 0)); if (t >= 0) errors += `\n${shape}: tile ${t} (${lengths[t]} atoms) = ${out[t]}, expected ${expected[t]}`; + if (fixupLevels !== undefined) { + let levels = 0; + for (let p = primitive.fixup; p; p = p.fixup) levels++; + if (levels !== fixupLevels) errors += `\n${shape}: ${levels} levels of fixup, expected ${fixupLevels}`; + } for (const b of Object.values(buffers)) b.destroy(); } primitive.destroy(); @@ -1569,6 +1576,15 @@ ${randomBitsWGSL(this.rng)} runSegmentedReduceTest({ shapes: allShapes, tiles: 30011, binop: new BinOpAdd({ datatype: "u32" }), options: { schedule, ...extra } }) ); } + await runTest("segred SegmentedReduce u32 add, mergepath, three levels of fixup", () => + runSegmentedReduceTest({ + shapes: ["huge tile"], + tiles: 3, + binop: new BinOpAdd({ datatype: "u32" }), + options: { schedule: "mergepath", itemsPerThread: 1 }, + fixupLevels: 3, + }) + ); for (const schedule of ["thread", "mergepath"]) { await runTest(`segred SegmentedReduce i32 max, ${schedule}`, () => runSegmentedReduceTest({ shapes: ["skewed", "all empty", "power law"], tiles: 5003, binop: new BinOpMax({ datatype: "i32" }), options: { schedule } }) diff --git a/misc/benchmark_spmv.mjs b/misc/benchmark_spmv.mjs index 49fe711..a984ac8 100644 --- a/misc/benchmark_spmv.mjs +++ b/misc/benchmark_spmv.mjs @@ -191,12 +191,12 @@ async function benchmark(name, m) { if (!(Math.abs(y[r] - expected[r]) <= 1e-4 * magnitude + 1e-30)) bad = r; } spmv.destroy(); - const fixup = kernels.find((k) => k.label?.includes("fixup")); + const fixupNS = kernels.filter((k) => /fixup|heads/.test(k.label ?? "")).reduce((a, k) => a + k.ns, 0); results.push({ label, ns }); console.log( ` ${label.padEnd(14)} ${(ns / 1e6).toFixed(3).padStart(9)} ms ${((2 * nnz) / ns).toFixed(2).padStart(7)} GFLOP/s ` + `${(bytes / ns).toFixed(1).padStart(7)} GB/s (${((100 * bytes) / ns / copyGBps).toFixed(0).padStart(3)}% of copy)` + - (fixup ? ` fixup ${(fixup.ns / 1e3).toFixed(1)} µs` : "") + + (fixupNS ? ` fixup ${(fixupNS / 1e3).toFixed(1)} µs` : "") + (bad >= 0 ? ` WRONG: row ${bad} = ${y[bad]}, expected ${expected[bad]}` : "") ); } diff --git a/src/segmentedreduce.mjs b/src/segmentedreduce.mjs index 50c5656..82d0366 100644 --- a/src/segmentedreduce.mjs +++ b/src/segmentedreduce.mjs @@ -32,8 +32,11 @@ * * Options: binop (required); tiles and atoms (required: the number of tiles, * and offsets[tiles]; both may change between executes); atom (required), - * emit, buffers; schedule ("mergepath", default, or "thread"); - * itemsPerThread (mergepath: path steps per invocation, default 7). + * emit, buffers; tileStart ((t) => WGSL expression of tile t's first atom, + * for t in [0, tiles], instead of an offsets buffer); wgsl (more WGSL + * declarations, e.g. functions that atom, emit or tileStart call); + * schedule ("mergepath", default, or "thread"); itemsPerThread (mergepath: + * path steps per invocation, default 4). * * Schedules (the paper's names): * thread one invocation per tile, striding over the tiles. Balanced @@ -43,12 +46,18 @@ * and Garland, SC 2016): tiles + atoms steps, split evenly, so * every invocation does itemsPerThread steps whatever the tile * sizes. A tile split between invocations is combined by a - * segmented scan in the workgroup, and between workgroups by a - * second kernel (fixup) that walks back over the workgroups the - * tile spans: each tile is emitted exactly once, with its whole - * value, so emit may be any WGSL (it is never asked to combine). + * segmented scan in the workgroup. Between workgroups, combining + * is itself a segmented reduce: each workgroup's carry is an + * atom, and the tiles are the runs of carries in the same tile, + * each ending at the workgroup that finishes it. So a second + * SegmentedReduce (merge-path, on the carries; recursively, until + * one workgroup suffices) combines them into the finishing + * workgroups' head values, however many workgroups a tile spans, + * and a last kernel emits the heads. Each tile is emitted exactly + * once, with its whole value, so emit may be any WGSL (it is + * never asked to combine). */ -import { BasePrimitive, Kernel, AllocateBuffer, UniformStruct } from "./primitive.mjs"; +import { BasePrimitive, Kernel, AllocateBuffer, ExecutePrimitive, UniformStruct } from "./primitive.mjs"; const WORKGROUP_SIZE = 256; const SCHEDULES = ["mergepath", "thread"]; @@ -60,7 +69,7 @@ export class SegmentedReduce extends BasePrimitive { if (typeof this.atom !== "function") throw new Error("SegmentedReduce: requires atom: (atom, tile) => WGSL expression"); this.schedule ??= "mergepath"; if (!SCHEDULES.includes(this.schedule)) throw new Error(`SegmentedReduce: schedule must be one of ${SCHEDULES.join(", ")}`); - this.itemsPerThread ??= 7; + this.itemsPerThread ??= 4; if (!(Number.isInteger(this.itemsPerThread) && this.itemsPerThread >= 1)) throw new Error("SegmentedReduce: itemsPerThread must be a positive integer"); this.datatype = this.binop.datatype; this.userBuffers = (this.buffers ?? []).map((b) => ({ access: "read", ...b })); @@ -68,7 +77,11 @@ export class SegmentedReduce extends BasePrimitive { this.emit = (t, v) => `outputBuffer[${t}] = ${v};`; this.userBuffers.push({ name: "outputBuffer", access: "read_write", type: `array<${this.datatype}>` }); } - this.knownBuffers = ["offsets", ...this.userBuffers.map((b) => b.name)]; + this.level ??= 0; /* 0 for the caller's; n for the nth level of fixup */ + this.wgTileName = `wgTile${this.level}`; + this.wgValueName = `wgValue${this.level}`; + this.knownBuffers = [...(this.tileStart ? [] : ["offsets"]), ...this.userBuffers.map((b) => b.name)]; + this.tileStart ??= (t) => `offsets[${t}]`; this.workgroupSize = WORKGROUP_SIZE; this.params = new UniformStruct("SRParams", { tiles: "u32", atoms: "u32", workgroups: "u32" }); if (!("label" in args)) this.label += this.makeParamString(["schedule"]); @@ -87,6 +100,8 @@ export class SegmentedReduce extends BasePrimitive { return /* wgsl */ ` ${this.binop.wgslfn} const IDENTITY: ${T} = ${T}(${this.binop.identity}); +${this.wgsl ?? ""} +fn tileStartOf(t: u32) -> u32 { return ${this.tileStart("t")}; } fn atomValue(atom: u32, tile: u32) -> ${T} { return ${this.atom("atom", "tile")}; } fn emitTile(tile: u32, value: ${T}) { ${this.emit("tile", "value")} }`; } @@ -99,8 +114,8 @@ fn thread_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNon ${this.fnDeclarations.computeLinearizedGridParametersSplit()} for (var tile = gid; tile < sr.tiles; tile += totalThreadCount) { var acc = IDENTITY; - let end = offsets[tile + 1u]; - for (var atom = offsets[tile]; atom < end; atom++) { acc = binop(acc, atomValue(atom, tile)); } + let end = tileStartOf(tile + 1u); + for (var atom = tileStartOf(tile); atom < end; atom++) { acc = binop(acc, atomValue(atom, tile)); } emitTile(tile, acc); } }`; @@ -138,11 +153,11 @@ fn mergepath_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: Builtins var hi = min(diag, sr.tiles); while (lo < hi) { let mid = (lo + hi) / 2u; - if (offsets[mid + 1u] <= diag - mid - 1u) { lo = mid + 1u; } else { hi = mid; } + if (tileStartOf(mid + 1u) <= diag - mid - 1u) { lo = mid + 1u; } else { hi = mid; } } var tile = lo; var atom = diag - lo; - var tileEnd = select(sr.atoms, offsets[tile + 1u], tile < sr.tiles); + var tileEnd = select(sr.atoms, tileStartOf(tile + 1u), tile < sr.tiles); let stepEnd = min(diag + K, total); var acc = IDENTITY; var firstDone = false; /* ended the tile it started in? */ @@ -157,7 +172,7 @@ fn mergepath_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: Builtins if (!firstDone) { firstDone = true; firstValue = acc; } else { emitTile(tile, acc); } acc = IDENTITY; tile++; - tileEnd = select(sr.atoms, offsets[tile + 1u], tile < sr.tiles); + tileEnd = select(sr.atoms, tileStartOf(tile + 1u), tile < sr.tiles); } } if (lidx == 0u) { @@ -184,7 +199,7 @@ fn mergepath_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: Builtins if (lidx > 0u && scanTile[lidx - 1u] == lo) { value = binop(scanValue[lidx - 1u], value); } if (wgid > 0u && lo == wgStart) { /* the workgroup's first tile (one invocation ends it), possibly begun - * in earlier workgroups: the fixup kernel adds theirs and emits it */ + * in earlier workgroups: the fixup adds theirs and emits it */ headDone = 1u; headValue = value; } else { @@ -195,33 +210,72 @@ fn mergepath_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: Builtins if (lidx == ${W - 1}u) { /* what this workgroup carries out (to the tile unfinished at its end), * and its first tile's value if it ended here */ - wgTile[wgid] = scanTile[${W - 1}u]; - wgValue[wgid] = scanValue[${W - 1}u]; - wgTile[sr.workgroups + wgid] = headDone; - wgValue[sr.workgroups + wgid] = headValue; + ${this.wgTileName}[wgid] = scanTile[${W - 1}u]; + ${this.wgValueName}[wgid] = scanValue[${W - 1}u]; + ${this.wgTileName}[sr.workgroups + wgid] = headDone; + ${this.wgValueName}[sr.workgroups + wgid] = headValue; } } -/* Workgroup w (> 0) that ended the tile it started in (its head) adds what - * the earlier workgroups carried in that tile, walking back while they - * carried it, and emits it. */ +/* after the fixup has added the earlier workgroups' carries into each head + * value: emit the heads */ @compute @workgroup_size(${W}) -fn fixup_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { +fn heads_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { ${this.fnDeclarations.computeLinearizedGridParametersSplit()} let w = gid; if (w == 0u || w >= sr.workgroups) { return; } - if (wgTile[sr.workgroups + w] == 0u) { return; } - let tile = wgTile[w - 1u]; - var value = wgValue[sr.workgroups + w]; - var v = w; - while (v > 0u && wgTile[v - 1u] == tile) { - v--; - value = binop(wgValue[v], value); + if (${this.wgTileName}[sr.workgroups + w] == 1u) { + emitTile(${this.wgTileName}[w - 1u], ${this.wgValueName}[sr.workgroups + w]); } - emitTile(tile, value); }`; }; + /** The fixup: a SegmentedReduce over this one's W workgroup carries, + * whose tiles are this one's workgroups. Tile w's atoms are the carries of + * the run (the same carried tile) that ends at workgroup w - 1, if w ends + * that tile (w's head); runs start where the carried tiles, which are + * nondecreasing, first equal w - 1's. Its emit adds the run's carries + * before w's head value, in place. (Tile W - 1 also takes the last + * workgroup's carry, past the last tile: the identity.) It reads and writes + * only this level's carry buffers, so every level binds the same few. */ + fixupPrimitive() { + const C = this.wgTileName; + const V = this.wgValueName; + const head = (w) => `${V}[sr.tiles + ${w}]`; + this.fixup ??= new SegmentedReduce({ + device: this.device, + adapter: this.adapter, + label: `${this.label} fixup`, + binop: this.binop, + schedule: "mergepath", + itemsPerThread: this.itemsPerThread, + level: this.level + 1, + tiles: 1, + atoms: 0, + buffers: [ + { name: C, type: "array" }, + { name: V, type: `array<${this.datatype}>`, access: "read_write" }, + ], + wgsl: /* wgsl */ ` +fn runStart(t: u32) -> u32 { + if (t == 0u) { return 0u; } + if (t >= sr.tiles) { return sr.tiles; } + let key = ${C}[t - 1u]; + var lo = 0u; + var hi = t - 1u; + while (lo < hi) { + let mid = (lo + hi) / 2u; + if (${C}[mid] < key) { lo = mid + 1u; } else { hi = mid; } + } + return lo; +}`, + tileStart: (t) => `runStart(${t})`, + atom: (a) => `${V}[${a}]`, + emit: (w, v) => `if (${w} > 0u && ${C}[sr.tiles + ${w}] == 1u) { ${head(w)} = binop(${v}, ${head(w)}); }`, + }); + return this.fixup; + } + compute() { if (!(Number.isInteger(this.tiles) && this.tiles >= 1)) throw new Error(`SegmentedReduce: tiles (${this.tiles}) must be a positive integer`); if (!(Number.isInteger(this.atoms) && this.atoms >= 0)) throw new Error(`SegmentedReduce: atoms (${this.atoms}) must be a nonnegative integer`); @@ -229,12 +283,12 @@ fn fixup_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonu const workgroups = this.workgroupCount(); const T = this.datatype; const buffers = [ - { name: "offsets", access: "read", type: "array" }, + ...(this.knownBuffers.includes("offsets") ? [{ name: "offsets", access: "read", type: "array" }] : []), ...this.userBuffers, ...(this.schedule === "mergepath" ? [ - { name: "wgTile", access: "read_write", type: "array" }, - { name: "wgValue", access: "read_write", type: `array<${T}>` }, + { name: this.wgTileName, access: "read_write", type: "array" }, + { name: this.wgValueName, access: "read_write", type: `array<${T}>` }, ] : []), { name: "sr", buffer: "srParams", uniform: this.params }, @@ -249,20 +303,34 @@ fn fixup_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonu } else { actions.push( /* per workgroup: carried tile and value, then head flag and value */ - new AllocateBuffer({ label: "wgTile", size: Math.max(16, 8 * workgroups), clearBufferOnReuse: false }), - new AllocateBuffer({ label: "wgValue", size: Math.max(16, 8 * workgroups), clearBufferOnReuse: false }), - new Kernel({ kernel: this.mergepathWGSL, entryPoint: "mergepath_main", label: `${label} path`, buffers, getDispatchGeometry: dispatch(workgroups) }), - new Kernel({ - kernel: this.mergepathWGSL, - entryPoint: "fixup_main", - label: `${label} fixup`, - buffers, - getDispatchGeometry: dispatch(Math.ceil(workgroups / WORKGROUP_SIZE)), - }) + new AllocateBuffer({ label: this.wgTileName, size: Math.max(16, 8 * workgroups), clearBufferOnReuse: false }), + new AllocateBuffer({ label: this.wgValueName, size: Math.max(16, 8 * workgroups), clearBufferOnReuse: false }), + new Kernel({ kernel: this.mergepathWGSL, entryPoint: "mergepath_main", label: `${label} path`, buffers, getDispatchGeometry: dispatch(workgroups) }) ); + /* one workgroup emits every tile itself; more need the fixup */ + if (workgroups > 1) { + const fixup = this.fixupPrimitive(); + Object.assign(fixup, { tiles: workgroups, atoms: workgroups }); + const names = [this.wgTileName, this.wgValueName]; + actions.push( + new ExecutePrimitive({ primitive: fixup, buffers: Object.fromEntries(names.map((n) => [n, n])) }), + new Kernel({ + kernel: this.mergepathWGSL, + entryPoint: "heads_main", + label: `${label} heads`, + buffers, + getDispatchGeometry: dispatch(Math.ceil(workgroups / WORKGROUP_SIZE)), + }) + ); + } } return actions; } + + destroy() { + this.fixup?.destroy(); + super.destroy(); + } } /** The reduction on the CPU: result[t] = binop over tile t's atoms of From f57abf5b632af0c952e90f05645d19ed6a9108ad Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 18:09:07 -0800 Subject: [PATCH 03/10] docs: how merge-path's fixup works, with a worked example Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- docs/segmented-reduce.md | 54 ++++++++++++++++++++++++++++++++-------- 1 file changed, 44 insertions(+), 10 deletions(-) diff --git a/docs/segmented-reduce.md b/docs/segmented-reduce.md index 01f6943..878c5d5 100644 --- a/docs/segmented-reduce.md +++ b/docs/segmented-reduce.md @@ -69,22 +69,56 @@ This is code generation by string, and it is deliberately so: JavaScript compose **`mergepath`** (Merrill and Garland, SC 2016) walks the merge of two lists: the tiles' ends and the atom indices. Each step of the path either consumes an atom or ends a tile, so there are tiles + atoms steps, and every invocation takes `itemsPerThread` of them, whatever the tile sizes. A binary search over the offsets finds where each invocation's steps start. -Tiles split between invocations are put back together in two stages: +### How merge-path puts split tiles back together -1. **Within a workgroup.** Each invocation emits the tiles it both starts and ends, and carries a partial value for the tile it is in at the end. A segmented inclusive scan of the carries, keyed by tile, gives each invocation what earlier invocations carried into the tile it started in. -2. **Between workgroups.** A tile that began in an earlier workgroup is not emitted by the workgroup that ends it. That workgroup records the tile's value as its *head* instead, and every workgroup records what it carries out. +Merge-path balances the work by cutting the path into equal pieces, wherever the cuts fall. So a tile can be cut between invocations, and a long tile between many workgroups. Each piece of a tile computes a partial value, and the partial values have to be combined, in order, into one value per tile before it is emitted. That happens in two stages: within a workgroup, then between workgroups. -Combining the workgroups' carries is itself a segmented reduce, so `SegmentedReduce` does it with a second `SegmentedReduce` (the *fixup*): +**Within a workgroup.** Each invocation walks its `itemsPerThread` steps: -- its atoms are the workgroups' carries; -- its tiles are the workgroups: tile w holds the run of carries in the tile that workgroup w ends (carried tiles are nondecreasing, so a binary search finds where the run starts), and none if w ends nothing it did not start; -- its `atom` reads a carry, its `tileStart` is that binary search, and its `emit` adds the run's sum in front of workgroup w's head value. +- a tile it both starts and ends is whole: it emits it right away; +- the tile it ends first may have started in an earlier invocation, so it holds that partial value back; +- the tile it is in when its steps run out is unfinished: its partial value is the invocation's *carry*. -The fixup uses the merge-path schedule too, so it is load-balanced however many workgroups a tile spans, and it has its own fixup, recursively, until one workgroup suffices. Each level shrinks the problem by a factor of 512 (with 4 items per invocation), so 2^28 atoms take three levels. A last kernel emits the heads through the caller's `emit`. +Tiles only increase along the path, so the invocations that carry the same tile sit next to each other. A segmented inclusive scan of the carries in workgroup memory, keyed by tile, gives each invocation the sum of what the invocations before it carried in its tile. The invocation that ends that tile adds it in front of its held-back value and emits the tile. The scan's loop bounds are constants, so its barriers satisfy WGSL's uniformity rules. -So every tile is emitted exactly once, with its whole value, and `emit` can be any WGSL: it is never asked to combine partial results. The fixup levels bind only their own carry buffers, never the caller's, so the number of buffers a kernel binds does not grow with depth (WebGPU guarantees only 8 storage buffers per shader stage). +**Between workgroups.** The first tile a workgroup ends may have started in an earlier workgroup, whose partial values this workgroup cannot see: workgroups don't wait for each other (WebGPU makes no promise that they run in any order, or at the same time). So the workgroup doesn't emit that tile. It records it as its *head* (a flag and the value so far), and records what it carries out of its last tile. (Workgroup 0 has nothing before it, so it emits its first tile.) What remains is to add, in front of each head, the carries of the earlier workgroups in that tile. -The first version walked back over the carries with one invocation per finishing workgroup. On a matrix with one row holding half of 2^26 nonzeros, that walk crossed 33,000 workgroups in one invocation and took half the time (7.4 of 14.9 ms on an M1 Max). +Here is an example with five workgroups and tiles A to E. Tile B is long: it starts in workgroup 0 and ends in workgroup 3. In SpMV the tiles are rows, so b0 is the part of row B's dot product computed in workgroup 0. + +| Workgroup | Starts in | Emits itself | Head (tile it ends first, value so far) | Carries out | +|---|---|---|---|---| +| 0 | A | A | (none: workgroup 0 emits A) | B: b0 | +| 1 | B | | (none: ends nothing) | B: b1 | +| 2 | B | | (none: ends nothing) | B: b2 | +| 3 | B | C | B: b3 | D: d3 | +| 4 | D | E | D: d4 | (past the end): identity | + +B's value is b0 + b1 + b2 + b3, and D's is d3 + d4. Each head needs the *run* of carries in its tile that ends just before it: workgroup 3's head needs carries 0–2, and workgroup 4's needs carry 3. + +**The fixup is a segmented reduce.** That last sentence describes a segmented reduce, so `SegmentedReduce` does it with another `SegmentedReduce`: + +- **Atoms:** the workgroups' carries, here b0, b1, b2, d3 and the identity. +- **Tiles:** one per workgroup. Tile w holds the run of carries in the tile that workgroup w-1 carried out. That run ends at carry w-1, and it starts where the carried tiles first equal carry w-1's tile. The carried tiles are nondecreasing, so a binary search finds that start. Where workgroup w carries the same tile as w-1, the run hasn't ended, and tile w is empty. +- **Its emit:** for each workgroup w that has a head, it adds the run's sum *in front of* w's head value, in place. + +In the example, the fixup's tile starts (its `tileStart`) are 0, 0, 0, 0, 3, 5, so: + +| Fixup tile (workgroup) | Its atoms (carries) | Its emit | +|---|---|---| +| 0, 1, 2 | none | (no heads) | +| 3 | 0–2: b0, b1, b2 | head 3 = (b0 + b1 + b2) + b3 = B | +| 4 | 3–4: d3, identity | head 4 = (d3 + identity) + d4 = D | + +The last tile also takes the last workgroup's carry, which is past the end of the tiles and holds the identity, so it changes nothing. Then one more kernel emits each head through the caller's `emit`: B from workgroup 3, D from workgroup 4. + +**Why this is the right structure:** + +- **It is load-balanced at every level.** The fixup's atoms and tiles go through the merge-path schedule too, so a tile spanning 30,000 workgroups costs the same as 30,000 short tiles. The first version instead had each head walk back over its run with one invocation. On a matrix with one row holding half of 2^26 nonzeros, that walk crossed 33,000 workgroups and took half the time (7.4 of 14.9 ms on an M1 Max). +- **It recurses.** The fixup splits its own runs across its workgroups and has its own fixup, until one workgroup suffices. A level of W workgroups has W tiles and W atoms, so its path has 2W steps, and with 256 invocations of 4 steps each the next level has W / 512 workgroups. A path of 2^28 steps (tiles plus atoms) takes two levels of fixup. +- **It reuses the abstraction.** The fixup is written exactly as a caller would write any `SegmentedReduce`: `atom` reads a carry, `tileStart` is the binary search, and `emit` writes into the head. Only the buffers differ. +- **It keeps order.** Each run's carries are combined in order, and in front of the head value, so the binop need not be commutative. +- **It binds few buffers.** A fixup level reads and writes only its parent's carry buffers, never the caller's, so the number of buffers a kernel binds doesn't grow with depth. (WebGPU guarantees only 8 storage buffers per shader stage.) +- **Every tile is emitted exactly once, with its whole value.** So `emit` can be any WGSL: it is never asked to combine partial results, and no atomics are needed. ### Fit to WebGPU From 35f6e845bc75fc145ff11da9824235c0e481283c Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 18:19:04 -0800 Subject: [PATCH 04/10] SegmentedReduce: fixup writes the whole head value to its own slot The fixup's emit added the run's carries into the head value in place, so a repeated kernel (benchmark()'s trials) added them again: wrong results on the M1 Max benchmark. It now writes to a third slot; every kernel can be repeated. Tests with kernels repeated (trials 2 and 3), which fail without the fix. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- docs/segmented-reduce.md | 1 + examples/regression.mjs | 21 +++++++++++++++++---- src/segmentedreduce.mjs | 13 ++++++++----- 3 files changed, 26 insertions(+), 9 deletions(-) diff --git a/docs/segmented-reduce.md b/docs/segmented-reduce.md index 878c5d5..33c89f6 100644 --- a/docs/segmented-reduce.md +++ b/docs/segmented-reduce.md @@ -118,6 +118,7 @@ The last tile also takes the last workgroup's carry, which is past the end of th - **It reuses the abstraction.** The fixup is written exactly as a caller would write any `SegmentedReduce`: `atom` reads a carry, `tileStart` is the binary search, and `emit` writes into the head. Only the buffers differ. - **It keeps order.** Each run's carries are combined in order, and in front of the head value, so the binop need not be commutative. - **It binds few buffers.** A fixup level reads and writes only its parent's carry buffers, never the caller's, so the number of buffers a kernel binds doesn't grow with depth. (WebGPU guarantees only 8 storage buffers per shader stage.) +- **Every kernel can be repeated.** The fixup's `emit` writes each head's whole value to a slot of its own, never back into the head value it reads, so running any kernel twice gives the same result. (`benchmark()` repeats kernels in place, and a first version that added into the head in place got wrong answers there.) - **Every tile is emitted exactly once, with its whole value.** So `emit` can be any WGSL: it is never asked to combine partial results, and no atomics are needed. ### Fit to WebGPU diff --git a/examples/regression.mjs b/examples/regression.mjs index 6396168..a04aabb 100644 --- a/examples/regression.mjs +++ b/examples/regression.mjs @@ -1149,7 +1149,7 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} /* SegmentedReduce: each tile's reduction of data[atom] (u32 add or i32 * max), equal to the CPU's; one primitive over several shapes in turn * (tiles and atoms change between executes) */ - async function runSegmentedReduceTest({ shapes, tiles, binop, options = {}, fixupLevels }) { + async function runSegmentedReduceTest({ shapes, tiles, binop, options = {}, fixupLevels, trials = 1 }) { const T = binop.datatype; const TypedArray = T === "i32" ? Int32Array : Uint32Array; const primitive = new SegmentedReduce({ @@ -1169,7 +1169,7 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} const data = TypedArray.from({ length: atoms }, (_, a) => (T === "i32" ? (srHash(a) | 0) >> 8 : srHash(a) & 0xffff)); const buffers = { offsets: srBuffer(offsets), data: srBuffer(data), outputBuffer: srBuffer(new TypedArray(lengths.length)) }; Object.assign(primitive, { tiles: lengths.length, atoms }); - await primitive.execute({ ...buffers }); + await primitive.execute({ ...buffers, trials }); const out = await srRead(buffers.outputBuffer, TypedArray, lengths.length); const expected = segmentedReduceReference({ offsets, tiles: lengths.length, binop, atomValue: (a) => data[a] }); const t = expected.findIndex((v, i) => out[i] !== (T === "i32" ? v | 0 : v >>> 0)); @@ -1186,7 +1186,7 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} } /* SpMV: y = A x within f32 rounding of the CPU's f64 sums */ - async function runSpMVTest({ shape, rows, cols, options = {} }) { + async function runSpMVTest({ shape, rows, cols, options = {}, trials = 1 }) { const lengths = srShapes[shape](rows); const offsets = srOffsets(lengths); const nnz = offsets[lengths.length]; @@ -1202,7 +1202,7 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} y: srBuffer(new Float32Array(lengths.length)), }; const primitive = new SpMV({ device, rows: lengths.length, nnz, ...options }); - await primitive.execute({ ...buffers }); + await primitive.execute({ ...buffers, trials }); const y = await srRead(buffers.y, Float32Array, lengths.length); const expected = spmvReference({ offsets, indices, values, x }); let errors = ""; @@ -1585,6 +1585,16 @@ ${randomBitsWGSL(this.rng)} fixupLevels: 3, }) ); + /* every kernel repeated in place (as benchmark() does) gives the same result */ + await runTest("segred SegmentedReduce u32 add, mergepath, kernels repeated (trials: 3)", () => + runSegmentedReduceTest({ + shapes: ["skewed", "power law", "huge tile"], + tiles: 30011, + binop: new BinOpAdd({ datatype: "u32" }), + options: { schedule: "mergepath", itemsPerThread: 1 }, + trials: 3, + }) + ); for (const schedule of ["thread", "mergepath"]) { await runTest(`segred SegmentedReduce i32 max, ${schedule}`, () => runSegmentedReduceTest({ shapes: ["skewed", "all empty", "power law"], tiles: 5003, binop: new BinOpMax({ datatype: "i32" }), options: { schedule } }) @@ -1594,6 +1604,9 @@ ${randomBitsWGSL(this.rng)} ); for (const shape of ["uniform", "power law", "skewed"]) { await runTest(`segred SpMV ${shape}, ${schedule}`, () => runSpMVTest({ shape, rows: 20000, cols: 5000, options: { schedule } })); + await runTest(`segred SpMV ${shape}, ${schedule}, trials: 2`, () => + runSpMVTest({ shape, rows: 20000, cols: 5000, options: { schedule }, trials: 2 }) + ); } } } diff --git a/src/segmentedreduce.mjs b/src/segmentedreduce.mjs index 82d0366..04aef4a 100644 --- a/src/segmentedreduce.mjs +++ b/src/segmentedreduce.mjs @@ -225,7 +225,7 @@ fn heads_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonu let w = gid; if (w == 0u || w >= sr.workgroups) { return; } if (${this.wgTileName}[sr.workgroups + w] == 1u) { - emitTile(${this.wgTileName}[w - 1u], ${this.wgValueName}[sr.workgroups + w]); + emitTile(${this.wgTileName}[w - 1u], ${this.wgValueName}[2u * sr.workgroups + w]); } }`; }; @@ -235,7 +235,8 @@ fn heads_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonu * the run (the same carried tile) that ends at workgroup w - 1, if w ends * that tile (w's head); runs start where the carried tiles, which are * nondecreasing, first equal w - 1's. Its emit adds the run's carries - * before w's head value, in place. (Tile W - 1 also takes the last + * before w's head value, into a third slot (never in place: a repeated + * kernel would add them twice). (Tile W - 1 also takes the last * workgroup's carry, past the last tile: the identity.) It reads and writes * only this level's carry buffers, so every level binds the same few. */ fixupPrimitive() { @@ -271,7 +272,7 @@ fn runStart(t: u32) -> u32 { }`, tileStart: (t) => `runStart(${t})`, atom: (a) => `${V}[${a}]`, - emit: (w, v) => `if (${w} > 0u && ${C}[sr.tiles + ${w}] == 1u) { ${head(w)} = binop(${v}, ${head(w)}); }`, + emit: (w, v) => `if (${w} > 0u && ${C}[sr.tiles + ${w}] == 1u) { ${V}[2u * sr.tiles + ${w}] = binop(${v}, ${head(w)}); }`, }); return this.fixup; } @@ -302,9 +303,11 @@ fn runStart(t: u32) -> u32 { ); } else { actions.push( - /* per workgroup: carried tile and value, then head flag and value */ + /* per workgroup: carried tile, then head flag; carried value, then + * head value, then the head's whole value (written by the fixup, not + * in place, so that every kernel can be repeated: trials) */ new AllocateBuffer({ label: this.wgTileName, size: Math.max(16, 8 * workgroups), clearBufferOnReuse: false }), - new AllocateBuffer({ label: this.wgValueName, size: Math.max(16, 8 * workgroups), clearBufferOnReuse: false }), + new AllocateBuffer({ label: this.wgValueName, size: Math.max(16, 12 * workgroups), clearBufferOnReuse: false }), new Kernel({ kernel: this.mergepathWGSL, entryPoint: "mergepath_main", label: `${label} path`, buffers, getDispatchGeometry: dispatch(workgroups) }) ); /* one workgroup emits every tile itself; more need the fixup */ From b1cdcf6de1719d6dbc5a505c382a7c7ea3d770ea Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 18:25:25 -0800 Subject: [PATCH 05/10] docs: SpMV measured on an M1 Max Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- docs/segmented-reduce.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/segmented-reduce.md b/docs/segmented-reduce.md index 33c89f6..482ddfe 100644 --- a/docs/segmented-reduce.md +++ b/docs/segmented-reduce.md @@ -133,6 +133,21 @@ The regression tests (`segred`) check `SegmentedReduce` (u32 add, i32 max) and ` `misc/benchmark_spmv.mjs` times both schedules on generated matrices (`uniform`, `rmat` scale-free graphs, `denserow` with one row holding half the nonzeros) and on Matrix Market files (`--mtx=`, e.g. from the [SuiteSparse collection](https://sparse.tamu.edu)), checks every result, and reports GFLOP/s and bandwidth against a copy kernel's. +## Measured performance + +`node misc/benchmark_spmv.mjs uniform,rmat,denserow 18,20,22 --items=2,4,7`, on an Apple M1 Max (355 GB/s copy bandwidth), merge-path with 4 steps per invocation (the fastest of 2, 4 and 7 on every matrix). Every result was checked against a CPU SpMV. + +| Matrix, 2^22 rows, 67M nonzeros | Longest row | Thread-mapped | Merge-path | Merge-path, of which fixup | Speedup | +|---|---|---|---|---|---| +| uniform (16 per row) | 16 | 16.7 ms (8.1 GFLOP/s) | 7.7 ms (17.5 GFLOP/s) | 19 µs | 2.2× | +| R-MAT (scale-free, 52% of rows empty) | 160,302 | 83.2 ms (1.6 GFLOP/s) | 7.0 ms (19.2 GFLOP/s) | 17 µs | 11.9× | +| one row with half the nonzeros | 33,554,440 | 15,398 ms (0.01 GFLOP/s) | 7.6 ms (17.7 GFLOP/s) | 14 µs | 2,027× | + +- **Merge-path's time does not depend on the matrix's shape.** It runs at 17–21 GFLOP/s on all three matrices at every size from 2^18 to 2^22 rows. Thread-mapped runs at the speed of its longest row. +- **Merge-path wins on the balanced matrix too.** Neighboring invocations' loads are `itemsPerThread` elements apart, against a row (16 elements) apart for thread-mapped. Of 2, 4 and 7 steps, 4 is fastest: more steps spread a subgroup's loads further, fewer cost more searches and carries. +- **The fixup is almost free.** It is 0.2–0.3% of the time, even with one row spanning about 33,000 workgroups. The first version's walk-back fixup took 7,356 µs on that matrix, half of a 14.9 ms SpMV. +- **Large matrices are bound by the random reads of x.** Each nonzero reads x at a random column. On a V100 and an A100, the random-read cost that `misc/device_rates.mjs` measures at 2^22 (0.107 and 0.052 ns per read) times 67M nonzeros is 7.2 and 3.5 ms, against 7.6 and 3.4 ms measured for merge-path SpMV on the uniform matrix. + ## Possible improvements not made - **Coalesced loads.** Each invocation reads its atoms consecutively, so a subgroup's loads are strided by `itemsPerThread`. CUB stages a tile's atoms through workgroup memory first. From cd0b5ed2f8cb6fdfcba2c2e110ffeffb58565e55 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 18:26:32 -0800 Subject: [PATCH 06/10] SegmentedReduce: label fixup kernels by level; docs: V100 and A100 results The fixup levels' kernels were labeled like the main kernel, so the benchmark's fixup time left them out. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- docs/segmented-reduce.md | 28 +++++++++++++++++----------- src/segmentedreduce.mjs | 2 +- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/docs/segmented-reduce.md b/docs/segmented-reduce.md index 482ddfe..0e01ed7 100644 --- a/docs/segmented-reduce.md +++ b/docs/segmented-reduce.md @@ -135,22 +135,28 @@ The regression tests (`segred`) check `SegmentedReduce` (u32 add, i32 max) and ` ## Measured performance -`node misc/benchmark_spmv.mjs uniform,rmat,denserow 18,20,22 --items=2,4,7`, on an Apple M1 Max (355 GB/s copy bandwidth), merge-path with 4 steps per invocation (the fastest of 2, 4 and 7 on every matrix). Every result was checked against a CPU SpMV. +`node misc/benchmark_spmv.mjs uniform,rmat,denserow 18,20,22 --items=2,4,7`, on an Apple M1 Max (355 GB/s copy bandwidth), an NVIDIA V100 (809 GB/s) and an A100 (PCIe, 1,351 GB/s), merge-path with 4 steps per invocation. Every result was checked against a CPU SpMV. -| Matrix, 2^22 rows, 67M nonzeros | Longest row | Thread-mapped | Merge-path | Merge-path, of which fixup | Speedup | -|---|---|---|---|---|---| -| uniform (16 per row) | 16 | 16.7 ms (8.1 GFLOP/s) | 7.7 ms (17.5 GFLOP/s) | 19 µs | 2.2× | -| R-MAT (scale-free, 52% of rows empty) | 160,302 | 83.2 ms (1.6 GFLOP/s) | 7.0 ms (19.2 GFLOP/s) | 17 µs | 11.9× | -| one row with half the nonzeros | 33,554,440 | 15,398 ms (0.01 GFLOP/s) | 7.6 ms (17.7 GFLOP/s) | 14 µs | 2,027× | +GFLOP/s (2 per nonzero) at 2^22 rows and 67M nonzeros, thread-mapped → merge-path: -- **Merge-path's time does not depend on the matrix's shape.** It runs at 17–21 GFLOP/s on all three matrices at every size from 2^18 to 2^22 rows. Thread-mapped runs at the speed of its longest row. -- **Merge-path wins on the balanced matrix too.** Neighboring invocations' loads are `itemsPerThread` elements apart, against a row (16 elements) apart for thread-mapped. Of 2, 4 and 7 steps, 4 is fastest: more steps spread a subgroup's loads further, fewer cost more searches and carries. -- **The fixup is almost free.** It is 0.2–0.3% of the time, even with one row spanning about 33,000 workgroups. The first version's walk-back fixup took 7,356 µs on that matrix, half of a 14.9 ms SpMV. -- **Large matrices are bound by the random reads of x.** Each nonzero reads x at a random column. On a V100 and an A100, the random-read cost that `misc/device_rates.mjs` measures at 2^22 (0.107 and 0.052 ns per read) times 67M nonzeros is 7.2 and 3.5 ms, against 7.6 and 3.4 ms measured for merge-path SpMV on the uniform matrix. +| Matrix | Longest row | M1 Max | V100 | A100 | +|---|---|---|---|---| +| uniform (16 per row) | 16 | 8.1 → 17.5 | 8.8 → 17.5 | 18.6 → 38.8 | +| R-MAT (scale-free, 52% of rows empty) | 160,302 | 1.6 → 19.2 | 2.3 → 18.4 | 2.6 → 39.9 | +| one row with half the nonzeros | 33,554,440 | 0.01 → 17.7 | 0.01 → 17.5 | 0.01 → 39.0 | + +That is 2.0–2.2× faster on the uniform matrix, 8–15× on R-MAT, and 1,500–2,800× with one dense row (thread-mapped takes 10–15 s there; merge-path 3.4–7.7 ms). + +- **Merge-path's speed does not depend on the matrix's shape.** On each GPU it runs within 10% of the same GFLOP/s on all three matrices, at every size from 2^18 to 2^22 rows. Thread-mapped runs at the speed of its longest row. +- **Merge-path also wins on the balanced matrix** on all three GPUs at 2^22. Neighboring invocations' loads are `itemsPerThread` elements apart, against a row (16 elements) apart for thread-mapped. Of 2, 4 and 7 steps, 4 was fastest or within 2% of it everywhere. (At 2^18 and 2^20 on the A100, the uniform matrix runs the same either way.) +- **The recursive fixup removed the cost of long rows.** With the first version's walk-back fixup, the dense-row matrix took 9.2 ms on the A100 (5.8 ms of it fixup) and 14.9 ms on the M1 Max (7.4 ms); now 3.4 and 7.6 ms, the same as the uniform matrix. +- **The fixup costs a few small kernels: about 30 µs.** Each level is a short merge-path kernel and a heads kernel. At 2^18 that is 15–25% of a 0.12–0.43 ms SpMV, where the walk-back fixup cost 3–8 µs on matrices without long rows; at 2^22 it is under 1%. +- **Large matrices are bound by the random reads of x.** Each nonzero reads x at a random column. The random-read cost that `misc/device_rates.mjs` measures at 2^22 (V100 0.107 ns, A100 0.052 ns per read) times 67M nonzeros is 7.2 and 3.5 ms, against 7.7 and 3.5 ms measured for merge-path on the uniform matrix. So at this size, merge-path SpMV runs as fast as its reads of x allow, and loading the atoms better would not help. ## Possible improvements not made -- **Coalesced loads.** Each invocation reads its atoms consecutively, so a subgroup's loads are strided by `itemsPerThread`. CUB stages a tile's atoms through workgroup memory first. +- **Coalesced loads.** Each invocation reads its atoms consecutively, so a subgroup's loads are strided by `itemsPerThread`. CUB stages a tile's atoms through workgroup memory first. This could help only where x fits in cache (small matrices): large ones are bound by the reads of x. +- **A cheaper fixup for small problems.** The recursive fixup's few kernels cost about 30 µs whatever the matrix. A single-kernel fixup when no tile spans many workgroups would save most of that on small matrices. - **One search per workgroup.** Each invocation searches the global offsets. CUB searches once per workgroup, then within a tile of row ends in workgroup memory. - **Atoms from the GPU.** The caller passes `atoms` (offsets[tiles]) to size the grid. An indirect dispatch could read it on the GPU instead, for offsets produced by an earlier kernel. - **More schedules.** Workgroup-mapped and subgroup-mapped schedules, as in the paper. diff --git a/src/segmentedreduce.mjs b/src/segmentedreduce.mjs index 04aef4a..5b23059 100644 --- a/src/segmentedreduce.mjs +++ b/src/segmentedreduce.mjs @@ -295,7 +295,7 @@ fn runStart(t: u32) -> u32 { { name: "sr", buffer: "srParams", uniform: this.params }, ]; const dispatch = (count) => () => this.getSimpleDispatchGeometry({ workgroupCount: count }); - const label = `SegmentedReduce (${this.schedule})`; + const label = this.level ? `SegmentedReduce fixup level ${this.level}` : `SegmentedReduce (${this.schedule})`; const actions = [this.params.allocate("srParams", { tiles: this.tiles, atoms: this.atoms, workgroups })]; if (this.schedule === "thread") { actions.push( From 610d518fc3fbc165536e63a53248fcb51e5082e5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 18:26:43 -0800 Subject: [PATCH 07/10] docs: ranges Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- docs/segmented-reduce.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/segmented-reduce.md b/docs/segmented-reduce.md index 0e01ed7..a5d71ad 100644 --- a/docs/segmented-reduce.md +++ b/docs/segmented-reduce.md @@ -145,12 +145,12 @@ GFLOP/s (2 per nonzero) at 2^22 rows and 67M nonzeros, thread-mapped → merge-p | R-MAT (scale-free, 52% of rows empty) | 160,302 | 1.6 → 19.2 | 2.3 → 18.4 | 2.6 → 39.9 | | one row with half the nonzeros | 33,554,440 | 0.01 → 17.7 | 0.01 → 17.5 | 0.01 → 39.0 | -That is 2.0–2.2× faster on the uniform matrix, 8–15× on R-MAT, and 1,500–2,800× with one dense row (thread-mapped takes 10–15 s there; merge-path 3.4–7.7 ms). +That is 2.0–2.2× faster on the uniform matrix, 8–15× on R-MAT, and 1,500–2,800× with one dense row (thread-mapped takes 9.6–15.4 s there; merge-path 3.4–7.7 ms). - **Merge-path's speed does not depend on the matrix's shape.** On each GPU it runs within 10% of the same GFLOP/s on all three matrices, at every size from 2^18 to 2^22 rows. Thread-mapped runs at the speed of its longest row. - **Merge-path also wins on the balanced matrix** on all three GPUs at 2^22. Neighboring invocations' loads are `itemsPerThread` elements apart, against a row (16 elements) apart for thread-mapped. Of 2, 4 and 7 steps, 4 was fastest or within 2% of it everywhere. (At 2^18 and 2^20 on the A100, the uniform matrix runs the same either way.) - **The recursive fixup removed the cost of long rows.** With the first version's walk-back fixup, the dense-row matrix took 9.2 ms on the A100 (5.8 ms of it fixup) and 14.9 ms on the M1 Max (7.4 ms); now 3.4 and 7.6 ms, the same as the uniform matrix. -- **The fixup costs a few small kernels: about 30 µs.** Each level is a short merge-path kernel and a heads kernel. At 2^18 that is 15–25% of a 0.12–0.43 ms SpMV, where the walk-back fixup cost 3–8 µs on matrices without long rows; at 2^22 it is under 1%. +- **The fixup costs a few small kernels: about 30 µs.** Each level is a short merge-path kernel and a heads kernel. At 2^18 that is 7–25% of a 0.12–0.43 ms SpMV, where the walk-back fixup cost 3–8 µs on matrices without long rows; at 2^22 it is under 1%. - **Large matrices are bound by the random reads of x.** Each nonzero reads x at a random column. The random-read cost that `misc/device_rates.mjs` measures at 2^22 (V100 0.107 ns, A100 0.052 ns per read) times 67M nonzeros is 7.2 and 3.5 ms, against 7.7 and 3.5 ms measured for merge-path on the uniform matrix. So at this size, merge-path SpMV runs as fast as its reads of x allow, and loading the atoms better would not help. ## Possible improvements not made From 11e88ac06aff3f5979f3706524946b75f16f6f99 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 18:45:58 -0800 Subject: [PATCH 08/10] SegmentedReduce: schedules are classes; add workgroup-mapped Each schedule (ThreadMapped, WorkgroupMapped, MergePath, in src/schedules.mjs) writes its kernels and says what runs; SegmentedReduce holds only the work (binop, tileStart, atom, emit) and helpers schedules use (kernel, allocate, child, run). Schedules call the work through binop/IDENTITY/tileStartOf/atomValue/emitTile and the sr uniform. A schedule is given as an instance or by a name in SegmentedReduce.schedules; adding one changes nothing in SegmentedReduce, as a test with a schedule defined in the test file shows. MergePath's fixup is a child SegmentedReduce. Adds WorkgroupMapped (the paper's block-mapped) to the tests and the benchmark. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- docs/segmented-reduce.md | 68 +++++++-- examples/regression.mjs | 34 ++++- index.d.ts | 54 ++++++- index.mjs | 1 + misc/benchmark_spmv.mjs | 22 +-- src/schedules.mjs | 318 +++++++++++++++++++++++++++++++++++++++ src/segmentedreduce.mjs | 316 ++++++++------------------------------ src/spmv.mjs | 3 +- 8 files changed, 535 insertions(+), 281 deletions(-) create mode 100644 src/schedules.mjs diff --git a/docs/segmented-reduce.md b/docs/segmented-reduce.md index a5d71ad..af0d70e 100644 --- a/docs/segmented-reduce.md +++ b/docs/segmented-reduce.md @@ -60,14 +60,64 @@ This is code generation by string, and it is deliberately so: JavaScript compose - `atom` (required), `emit`, `buffers`: above; - `tileStart: (t) => expression`: tile t's first atom, for t from 0 to `tiles`, in place of an `offsets` buffer (the paper's iterators over a sparse format); - `wgsl`: more WGSL declarations, such as functions that `atom`, `emit` or `tileStart` call; -- `schedule`: `"mergepath"` (default) or `"thread"`; -- `itemsPerThread`: merge-path steps per invocation (default 4, the fastest on an M1 Max, a V100 and an A100). +- `schedule`: a `Schedule` (`new MergePath({ itemsPerThread: 7 })`), or the name of one: `"mergepath"` (default), `"thread"` or `"workgroup"`; +- `itemsPerThread`: passed to a schedule given by name; merge-path's steps per invocation (default 4, the fastest on an M1 Max, a V100 and an A100). ## Schedules -**`thread`** gives each invocation one tile at a time, striding over the tiles. It is balanced only when tiles are about the same size. One long tile holds up the whole kernel, because a single invocation walks all of it. +A schedule is a class of its own (`src/schedules.mjs`). `SegmentedReduce` holds the work and asks its schedule for the kernels; it knows no schedule by more than its name. There are three: -**`mergepath`** (Merrill and Garland, SC 2016) walks the merge of two lists: the tiles' ends and the atom indices. Each step of the path either consumes an atom or ends a tile, so there are tiles + atoms steps, and every invocation takes `itemsPerThread` of them, whatever the tile sizes. A binary search over the offsets finds where each invocation's steps start. +**`ThreadMapped`** (`"thread"`) gives each invocation one tile at a time, striding over the tiles. It is balanced only when tiles are about the same size. One long tile holds up the whole kernel, because a single invocation walks all of it. + +**`WorkgroupMapped`** (`"workgroup"`, the paper's block-mapped) gives each workgroup one tile at a time. Its invocations stride over the tile's atoms, then reduce their partial values in workgroup memory. It suits long tiles of about the same size. A short tile leaves most of the workgroup idle, and one very long tile still holds up the kernel, though 256 times less than with thread-mapped. + +**`MergePath`** (`"mergepath"`; Merrill and Garland, SC 2016) walks the merge of two lists: the tiles' ends and the atom indices. Each step of the path either consumes an atom or ends a tile, so there are tiles + atoms steps, and every invocation takes `itemsPerThread` of them, whatever the tile sizes. A binary search over the offsets finds where each invocation's steps start. + +### Writing a schedule + +A schedule writes the loops and never sees the work. It calls the work through WGSL that `SegmentedReduce` puts in the same module: + +| Name | What it is | +|---|---| +| `binop(a, b)`, `IDENTITY` | the reduction | +| `tileStartOf(t)` | tile t's first atom, for t from 0 to `sr.tiles` | +| `atomValue(atom, tile)` | an atom's value | +| `emitTile(tile, value)` | a tile's result, to be called once per tile | +| `sr.tiles`, `sr.atoms`, `sr.workgroups` | uniform parameters | + +A schedule is a class with four methods: + +- `workgroups({ tiles, atoms })`: how many workgroups its main kernel needs; +- `wgsl(sr)`: the WGSL of its kernels (256 invocations per workgroup); +- `buffers(sr)` (optional): buffers its kernels bind beyond the caller's; +- `actions(sr, { workgroups })`: what runs in one execute, made with `sr.kernel({ entryPoint, workgroups })`, `sr.allocate()` for scratch buffers, and `sr.child()` and `sr.run()` for another `SegmentedReduce` it runs (as merge-path runs its fixup). + +The regression tests define one outside the library, in which one invocation does every tile: + +```js +class SerialSchedule extends Schedule { + name = "serial"; + workgroups() { return 1; } + wgsl(sr) { + return /* wgsl */ ` +@compute @workgroup_size(256) +fn serial_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${sr.fnDeclarations.computeLinearizedGridParametersSplit()} + if (gid != 0u) { return; } + for (var tile = 0u; tile < sr.tiles; tile++) { + var acc = IDENTITY; + for (var atom = tileStartOf(tile); atom < tileStartOf(tile + 1u); atom++) { acc = binop(acc, atomValue(atom, tile)); } + emitTile(tile, acc); + } +}`; + } + actions(sr, { workgroups }) { return [sr.kernel({ entryPoint: "serial_main", workgroups })]; } +} + +new SegmentedReduce({ ..., schedule: new SerialSchedule() }); +``` + +To use one by name, register it: `SegmentedReduce.schedules.serial = SerialSchedule`. ### How merge-path puts split tiles back together @@ -123,15 +173,15 @@ The last tile also takes the last workgroup's carry, which is past the end of th ### Fit to WebGPU -- **Barriers in uniform control flow.** WGSL requires `workgroupBarrier()` where control flow is provably uniform across the workgroup. Merge-path's per-invocation loops have no barriers; its only barriers are in the workgroup scan, whose loop bounds are constants. A block-mapped schedule (a workgroup per tile, scanning atom counts per tile) would need its loop bounds made uniform with `workgroupUniformLoad`. +- **Barriers in uniform control flow.** WGSL requires `workgroupBarrier()` where control flow is provably uniform across the workgroup. Merge-path's per-invocation loops have no barriers; its only barriers are in the workgroup scan, whose loop bounds are constants. Workgroup-mapped has barriers inside its loop over tiles, which is allowed because that loop's index comes from the workgroup id and its bound is a uniform parameter; the atom loop, whose bounds come from storage, has none. - **No float atomics.** CUDA merge-path SpMV can finish split rows with `atomicAdd` on floats. WebGPU atomics are 32-bit integers only, so the fixup combines split tiles instead, which also keeps the result deterministic. -- **No subgroups needed.** Both schedules use only workgroup memory and barriers, so they run the same on every device, including those without subgroups. +- **No subgroups needed.** All three schedules use only workgroup memory and barriers, so they run the same on every device, including those without subgroups. ## Testing and measuring -The regression tests (`segred`) check `SegmentedReduce` (u32 add, i32 max) and `SpMV` against the CPU on several shapes: uniform tiles, a power law, a few tiles spanning many workgroups, one tile, all tiles empty. They cover both schedules, 1, 7 and 16 items per invocation, three levels of fixup (one tile of 5 million atoms, one item per invocation), and change `tiles` and `atoms` between executes. +The regression tests (`segred`) check `SegmentedReduce` (u32 add, i32 max) and `SpMV` against the CPU on several shapes: uniform tiles, a power law, a few tiles spanning many workgroups, one tile, all tiles empty. They cover all three schedules and one defined outside the library, 1, 7 and 16 items per invocation, kernels repeated in place, three levels of fixup (one tile of 5 million atoms, one item per invocation), and change `tiles` and `atoms` between executes. -`misc/benchmark_spmv.mjs` times both schedules on generated matrices (`uniform`, `rmat` scale-free graphs, `denserow` with one row holding half the nonzeros) and on Matrix Market files (`--mtx=`, e.g. from the [SuiteSparse collection](https://sparse.tamu.edu)), checks every result, and reports GFLOP/s and bandwidth against a copy kernel's. +`misc/benchmark_spmv.mjs` times the three schedules on generated matrices (`uniform`, `rmat` scale-free graphs, `denserow` with one row holding half the nonzeros) and on Matrix Market files (`--mtx=`, e.g. from the [SuiteSparse collection](https://sparse.tamu.edu)), checks every result, and reports GFLOP/s and bandwidth against a copy kernel's. ## Measured performance @@ -159,4 +209,4 @@ That is 2.0–2.2× faster on the uniform matrix, 8–15× on R-MAT, and 1,500 - **A cheaper fixup for small problems.** The recursive fixup's few kernels cost about 30 µs whatever the matrix. A single-kernel fixup when no tile spans many workgroups would save most of that on small matrices. - **One search per workgroup.** Each invocation searches the global offsets. CUB searches once per workgroup, then within a tile of row ends in workgroup memory. - **Atoms from the GPU.** The caller passes `atoms` (offsets[tiles]) to size the grid. An indirect dispatch could read it on the GPU instead, for offsets produced by an earlier kernel. -- **More schedules.** Workgroup-mapped and subgroup-mapped schedules, as in the paper. +- **More schedules.** A subgroup-mapped schedule, as in the paper: it would depend on the device's subgroup size, which WebGPU reports only as a range. diff --git a/examples/regression.mjs b/examples/regression.mjs index a04aabb..8c793f9 100644 --- a/examples/regression.mjs +++ b/examples/regression.mjs @@ -22,6 +22,7 @@ import { ShunPerm } from "../src/shunperm.mjs"; import { SuffixArray, suffixArrayReference, checkSuffixArray } from "../src/suffixarray.mjs"; import { Histogram } from "../src/histogram.mjs"; import { SegmentedReduce, segmentedReduceReference } from "../src/segmentedreduce.mjs"; +import { Schedule } from "../src/schedules.mjs"; import { SpMV, spmvReference } from "../src/spmv.mjs"; import { EvenBins, RangeBins, RadixBins } from "../src/bins.mjs"; @@ -1176,7 +1177,7 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} if (t >= 0) errors += `\n${shape}: tile ${t} (${lengths[t]} atoms) = ${out[t]}, expected ${expected[t]}`; if (fixupLevels !== undefined) { let levels = 0; - for (let p = primitive.fixup; p; p = p.fixup) levels++; + for (let p = primitive.children.fixup; p; p = p.children.fixup) levels++; if (levels !== fixupLevels) errors += `\n${shape}: ${levels} levels of fixup, expected ${fixupLevels}`; } for (const b of Object.values(buffers)) b.destroy(); @@ -1185,6 +1186,31 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} return errors; } + /* A schedule written outside SegmentedReduce, to check that one can be: + * one invocation does every tile */ + class SerialSchedule extends Schedule { + name = "serial"; + workgroups() { + return 1; + } + wgsl(sr) { + return /* wgsl */ ` +@compute @workgroup_size(256) +fn serial_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${sr.fnDeclarations.computeLinearizedGridParametersSplit()} + if (gid != 0u) { return; } + for (var tile = 0u; tile < sr.tiles; tile++) { + var acc = IDENTITY; + for (var atom = tileStartOf(tile); atom < tileStartOf(tile + 1u); atom++) { acc = binop(acc, atomValue(atom, tile)); } + emitTile(tile, acc); + } +}`; + } + actions(sr, { workgroups }) { + return [sr.kernel({ entryPoint: "serial_main", workgroups })]; + } + } + /* SpMV: y = A x within f32 rounding of the CPU's f64 sums */ async function runSpMVTest({ shape, rows, cols, options = {}, trials = 1 }) { const lengths = srShapes[shape](rows); @@ -1568,6 +1594,7 @@ ${randomBitsWGSL(this.rng)} const allShapes = ["uniform", "skewed", "one long tile", "all empty", "power law"]; for (const [label, schedule, extra] of [ ["thread", "thread", {}], + ["workgroup", "workgroup", {}], ["mergepath", "mergepath", {}], ["mergepath, 1 item per thread", "mergepath", { itemsPerThread: 1 }], ["mergepath, 16 items per thread", "mergepath", { itemsPerThread: 16 }], @@ -1595,7 +1622,10 @@ ${randomBitsWGSL(this.rng)} trials: 3, }) ); - for (const schedule of ["thread", "mergepath"]) { + await runTest("segred SegmentedReduce u32 add, a schedule defined outside (serial)", () => + runSegmentedReduceTest({ shapes: ["skewed", "all empty"], tiles: 2003, binop: new BinOpAdd({ datatype: "u32" }), options: { schedule: new SerialSchedule() } }) + ); + for (const schedule of ["thread", "workgroup", "mergepath"]) { await runTest(`segred SegmentedReduce i32 max, ${schedule}`, () => runSegmentedReduceTest({ shapes: ["skewed", "all empty", "power law"], tiles: 5003, binop: new BinOpMax({ datatype: "i32" }), options: { schedule } }) ); diff --git a/index.d.ts b/index.d.ts index 8c59da1..a627da8 100644 --- a/index.d.ts +++ b/index.d.ts @@ -498,18 +498,37 @@ export interface SegmentedReduceOptions { emit?: (tile: string, value: string) => string; /** the buffers atom and emit use, bound by name in execute() (access "read" unless given) */ buffers?: { name: string; type: string; access?: "read" | "read_write" }[]; - /** "mergepath" (default) or "thread" */ - schedule?: "mergepath" | "thread"; - /** mergepath: path steps per invocation (default 7) */ + /** a Schedule, or a name in SegmentedReduce.schedules: "mergepath" (default), "thread", "workgroup" */ + schedule?: Schedule | string; + /** passed to a schedule given by name (merge-path: steps per invocation, default 4) */ itemsPerThread?: number; + /** tile t's first atom (t in [0, tiles]) as a WGSL expression, in place of an offsets buffer */ + tileStart?: (t: string) => string; + /** more WGSL declarations (e.g. functions atom, emit or tileStart call) */ + wgsl?: string; adapter?: GPUAdapter; [key: string]: any; } export class SegmentedReduce extends BasePrimitive { + /** schedules by name */ + static schedules: { [name: string]: new (options?: { [key: string]: any }) => Schedule }; tiles: number; atoms: number; - schedule: "mergepath" | "thread"; + schedule: Schedule; + datatype: string; + /** nesting depth (0 for the caller's; schedules' children are deeper) */ + level: number; + /** child primitives a schedule made (e.g. merge-path's fixup) */ + children: { [key: string]: SegmentedReduce }; constructor(args: SegmentedReduceOptions); + /** for schedules: a Kernel running one entry point of the schedule's WGSL */ + kernel(args: { entryPoint: string; workgroups: number; label?: string }): Kernel; + /** for schedules: a scratch buffer */ + allocate(args: { label: string; size: number }): AllocateBuffer; + /** for schedules: a child SegmentedReduce, made once and kept under key */ + child(key: string, options: { [key: string]: any }): SegmentedReduce; + /** for schedules: run a child primitive, binding these buffers by the same names */ + run(primitive: BasePrimitive, names: string[]): ExecutePrimitive; /** offsets (tiles + 1 u32s), the buffers named in options.buffers, and * outputBuffer when there is no emit */ execute(args: { offsets: GPUBuffer; [buffer: string]: any }): Promise; @@ -522,6 +541,31 @@ export function segmentedReduceReference(args: { atomValue: (atom: number, tile: number) => number; }): number[]; +/** A load-balancing schedule for SegmentedReduce: writes the kernels; calls + * the work through binop, IDENTITY, tileStartOf, atomValue, emitTile and the + * uniform sr (tiles, atoms, workgroups). schedules.mjs */ +export class Schedule { + name: string; + constructor(options?: { [key: string]: any }); + /** workgroups of the main kernel */ + workgroups(args: { tiles: number; atoms: number }): number; + /** WGSL of the schedule's kernels */ + wgsl(sr: SegmentedReduce): string; + /** buffers its kernels bind beyond the caller's */ + buffers(sr: SegmentedReduce): KernelBuffer[]; + /** the actions of one execute */ + actions(sr: SegmentedReduce, args: { workgroups: number }): any[]; +} +/** one invocation per tile */ +export class ThreadMapped extends Schedule {} +/** one workgroup per tile (the paper's block-mapped) */ +export class WorkgroupMapped extends Schedule {} +/** merge path over tile ends and atoms, with a recursive SegmentedReduce fixup */ +export class MergePath extends Schedule { + itemsPerThread: number; + constructor(options?: { itemsPerThread?: number }); +} + /** y = A x, A in CSR form (f32): a SegmentedReduce over A's rows. docs/segmented-reduce.md */ export class SpMV extends SegmentedReduce { rows: number; @@ -530,7 +574,7 @@ export class SpMV extends SegmentedReduce { device: GPUDevice; rows: number; nnz: number; - schedule?: "mergepath" | "thread"; + schedule?: Schedule | string; itemsPerThread?: number; adapter?: GPUAdapter; [key: string]: any; diff --git a/index.mjs b/index.mjs index f1afb94..f0c1f3f 100644 --- a/index.mjs +++ b/index.mjs @@ -26,6 +26,7 @@ export { suffixArrayPhase, } from "./src/suffixarray.mjs"; export { SegmentedReduce, segmentedReduceReference } from "./src/segmentedreduce.mjs"; +export { Schedule, ThreadMapped, WorkgroupMapped, MergePath } from "./src/schedules.mjs"; export { SpMV, spmvReference } from "./src/spmv.mjs"; export { randomBitsWGSL, diff --git a/misc/benchmark_spmv.mjs b/misc/benchmark_spmv.mjs index a984ac8..17f1b02 100644 --- a/misc/benchmark_spmv.mjs +++ b/misc/benchmark_spmv.mjs @@ -1,5 +1,6 @@ -/* Benchmark SpMV (SegmentedReduce) schedules in Node (Dawn): thread-mapped - * against merge-path, on balanced and unbalanced sparse matrices. +/* Benchmark SpMV (SegmentedReduce) schedules in Node (Dawn): thread-mapped, + * workgroup-mapped and merge-path, on balanced and unbalanced sparse + * matrices. * * node misc/benchmark_spmv.mjs [matrices] [log2 sizes] [--mtx=file.mtx,...] [--items=K,...] * node misc/benchmark_spmv.mjs uniform,rmat,denserow 18,20,22 @@ -177,9 +178,12 @@ async function benchmark(name, m) { const expected = spmvReference({ offsets, indices, values, x }); const bytes = 12 * nnz + 8 * rows; const results = []; - const runs = [["thread", {}], ...itemsList.map((K) => [`mergepath${itemsList.length > 1 ? ` K=${K}` : ""}`, { itemsPerThread: K }])]; - for (const [label, extra] of runs) { - const schedule = label.startsWith("thread") ? "thread" : "mergepath"; + const runs = [ + ["thread", "thread", {}], + ["workgroup", "workgroup", {}], + ...itemsList.map((K) => [`mergepath${itemsList.length > 1 ? ` K=${K}` : ""}`, "mergepath", { itemsPerThread: K }]), + ]; + for (const [label, schedule, extra] of runs) { const spmv = new SpMV({ device, adapter, rows, nnz, schedule, ...extra }); device.queue.writeBuffer(buffers.y, 0, new Float32Array(rows)); const { ns, kernels } = await spmv.benchmark({ ...buffers, warmMs: 1000 }); @@ -192,7 +196,7 @@ async function benchmark(name, m) { } spmv.destroy(); const fixupNS = kernels.filter((k) => /fixup|heads/.test(k.label ?? "")).reduce((a, k) => a + k.ns, 0); - results.push({ label, ns }); + results.push({ schedule, ns }); console.log( ` ${label.padEnd(14)} ${(ns / 1e6).toFixed(3).padStart(9)} ms ${((2 * nnz) / ns).toFixed(2).padStart(7)} GFLOP/s ` + `${(bytes / ns).toFixed(1).padStart(7)} GB/s (${((100 * bytes) / ns / copyGBps).toFixed(0).padStart(3)}% of copy)` + @@ -200,9 +204,9 @@ async function benchmark(name, m) { (bad >= 0 ? ` WRONG: row ${bad} = ${y[bad]}, expected ${expected[bad]}` : "") ); } - const thread = results[0].ns; - const best = Math.min(...results.slice(1).map((r) => r.ns)); - console.log(` merge-path / thread: ${(thread / best).toFixed(2)}x faster\n`); + const best = (s) => Math.min(...results.filter((r) => r.schedule === s).map((r) => r.ns)); + const mergepath = best("mergepath"); + console.log(` merge-path is ${(best("thread") / mergepath).toFixed(2)}x thread-mapped's speed, ${(best("workgroup") / mergepath).toFixed(2)}x workgroup-mapped's\n`); for (const b of Object.values(buffers)) b.destroy(); } diff --git a/src/schedules.mjs b/src/schedules.mjs new file mode 100644 index 0000000..11bf754 --- /dev/null +++ b/src/schedules.mjs @@ -0,0 +1,318 @@ +/** + * Load-balancing schedules for SegmentedReduce: which invocation handles + * which atoms, separate from what the work is (Osama, Porumbescu and Owens, + * PPoPP 2023). A schedule knows nothing about the work: it writes the loops, + * and calls the work through WGSL that SegmentedReduce defines in the same + * module: + * + * binop(a, b), IDENTITY the reduction + * tileStartOf(t) -> u32 tile t's first atom, t in [0, sr.tiles] + * atomValue(atom, tile) an atom's value + * emitTile(tile, value) consume a tile's result (once per tile) + * sr.tiles, sr.atoms, sr.workgroups (uniform) + * WORKGROUP_SIZE invocations per workgroup; computeLinearizedGridParametersSplit + * + * A schedule is a class with: + * + * workgroups({ tiles, atoms }) workgroups of its main kernel + * wgsl(sr) WGSL of its kernels (entry points) + * buffers(sr) buffers its kernels bind beyond the + * caller's (optional) + * actions(sr, { workgroups }) the actions of one execute; sr.kernel() + * makes a Kernel of one entry point + * + * To add one, write the class and pass an instance (schedule: new Mine()) + * or register a name (SegmentedReduce.schedules.mine = Mine). + */ + +export const WORKGROUP_SIZE = 256; + +export class Schedule { + constructor(options = {}) { + Object.assign(this, options); + } + toString() { + return this.name; + } + buffers() { + return []; + } +} + +/** One invocation per tile, striding over the tiles. Balanced only when + * tiles are about the same size: one long tile holds up the kernel. */ +export class ThreadMapped extends Schedule { + name = "thread"; + workgroups({ tiles }) { + return Math.ceil(tiles / WORKGROUP_SIZE); + } + wgsl(sr) { + return /* wgsl */ ` +@compute @workgroup_size(${WORKGROUP_SIZE}) +fn thread_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${sr.fnDeclarations.computeLinearizedGridParametersSplit()} + for (var tile = gid; tile < sr.tiles; tile += totalThreadCount) { + var acc = IDENTITY; + let end = tileStartOf(tile + 1u); + for (var atom = tileStartOf(tile); atom < end; atom++) { acc = binop(acc, atomValue(atom, tile)); } + emitTile(tile, acc); + } +}`; + } + actions(sr, { workgroups }) { + return [sr.kernel({ entryPoint: "thread_main", workgroups })]; + } +} + +/** One workgroup per tile, striding over the tiles; its invocations stride + * over the tile's atoms, then reduce in workgroup memory. Balanced for long + * tiles of about the same size; a short tile leaves most of a workgroup + * idle. (The paper's block-mapped schedule.) */ +export class WorkgroupMapped extends Schedule { + name = "workgroup"; + workgroups({ tiles }) { + return Math.min(tiles, 65536); + } + wgsl(sr) { + const W = WORKGROUP_SIZE; + return /* wgsl */ ` +var partial: array<${sr.datatype}, ${W}>; +@compute @workgroup_size(${W}) +fn workgroup_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${sr.fnDeclarations.computeLinearizedGridParametersSplit()} + let lidx = builtinsNonuniform.lidx; + /* wgid, workgroupCount and sr.tiles are uniform, so the barriers below + * are in uniform control flow */ + for (var tile = wgid; tile < sr.tiles; tile += workgroupCount) { + var acc = IDENTITY; + let end = tileStartOf(tile + 1u); + for (var atom = tileStartOf(tile) + lidx; atom < end; atom += ${W}u) { acc = binop(acc, atomValue(atom, tile)); } + partial[lidx] = acc; + workgroupBarrier(); + /* tree reduction, keeping order: each step combines neighbors */ + for (var stride = 1u; stride < ${W}u; stride <<= 1u) { + var combined = IDENTITY; + let mine = lidx % (2u * stride) == 0u; + if (mine) { combined = binop(partial[lidx], partial[lidx + stride]); } + workgroupBarrier(); + if (mine) { partial[lidx] = combined; } + workgroupBarrier(); + } + if (lidx == 0u) { emitTile(tile, partial[0]); } + workgroupBarrier(); + } +}`; + } + actions(sr, { workgroups }) { + return [sr.kernel({ entryPoint: "workgroup_main", workgroups })]; + } +} + +/** The merge path of the tile ends and the atom indices (Merrill and + * Garland, SC 2016): tiles + atoms steps, split evenly, so every invocation + * does itemsPerThread steps (default 4) whatever the tile sizes. + * + * A tile split between invocations is combined by a segmented scan in the + * workgroup. Between workgroups, combining is itself a segmented reduce: each + * workgroup's carry (its partial value of the tile unfinished at its end) is + * an atom, and each workgroup is a tile, holding the run of carries in the + * tile it ends. So a second SegmentedReduce (merge-path too; recursively, + * until one workgroup suffices) combines them into the finishing workgroups' + * head values, however many workgroups a tile spans, and a last kernel emits + * the heads. Each tile is emitted exactly once, with its whole value, so emit + * may be any WGSL (it is never asked to combine). See + * docs/segmented-reduce.md for a worked example. */ +export class MergePath extends Schedule { + name = "mergepath"; + constructor(options = {}) { + super({ itemsPerThread: 4, ...options }); + if (!(Number.isInteger(this.itemsPerThread) && this.itemsPerThread >= 1)) { + throw new Error("MergePath: itemsPerThread must be a positive integer"); + } + } + workgroups({ tiles, atoms }) { + return Math.ceil((tiles + atoms) / (WORKGROUP_SIZE * this.itemsPerThread)); + } + /* per-level names, so that each level of fixup binds its parent's carries + * under different names from its own */ + carryNames(sr) { + return { C: `wgTile${sr.level}`, V: `wgValue${sr.level}` }; + } + buffers(sr) { + const { C, V } = this.carryNames(sr); + return [ + { name: C, access: "read_write", type: "array" }, + { name: V, access: "read_write", type: `array<${sr.datatype}>` }, + ]; + } + + /* Step d of the path either consumes an atom (atom < the current tile's + * end) or ends the current tile. Invocation i takes steps [i K, (i + 1) K); + * a binary search finds where its path starts. */ + wgsl(sr) { + const T = sr.datatype; + const W = WORKGROUP_SIZE; + const { C, V } = this.carryNames(sr); + return /* wgsl */ ` +const K = ${this.itemsPerThread}u; + +/* the carried (unfinished) tile and value of each invocation, then their + * segmented inclusive scan */ +var scanTile: array; +var scanValue: array<${T}, ${W}>; +var wgStart: u32; /* the tile the workgroup's path starts in */ +var headDone: u32; /* did the workgroup end it? */ +var headValue: ${T}; + +@compute @workgroup_size(${W}) +fn mergepath_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${sr.fnDeclarations.computeLinearizedGridParametersSplit()} + if (wgid >= sr.workgroups) { return; } + let lidx = builtinsNonuniform.lidx; + let total = sr.tiles + sr.atoms; + let diag = min(gid * K, total); + /* where this invocation's path starts: the most tiles ended (tile) such + * that tile ends come no later than the atoms before them */ + var lo = select(0u, diag - sr.atoms, diag > sr.atoms); + var hi = min(diag, sr.tiles); + while (lo < hi) { + let mid = (lo + hi) / 2u; + if (tileStartOf(mid + 1u) <= diag - mid - 1u) { lo = mid + 1u; } else { hi = mid; } + } + var tile = lo; + var atom = diag - lo; + var tileEnd = select(sr.atoms, tileStartOf(tile + 1u), tile < sr.tiles); + let stepEnd = min(diag + K, total); + var acc = IDENTITY; + var firstDone = false; /* ended the tile it started in? */ + var firstValue = IDENTITY; + for (var d = diag; d < stepEnd; d++) { + if (atom < tileEnd) { + acc = binop(acc, atomValue(atom, tile)); + atom++; + } else { + /* end of tile: the first tile ended may have started in an earlier + * invocation, so it waits for the scan; later ones are whole */ + if (!firstDone) { firstDone = true; firstValue = acc; } else { emitTile(tile, acc); } + acc = IDENTITY; + tile++; + tileEnd = select(sr.atoms, tileStartOf(tile + 1u), tile < sr.tiles); + } + } + if (lidx == 0u) { + wgStart = lo; + headDone = 0u; + } + /* segmented inclusive scan of the carried values, keyed by tile: tiles are + * nondecreasing across invocations, so a tile's carries are contiguous */ + scanTile[lidx] = tile; + scanValue[lidx] = acc; + workgroupBarrier(); + for (var off = 1u; off < ${W}u; off <<= 1u) { + let joins = lidx >= off && scanTile[lidx - off] == scanTile[lidx]; + var earlier = IDENTITY; + if (joins) { earlier = scanValue[lidx - off]; } + workgroupBarrier(); + if (joins) { scanValue[lidx] = binop(earlier, scanValue[lidx]); } + workgroupBarrier(); + } + if (firstDone) { + /* the tile this invocation started in: add what earlier invocations of + * this workgroup carried in it */ + var value = firstValue; + if (lidx > 0u && scanTile[lidx - 1u] == lo) { value = binop(scanValue[lidx - 1u], value); } + if (wgid > 0u && lo == wgStart) { + /* the workgroup's first tile (one invocation ends it), possibly begun + * in earlier workgroups: the fixup adds theirs and emits it */ + headDone = 1u; + headValue = value; + } else { + emitTile(lo, value); + } + } + workgroupBarrier(); + if (lidx == ${W - 1}u) { + /* what this workgroup carries out (to the tile unfinished at its end), + * and its first tile's value if it ended here */ + ${C}[wgid] = scanTile[${W - 1}u]; + ${V}[wgid] = scanValue[${W - 1}u]; + ${C}[sr.workgroups + wgid] = headDone; + ${V}[sr.workgroups + wgid] = headValue; + } +} + +/* after the fixup has written each head's whole value: emit the heads */ +@compute @workgroup_size(${W}) +fn heads_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${sr.fnDeclarations.computeLinearizedGridParametersSplit()} + let w = gid; + if (w == 0u || w >= sr.workgroups) { return; } + if (${C}[sr.workgroups + w] == 1u) { + emitTile(${C}[w - 1u], ${V}[2u * sr.workgroups + w]); + } +}`; + } + + actions(sr, { workgroups }) { + const { C, V } = this.carryNames(sr); + const actions = [ + /* per workgroup: carried tile, then head flag; carried value, then + * head value, then the head's whole value (written by the fixup, not in + * place, so that every kernel can be repeated: trials) */ + sr.allocate({ label: C, size: Math.max(16, 8 * workgroups) }), + sr.allocate({ label: V, size: Math.max(16, 12 * workgroups) }), + sr.kernel({ entryPoint: "mergepath_main", workgroups, label: "path" }), + ]; + /* one workgroup emits every tile itself; more need the fixup */ + if (workgroups > 1) { + const fixup = this.fixup(sr); + Object.assign(fixup, { tiles: workgroups, atoms: workgroups }); + actions.push( + sr.run(fixup, [C, V]), + sr.kernel({ entryPoint: "heads_main", workgroups: Math.ceil(workgroups / WORKGROUP_SIZE), label: "heads" }) + ); + } + return actions; + } + + /** The fixup: a SegmentedReduce over sr's W workgroup carries, whose + * tiles are sr's workgroups. Tile w's atoms are the carries of the run (the + * same carried tile) that ends at workgroup w - 1, if w ends that tile (w's + * head); runs start where the carried tiles, which are nondecreasing, first + * equal w - 1's. Its emit adds the run's carries before w's head value, + * into a third slot (never in place: a repeated kernel would add them + * twice). (Tile W - 1 also takes the last workgroup's carry, past the last + * tile: the identity.) It binds only sr's carry buffers, so every level + * binds the same few. */ + fixup(sr) { + const { C, V } = this.carryNames(sr); + const head = (w) => `${V}[sr.tiles + ${w}]`; + return sr.child("fixup", { + label: `${sr.label} fixup`, + schedule: new MergePath({ itemsPerThread: this.itemsPerThread }), + level: sr.level + 1, + tiles: 1, + atoms: 0, + buffers: [ + { name: C, type: "array" }, + { name: V, type: `array<${sr.datatype}>`, access: "read_write" }, + ], + wgsl: /* wgsl */ ` +fn runStart(t: u32) -> u32 { + if (t == 0u) { return 0u; } + if (t >= sr.tiles) { return sr.tiles; } + let key = ${C}[t - 1u]; + var lo = 0u; + var hi = t - 1u; + while (lo < hi) { + let mid = (lo + hi) / 2u; + if (${C}[mid] < key) { lo = mid + 1u; } else { hi = mid; } + } + return lo; +}`, + tileStart: (t) => `runStart(${t})`, + atom: (a) => `${V}[${a}]`, + emit: (w, v) => `if (${w} > 0u && ${C}[sr.tiles + ${w}] == 1u) { ${V}[2u * sr.tiles + ${w}] = binop(${v}, ${head(w)}); }`, + }); + } +} diff --git a/src/segmentedreduce.mjs b/src/segmentedreduce.mjs index 5b23059..27e19f5 100644 --- a/src/segmentedreduce.mjs +++ b/src/segmentedreduce.mjs @@ -35,303 +35,109 @@ * emit, buffers; tileStart ((t) => WGSL expression of tile t's first atom, * for t in [0, tiles], instead of an offsets buffer); wgsl (more WGSL * declarations, e.g. functions that atom, emit or tileStart call); - * schedule ("mergepath", default, or "thread"); itemsPerThread (mergepath: - * path steps per invocation, default 4). + * schedule (a Schedule, or the name of one: "mergepath", default, + * "thread" or "workgroup"; see schedules.mjs); itemsPerThread (passed to a + * schedule given by name, e.g. merge-path's steps per invocation). * - * Schedules (the paper's names): - * thread one invocation per tile, striding over the tiles. Balanced - * only when tiles are about the same size: one long tile - * holds up the kernel. - * mergepath the merge path of the tile ends and the atom indices (Merrill - * and Garland, SC 2016): tiles + atoms steps, split evenly, so - * every invocation does itemsPerThread steps whatever the tile - * sizes. A tile split between invocations is combined by a - * segmented scan in the workgroup. Between workgroups, combining - * is itself a segmented reduce: each workgroup's carry is an - * atom, and the tiles are the runs of carries in the same tile, - * each ending at the workgroup that finishes it. So a second - * SegmentedReduce (merge-path, on the carries; recursively, until - * one workgroup suffices) combines them into the finishing - * workgroups' head values, however many workgroups a tile spans, - * and a last kernel emits the heads. Each tile is emitted exactly - * once, with its whole value, so emit may be any WGSL (it is - * never asked to combine). + * SegmentedReduce holds the work; the schedule writes the kernels and says + * what runs. It knows no schedule by more than its name. */ import { BasePrimitive, Kernel, AllocateBuffer, ExecutePrimitive, UniformStruct } from "./primitive.mjs"; - -const WORKGROUP_SIZE = 256; -const SCHEDULES = ["mergepath", "thread"]; +import { WORKGROUP_SIZE, Schedule, ThreadMapped, WorkgroupMapped, MergePath } from "./schedules.mjs"; export class SegmentedReduce extends BasePrimitive { + /** schedules by name; add yours here, or pass an instance */ + static schedules = { thread: ThreadMapped, workgroup: WorkgroupMapped, mergepath: MergePath }; + constructor(args) { super(args); - if (!this.binop) throw new Error("SegmentedReduce: requires a binop (e.g. new BinOpAdd({ datatype: \"f32\" }))"); + if (!this.binop) throw new Error('SegmentedReduce: requires a binop (e.g. new BinOpAdd({ datatype: "f32" }))'); if (typeof this.atom !== "function") throw new Error("SegmentedReduce: requires atom: (atom, tile) => WGSL expression"); this.schedule ??= "mergepath"; - if (!SCHEDULES.includes(this.schedule)) throw new Error(`SegmentedReduce: schedule must be one of ${SCHEDULES.join(", ")}`); - this.itemsPerThread ??= 4; - if (!(Number.isInteger(this.itemsPerThread) && this.itemsPerThread >= 1)) throw new Error("SegmentedReduce: itemsPerThread must be a positive integer"); + if (typeof this.schedule === "string") { + const Named = SegmentedReduce.schedules[this.schedule]; + if (!Named) throw new Error(`SegmentedReduce: unknown schedule ${this.schedule}; choose from ${Object.keys(SegmentedReduce.schedules).join(", ")}`); + this.schedule = new Named(this.itemsPerThread !== undefined ? { itemsPerThread: this.itemsPerThread } : {}); + } + if (!(this.schedule instanceof Schedule)) throw new Error("SegmentedReduce: schedule must be a Schedule or a schedule's name"); this.datatype = this.binop.datatype; this.userBuffers = (this.buffers ?? []).map((b) => ({ access: "read", ...b })); if (!this.emit) { this.emit = (t, v) => `outputBuffer[${t}] = ${v};`; this.userBuffers.push({ name: "outputBuffer", access: "read_write", type: `array<${this.datatype}>` }); } - this.level ??= 0; /* 0 for the caller's; n for the nth level of fixup */ - this.wgTileName = `wgTile${this.level}`; - this.wgValueName = `wgValue${this.level}`; + /* nesting depth: 0 for the caller's, n for a schedule's nth level of + * child primitives (e.g. merge-path's fixup), which name buffers by it */ + this.level ??= 0; this.knownBuffers = [...(this.tileStart ? [] : ["offsets"]), ...this.userBuffers.map((b) => b.name)]; this.tileStart ??= (t) => `offsets[${t}]`; this.workgroupSize = WORKGROUP_SIZE; + this.children = {}; this.params = new UniformStruct("SRParams", { tiles: "u32", atoms: "u32", workgroups: "u32" }); - if (!("label" in args)) this.label += this.makeParamString(["schedule"]); - } - - /** workgroups of the main kernel */ - workgroupCount() { - return this.schedule === "thread" - ? Math.ceil(this.tiles / WORKGROUP_SIZE) - : Math.ceil((this.tiles + this.atoms) / (WORKGROUP_SIZE * this.itemsPerThread)); + if (!("label" in args)) this.label += ` [ schedule: ${this.schedule.name} ]`; } - /* the caller's pieces, as WGSL functions the schedules call */ - definitions() { + /** The work, as WGSL functions the schedules call, then the schedule's + * kernels: one module for all of a schedule's entry points. */ + moduleWGSL = () => { const T = this.datatype; return /* wgsl */ ` +${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} ${this.binop.wgslfn} const IDENTITY: ${T} = ${T}(${this.binop.identity}); ${this.wgsl ?? ""} fn tileStartOf(t: u32) -> u32 { return ${this.tileStart("t")}; } fn atomValue(atom: u32, tile: u32) -> ${T} { return ${this.atom("atom", "tile")}; } -fn emitTile(tile: u32, value: ${T}) { ${this.emit("tile", "value")} }`; - } - - threadWGSL = () => /* wgsl */ ` -${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} -${this.definitions()} -@compute @workgroup_size(${WORKGROUP_SIZE}) -fn thread_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { - ${this.fnDeclarations.computeLinearizedGridParametersSplit()} - for (var tile = gid; tile < sr.tiles; tile += totalThreadCount) { - var acc = IDENTITY; - let end = tileStartOf(tile + 1u); - for (var atom = tileStartOf(tile); atom < end; atom++) { acc = binop(acc, atomValue(atom, tile)); } - emitTile(tile, acc); - } -}`; - - /* Merge path: step d of the path either consumes an atom (atom < the - * current tile's end) or ends the current tile. Invocation i takes steps - * [i K, (i + 1) K); a binary search finds where its path starts. */ - mergepathWGSL = () => { - const T = this.datatype; - const W = WORKGROUP_SIZE; - return /* wgsl */ ` -${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} -${this.definitions()} -const K = ${this.itemsPerThread}u; -const NONE = 0xffffffffu; - -/* the carried (unfinished) tile and value of each invocation, then their - * segmented inclusive scan */ -var scanTile: array; -var scanValue: array<${T}, ${W}>; -var wgStart: u32; /* the tile the workgroup's path starts in */ -var headDone: u32; /* did the workgroup end it? */ -var headValue: ${T}; - -@compute @workgroup_size(${W}) -fn mergepath_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { - ${this.fnDeclarations.computeLinearizedGridParametersSplit()} - if (wgid >= sr.workgroups) { return; } - let lidx = builtinsNonuniform.lidx; - let total = sr.tiles + sr.atoms; - let diag = min(gid * K, total); - /* where this invocation's path starts: the most tiles ended (tile) such - * that tile ends come no later than the atoms before them */ - var lo = select(0u, diag - sr.atoms, diag > sr.atoms); - var hi = min(diag, sr.tiles); - while (lo < hi) { - let mid = (lo + hi) / 2u; - if (tileStartOf(mid + 1u) <= diag - mid - 1u) { lo = mid + 1u; } else { hi = mid; } - } - var tile = lo; - var atom = diag - lo; - var tileEnd = select(sr.atoms, tileStartOf(tile + 1u), tile < sr.tiles); - let stepEnd = min(diag + K, total); - var acc = IDENTITY; - var firstDone = false; /* ended the tile it started in? */ - var firstValue = IDENTITY; - for (var d = diag; d < stepEnd; d++) { - if (atom < tileEnd) { - acc = binop(acc, atomValue(atom, tile)); - atom++; - } else { - /* end of tile: the first tile ended may have started in an earlier - * invocation, so it waits for the scan; later ones are whole */ - if (!firstDone) { firstDone = true; firstValue = acc; } else { emitTile(tile, acc); } - acc = IDENTITY; - tile++; - tileEnd = select(sr.atoms, tileStartOf(tile + 1u), tile < sr.tiles); - } - } - if (lidx == 0u) { - wgStart = lo; - headDone = 0u; - } - /* segmented inclusive scan of the carried values, keyed by tile: tiles are - * nondecreasing across invocations, so a tile's carries are contiguous */ - scanTile[lidx] = tile; - scanValue[lidx] = acc; - workgroupBarrier(); - for (var off = 1u; off < ${W}u; off <<= 1u) { - let joins = lidx >= off && scanTile[lidx - off] == scanTile[lidx]; - var earlier = IDENTITY; - if (joins) { earlier = scanValue[lidx - off]; } - workgroupBarrier(); - if (joins) { scanValue[lidx] = binop(earlier, scanValue[lidx]); } - workgroupBarrier(); - } - if (firstDone) { - /* the tile this invocation started in: add what earlier invocations of - * this workgroup carried in it */ - var value = firstValue; - if (lidx > 0u && scanTile[lidx - 1u] == lo) { value = binop(scanValue[lidx - 1u], value); } - if (wgid > 0u && lo == wgStart) { - /* the workgroup's first tile (one invocation ends it), possibly begun - * in earlier workgroups: the fixup adds theirs and emits it */ - headDone = 1u; - headValue = value; - } else { - emitTile(lo, value); - } - } - workgroupBarrier(); - if (lidx == ${W - 1}u) { - /* what this workgroup carries out (to the tile unfinished at its end), - * and its first tile's value if it ended here */ - ${this.wgTileName}[wgid] = scanTile[${W - 1}u]; - ${this.wgValueName}[wgid] = scanValue[${W - 1}u]; - ${this.wgTileName}[sr.workgroups + wgid] = headDone; - ${this.wgValueName}[sr.workgroups + wgid] = headValue; - } -} - -/* after the fixup has added the earlier workgroups' carries into each head - * value: emit the heads */ -@compute @workgroup_size(${W}) -fn heads_main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { - ${this.fnDeclarations.computeLinearizedGridParametersSplit()} - let w = gid; - if (w == 0u || w >= sr.workgroups) { return; } - if (${this.wgTileName}[sr.workgroups + w] == 1u) { - emitTile(${this.wgTileName}[w - 1u], ${this.wgValueName}[2u * sr.workgroups + w]); - } -}`; +fn emitTile(tile: u32, value: ${T}) { ${this.emit("tile", "value")} } +${this.schedule.wgsl(this)}`; }; - /** The fixup: a SegmentedReduce over this one's W workgroup carries, - * whose tiles are this one's workgroups. Tile w's atoms are the carries of - * the run (the same carried tile) that ends at workgroup w - 1, if w ends - * that tile (w's head); runs start where the carried tiles, which are - * nondecreasing, first equal w - 1's. Its emit adds the run's carries - * before w's head value, into a third slot (never in place: a repeated - * kernel would add them twice). (Tile W - 1 also takes the last - * workgroup's carry, past the last tile: the identity.) It reads and writes - * only this level's carry buffers, so every level binds the same few. */ - fixupPrimitive() { - const C = this.wgTileName; - const V = this.wgValueName; - const head = (w) => `${V}[sr.tiles + ${w}]`; - this.fixup ??= new SegmentedReduce({ - device: this.device, - adapter: this.adapter, - label: `${this.label} fixup`, - binop: this.binop, - schedule: "mergepath", - itemsPerThread: this.itemsPerThread, - level: this.level + 1, - tiles: 1, - atoms: 0, + /* ---- for schedules' actions() */ + + /** a Kernel running one of the schedule's entry points, binding the + * caller's buffers, the schedule's and the parameters */ + kernel({ entryPoint, workgroups, label }) { + return new Kernel({ + kernel: this.moduleWGSL, + entryPoint, + label: [this.level ? `SegmentedReduce fixup level ${this.level}` : `SegmentedReduce (${this.schedule.name})`, label].filter(Boolean).join(" "), buffers: [ - { name: C, type: "array" }, - { name: V, type: `array<${this.datatype}>`, access: "read_write" }, + ...(this.knownBuffers.includes("offsets") ? [{ name: "offsets", access: "read", type: "array" }] : []), + ...this.userBuffers, + ...this.schedule.buffers(this), + { name: "sr", buffer: "srParams", uniform: this.params }, ], - wgsl: /* wgsl */ ` -fn runStart(t: u32) -> u32 { - if (t == 0u) { return 0u; } - if (t >= sr.tiles) { return sr.tiles; } - let key = ${C}[t - 1u]; - var lo = 0u; - var hi = t - 1u; - while (lo < hi) { - let mid = (lo + hi) / 2u; - if (${C}[mid] < key) { lo = mid + 1u; } else { hi = mid; } - } - return lo; -}`, - tileStart: (t) => `runStart(${t})`, - atom: (a) => `${V}[${a}]`, - emit: (w, v) => `if (${w} > 0u && ${C}[sr.tiles + ${w}] == 1u) { ${V}[2u * sr.tiles + ${w}] = binop(${v}, ${head(w)}); }`, + getDispatchGeometry: () => this.getSimpleDispatchGeometry({ workgroupCount: workgroups }), }); - return this.fixup; + } + /** a scratch buffer, not cleared between executes */ + allocate({ label, size }) { + return new AllocateBuffer({ label, size, clearBufferOnReuse: false }); + } + /** a SegmentedReduce the schedule runs (made once, kept under key, and + * destroyed with this one): same device and binop */ + child(key, options) { + this.children[key] ??= new SegmentedReduce({ device: this.device, adapter: this.adapter, binop: this.binop, ...options }); + return this.children[key]; + } + /** run a child primitive here, binding these of our buffers by the same names */ + run(primitive, names) { + return new ExecutePrimitive({ primitive, buffers: Object.fromEntries(names.map((n) => [n, n])) }); } compute() { if (!(Number.isInteger(this.tiles) && this.tiles >= 1)) throw new Error(`SegmentedReduce: tiles (${this.tiles}) must be a positive integer`); if (!(Number.isInteger(this.atoms) && this.atoms >= 0)) throw new Error(`SegmentedReduce: atoms (${this.atoms}) must be a nonnegative integer`); - if (this.tiles + this.atoms >= 2 ** 32 - WORKGROUP_SIZE * this.itemsPerThread) throw new Error("SegmentedReduce: tiles + atoms must fit in 32 bits"); - const workgroups = this.workgroupCount(); - const T = this.datatype; - const buffers = [ - ...(this.knownBuffers.includes("offsets") ? [{ name: "offsets", access: "read", type: "array" }] : []), - ...this.userBuffers, - ...(this.schedule === "mergepath" - ? [ - { name: this.wgTileName, access: "read_write", type: "array" }, - { name: this.wgValueName, access: "read_write", type: `array<${T}>` }, - ] - : []), - { name: "sr", buffer: "srParams", uniform: this.params }, + if (this.tiles + this.atoms >= 2 ** 32 - 2 ** 20) throw new Error("SegmentedReduce: tiles + atoms must fit in 32 bits"); + const workgroups = this.schedule.workgroups({ tiles: this.tiles, atoms: this.atoms }); + return [ + this.params.allocate("srParams", { tiles: this.tiles, atoms: this.atoms, workgroups }), + ...this.schedule.actions(this, { workgroups }), ]; - const dispatch = (count) => () => this.getSimpleDispatchGeometry({ workgroupCount: count }); - const label = this.level ? `SegmentedReduce fixup level ${this.level}` : `SegmentedReduce (${this.schedule})`; - const actions = [this.params.allocate("srParams", { tiles: this.tiles, atoms: this.atoms, workgroups })]; - if (this.schedule === "thread") { - actions.push( - new Kernel({ kernel: this.threadWGSL, entryPoint: "thread_main", label, buffers, getDispatchGeometry: dispatch(workgroups) }) - ); - } else { - actions.push( - /* per workgroup: carried tile, then head flag; carried value, then - * head value, then the head's whole value (written by the fixup, not - * in place, so that every kernel can be repeated: trials) */ - new AllocateBuffer({ label: this.wgTileName, size: Math.max(16, 8 * workgroups), clearBufferOnReuse: false }), - new AllocateBuffer({ label: this.wgValueName, size: Math.max(16, 12 * workgroups), clearBufferOnReuse: false }), - new Kernel({ kernel: this.mergepathWGSL, entryPoint: "mergepath_main", label: `${label} path`, buffers, getDispatchGeometry: dispatch(workgroups) }) - ); - /* one workgroup emits every tile itself; more need the fixup */ - if (workgroups > 1) { - const fixup = this.fixupPrimitive(); - Object.assign(fixup, { tiles: workgroups, atoms: workgroups }); - const names = [this.wgTileName, this.wgValueName]; - actions.push( - new ExecutePrimitive({ primitive: fixup, buffers: Object.fromEntries(names.map((n) => [n, n])) }), - new Kernel({ - kernel: this.mergepathWGSL, - entryPoint: "heads_main", - label: `${label} heads`, - buffers, - getDispatchGeometry: dispatch(Math.ceil(workgroups / WORKGROUP_SIZE)), - }) - ); - } - } - return actions; } destroy() { - this.fixup?.destroy(); + for (const child of Object.values(this.children)) child.destroy(); super.destroy(); } } diff --git a/src/spmv.mjs b/src/spmv.mjs index cd4a288..0af0de5 100644 --- a/src/spmv.mjs +++ b/src/spmv.mjs @@ -9,7 +9,8 @@ * [offsets[r], offsets[r + 1])), indices (nnz u32 column indices), values * (nnz f32), x (f32, one per column), y (rows f32, written). * Options: rows, nnz (both may change between executes), schedule - * ("mergepath", default, or "thread"), itemsPerThread; see SegmentedReduce. + * ("mergepath", default, "thread", "workgroup", or a Schedule: see + * schedules.mjs), itemsPerThread; see SegmentedReduce. */ import { SegmentedReduce } from "./segmentedreduce.mjs"; import { BinOpAdd } from "./binop.mjs"; From 6fd6f5e8e74c87d786c3aa9885f4bd6719085f3f Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 19:01:32 -0800 Subject: [PATCH 09/10] benchmark_spmv: merge-path at the library's default steps unless --items Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- misc/benchmark_spmv.mjs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/misc/benchmark_spmv.mjs b/misc/benchmark_spmv.mjs index 17f1b02..47ae154 100644 --- a/misc/benchmark_spmv.mjs +++ b/misc/benchmark_spmv.mjs @@ -16,7 +16,7 @@ * --mtx: Matrix Market files (e.g. from the SuiteSparse collection, * sparse.tamu.edu): coordinate real, integer or pattern (values 1), * general or symmetric. --items: merge-path steps per invocation to try - * (default 7). + * (default: MergePath's, 4). * * Per matrix: rows, nonzeros, mean and longest row; then per schedule the * GPU time (median after a 1 s warm-up), GFLOP/s (2 per nonzero), and the @@ -32,7 +32,7 @@ import { copyBandwidth } from "./device_rates.mjs"; const positional = process.argv.slice(2).filter((a) => !a.startsWith("--")); const option = (name) => process.argv.find((a) => a.startsWith(`--${name}=`))?.slice(name.length + 3); const mtxFiles = option("mtx")?.split(",") ?? []; -const itemsList = (option("items") ?? "7").split(",").map(Number); +const itemsList = option("items")?.split(",").map(Number) ?? [undefined]; /* undefined: the library's default */ const generated = mtxFiles.length && !positional.length ? [] : (positional[0] ?? "uniform,rmat,denserow").split(","); const logSizes = (positional[1] ?? "18,20,22").split(",").map(Number); @@ -181,7 +181,7 @@ async function benchmark(name, m) { const runs = [ ["thread", "thread", {}], ["workgroup", "workgroup", {}], - ...itemsList.map((K) => [`mergepath${itemsList.length > 1 ? ` K=${K}` : ""}`, "mergepath", { itemsPerThread: K }]), + ...itemsList.map((K) => [`mergepath${itemsList.length > 1 ? ` K=${K}` : ""}`, "mergepath", K === undefined ? {} : { itemsPerThread: K }]), ]; for (const [label, schedule, extra] of runs) { const spmv = new SpMV({ device, adapter, rows, nnz, schedule, ...extra }); From e81367e4b25328e334433c0a566720d38c0fdd99 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 19:26:26 -0800 Subject: [PATCH 10/10] docs: workgroup-mapped measured (never the one to use) Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- docs/segmented-reduce.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/segmented-reduce.md b/docs/segmented-reduce.md index af0d70e..ea50004 100644 --- a/docs/segmented-reduce.md +++ b/docs/segmented-reduce.md @@ -197,6 +197,16 @@ GFLOP/s (2 per nonzero) at 2^22 rows and 67M nonzeros, thread-mapped → merge-p That is 2.0–2.2× faster on the uniform matrix, 8–15× on R-MAT, and 1,500–2,800× with one dense row (thread-mapped takes 9.6–15.4 s there; merge-path 3.4–7.7 ms). +Workgroup-mapped, GFLOP/s at 2^22 rows, beside the best of the other two (the A100's run did not reach 2^22): + +| Matrix | M1 Max: workgroup / thread / merge-path | V100: workgroup / thread / merge-path | +|---|---|---| +| uniform (16 per row) | 0.90 / 8.1 / 17.5 | 3.40 / 8.8 / 17.5 | +| R-MAT | 0.90 / 1.6 / 19.2 | 3.47 / 2.3 / 18.4 | +| one row with half the nonzeros | 0.90 / 0.01 / 17.7 | 0.89 / 0.01 / 17.5 | + +Workgroup-mapped beats thread-mapped where long rows hold thread-mapped up (R-MAT on the V100; the dense row, 152 ms against 11.8 s), and loses to it on short rows, where one row leaves most of a workgroup idle. Merge-path is 4–20× faster than it on every matrix, so there is no reason to choose it: it is here as a third schedule, to show that one can be added without changing `SegmentedReduce`. + - **Merge-path's speed does not depend on the matrix's shape.** On each GPU it runs within 10% of the same GFLOP/s on all three matrices, at every size from 2^18 to 2^22 rows. Thread-mapped runs at the speed of its longest row. - **Merge-path also wins on the balanced matrix** on all three GPUs at 2^22. Neighboring invocations' loads are `itemsPerThread` elements apart, against a row (16 elements) apart for thread-mapped. Of 2, 4 and 7 steps, 4 was fastest or within 2% of it everywhere. (At 2^18 and 2^20 on the A100, the uniform matrix runs the same either way.) - **The recursive fixup removed the cost of long rows.** With the first version's walk-back fixup, the dense-row matrix took 9.2 ms on the A100 (5.8 ms of it fixup) and 14.9 ms on the M1 Max (7.4 ms); now 3.4 and 7.6 ms, the same as the uniform matrix.