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..ea50004 --- /dev/null +++ b/docs/segmented-reduce.md @@ -0,0 +1,222 @@ +--- +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; +- `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`: 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 + +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: + +**`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 + +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. + +**Within a workgroup.** Each invocation walks its `itemsPerThread` steps: + +- 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*. + +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. + +**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. + +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 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 + +- **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.** 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 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 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 + +`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. + +GFLOP/s (2 per nonzero) at 2^22 rows and 67M nonzeros, thread-mapped → merge-path: + +| 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 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. +- **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 + +- **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.** 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 47adf13..8c793f9 100644 --- a/examples/regression.mjs +++ b/examples/regression.mjs @@ -21,6 +21,9 @@ 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 { Schedule } from "../src/schedules.mjs"; +import { SpMV, spmvReference } from "../src/spmv.mjs"; import { EvenBins, RangeBins, RadixBins } from "../src/bins.mjs"; export async function main(navigator) { @@ -1110,6 +1113,135 @@ ${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], + /* 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)), + /* 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 = {}, fixupLevels, trials = 1 }) { + 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, 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)); + 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.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(); + } + primitive.destroy(); + 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); + 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, trials }); + 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 +1590,57 @@ ${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", {}], + ["workgroup", "workgroup", {}], + ["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 } }) + ); + } + 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, + }) + ); + /* 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, + }) + ); + 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 } }) + ); + 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 } })); + await runTest(`segred SpMV ${shape}, ${schedule}, trials: 2`, () => + runSpMVTest({ shape, rows: 20000, cols: 5000, options: { schedule }, trials: 2 }) + ); + } + } + } + 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..a627da8 100644 --- a/index.d.ts +++ b/index.d.ts @@ -482,6 +482,117 @@ 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" }[]; + /** 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: 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; +} +/** the reduction on the CPU, for checking */ +export function segmentedReduceReference(args: { + offsets: Uint32Array; + tiles: number; + binop: BinOp; + 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; + nnz: number; + constructor(args: { + device: GPUDevice; + rows: number; + nnz: number; + schedule?: Schedule | string; + 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..f0c1f3f 100644 --- a/index.mjs +++ b/index.mjs @@ -25,6 +25,9 @@ export { suffixArrayTraffic, 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, randomBitsJS, diff --git a/misc/benchmark_spmv.mjs b/misc/benchmark_spmv.mjs new file mode 100644 index 0000000..47ae154 --- /dev/null +++ b/misc/benchmark_spmv.mjs @@ -0,0 +1,228 @@ +/* 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 + * 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: 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 + * 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")?.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); + +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", "thread", {}], + ["workgroup", "workgroup", {}], + ...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 }); + 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 fixupNS = kernels.filter((k) => /fixup|heads/.test(k.label ?? "")).reduce((a, k) => a + k.ns, 0); + 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)` + + (fixupNS ? ` fixup ${(fixupNS / 1e3).toFixed(1)} µs` : "") + + (bad >= 0 ? ` WRONG: row ${bad} = ${y[bad]}, expected ${expected[bad]}` : "") + ); + } + 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(); +} + +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/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 new file mode 100644 index 0000000..27e19f5 --- /dev/null +++ b/src/segmentedreduce.mjs @@ -0,0 +1,155 @@ +/** + * 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; 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 (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). + * + * 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"; +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 (typeof this.atom !== "function") throw new Error("SegmentedReduce: requires atom: (atom, tile) => WGSL expression"); + this.schedule ??= "mergepath"; + 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}>` }); + } + /* 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 += ` [ schedule: ${this.schedule.name} ]`; + } + + /** 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")} } +${this.schedule.wgsl(this)}`; + }; + + /* ---- 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: [ + ...(this.knownBuffers.includes("offsets") ? [{ name: "offsets", access: "read", type: "array" }] : []), + ...this.userBuffers, + ...this.schedule.buffers(this), + { name: "sr", buffer: "srParams", uniform: this.params }, + ], + getDispatchGeometry: () => this.getSimpleDispatchGeometry({ workgroupCount: workgroups }), + }); + } + /** 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 - 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 }), + ]; + } + + destroy() { + for (const child of Object.values(this.children)) child.destroy(); + super.destroy(); + } +} + +/** 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..0af0de5 --- /dev/null +++ b/src/spmv.mjs @@ -0,0 +1,57 @@ +/** + * 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, "thread", "workgroup", or a Schedule: see + * schedules.mjs), 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; +}