---
library_name: kernels
license: apache-2.0
---

# graph-reduce

Segmented reductions for geometric deep learning, loadable through
`kernels`. Message passing reduces a per-edge tensor into a per-node one,
and the two access patterns that cover it want different kernels: scatter
(arbitrary output index per row, atomic writes) and segment (rows grouped by
a CSR pointer, each output owned by one block, no atomics). Both ship here.
The reference baselines are `torch.index_add_` and `torch.segment_reduce`,
the latter matched bitwise. Companion to
[graph-spmm](https://huggingface.co/kernels/phanerozoic/graph-spmm).

Edge lists are commonly stored sorted by destination, which makes the
atomic-free path available, and that path buys two things at once: it is
faster than both the atomic kernels and `torch.index_add_` in the regime
graph networks run, and its summation order is fixed, so a reduction that
has to be reproducible stops being a compromise. `scatter_max` and
`scatter_min` also return the arg, which torch's native `scatter_reduce`
does not.

![Two particle funnels: rows jam above the atomic bins while the segment lanes drain straight through, with live times and the run-to-run spread](https://huggingface.co/kernels/phanerozoic/graph-reduce/resolve/main/media/hero.gif)

*One million rows into 10,000 outputs, drain rates set by the measured
times: the atomic path at 0.047 ms with a run-to-run spread of 8e-6, the
segment path at 0.016 ms, bitwise identical on every run, 6.1x
`torch.index_add_`.*

## Usage

```python
import torch
from kernels import get_kernel

ts = get_kernel("phanerozoic/graph-reduce", version=1, trust_remote_code=True)

# general case: arbitrary destination per edge
out = ts.scatter_sum(messages, dst_index, num_nodes)
mx, argmx = ts.scatter_max(messages, dst_index, num_nodes)

# sorted edge list: build the CSR pointer once, then reduce without atomics
counts = torch.bincount(dst_index, minlength=num_nodes)
indptr = torch.cat([counts.new_zeros(1), counts.cumsum(0)])
out = ts.segment_sum_csr(messages_sorted_by_dst, indptr)
```

`version` selects the release branch; `trust_remote_code` is required by
`kernels` for publishers without the trusted-publisher mark.

## API

| Symbol | Purpose |
|---|---|
| `scatter_sum(src, index, dim_size)` / `scatter_add` | sum rows into the output named by `index` |
| `scatter_mean(src, index, dim_size)` | mean; empty outputs are 0 |
| `scatter_max(src, index, dim_size)` | `(values, argmax)` |
| `scatter_min(src, index, dim_size)` | `(values, argmin)` |
| `scatter(src, index, dim_size, reduce)` | dispatch by name |
| `segment_sum_csr(src, indptr)` | atomic-free sum over contiguous groups |
| `segment_mean_csr(src, indptr)` | atomic-free mean |
| `segment_max_csr(src, indptr)` / `segment_min_csr` | `(values, arg)` |

`src` is `[N, ...]`; `index` is `[N]` int64; `indptr` is `[S+1]` int64.

## Method

The scatter path is one atomic per element. For max and min the value and
its index resolve inside a single atomic: a float's bit pattern made
monotone as an int32, packed as `(mapped_value << 32) | ~index`, so the
largest 64-bit word carries both the maximum and the row that produced it,
with ties to the lower row. The segment path assigns one block per output
row; a segment's rows are contiguous and no other block touches that
output, so the reduction runs in registers and the store is a single write.

## Measured

fp32 sum of `[M, C]` into `S` outputs, median of 50, against
`torch.index_add_` (current torch and driver on the same host; atomic
throughput has improved since earlier published rows, and the table reflects
the current stack):

| M | C | S | `index_add_` | scatter | segment | segment vs `index_add_` |
|---|---|---|---|---|---|---|
| 1,000,000 | 16 | 10,000 | 0.087 ms | 0.047 ms | 0.016 ms | 5.4x |
| 1,000,000 | 64 | 10,000 | 0.415 ms | 0.338 ms | 0.358 ms | 1.16x |
| 4,000,000 | 16 | 50,000 | 0.482 ms | 0.357 ms | 0.366 ms | 1.32x |
| 1,000,000 | 16 | 100 | 0.250 ms | 0.273 ms | 0.182 ms | 1.37x |

The segment margin is largest on narrow features and moderate output
counts; at wide channels all paths are bandwidth-bound and converge.

## Correctness

- Sum and mean agree with `index_add_` to 1.7e-5 absolute over `M` to 10^6
  (fp32 atomic ordering); max and min are `torch.equal` to `index_reduce`'s
  `amax`/`amin`.
- The arg is checked, not assumed: every returned index is gathered back out
  of the source, must equal the reported extremum, and must belong to that
  row's segment; ties resolve to the lower row.
- `segment_csr` is bitwise equal to `torch.segment_reduce` (fixed summation
  order); empty outputs are 0 for sum/mean and `-inf`/`+inf` with arg `-1`
  for max/min.
- Reproducibility, measured: over repeated runs the CSR path is bitwise
  identical every time; the atomic path is not, because arrival order is not
  fixed. Use the CSR path when a reduction has to be reproducible.

## Requirements and limits

- NVIDIA GPU with compute capability 8.0+; fp32 only (the packed argmax
  atomic relies on the fp32 bit layout).
- The reduction is over dimension 0; permute first for other axes.
- The segment path requires input already sorted by output; sorting is worth
  it when amortized over many reductions on a fixed graph.
- Forward only; backwards compose from gathers and scatters.

## References

Merrill and Garland, "Merge-based parallel sparse matrix-vector
multiplication" (2016); the packed value-index atomic shared with
phanerozoic/dpx-decode.

## License

Apache-2.0.