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

# graph-spmm

Sparse-dense matrix products for message passing, loadable through
`kernels`: a CSR adjacency applied to a dense feature matrix with sum, mean,
or max aggregation. The reference baseline is `torch.sparse.mm` (cuSPARSE),
matched to fp32 summation order and beaten up to 2.1x on the sparse,
narrow-feature shapes graph networks actually run. Companion to
[graph-reduce](https://huggingface.co/kernels/phanerozoic/graph-reduce) for
unsorted destinations.

A graph convolution is a sparse adjacency times a dense feature matrix,
applied once per layer over a graph that does not change. General sparse
libraries pay for generality, and atomic-based scatter alternatives return
different bits on every run. This kernel keeps one warp per output row with
no atomics anywhere, so message passing over million-node graphs is both
faster than cuSPARSE in its home regime and bitwise reproducible, and
max-aggregation returns the winning neighbour's index in the same sweep, the
piece a max-pool backward needs.

<video src="https://huggingface.co/kernels/phanerozoic/graph-spmm/resolve/main/media/hero.mp4" autoplay loop muted playsinline controls width="100%">
<img src="https://huggingface.co/kernels/phanerozoic/graph-spmm/resolve/main/media/hero.gif" alt="Heat from four sources spreads along the arms of a spiral graph, one spmm application per step">
</video>

*Heat from four sources diffusing through an 8,000-node kNN graph, one
`spmm` mean-aggregation per step (log-scale coloring). On 1M nodes and 10M
edges at 16 channels the product measures 0.39 ms against cuSPARSE's 0.85
ms, bitwise reproducible with no atomics.*

## Usage

```python
import torch
from kernels import get_kernel

tsp = get_kernel("phanerozoic/graph-spmm", version=1, trust_remote_code=True)

# build the CSR once for a fixed graph, then apply it per layer
adj = tsp.SparseAdj.from_edge_index(edge_index, num_nodes, edge_weight)
h = adj @ x                                  # sum aggregation
h = adj.reduce(x, "mean")                    # mean
h, arg = adj.reduce(x, "max")                # max, plus the winning neighbour
```

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

## API

| Symbol | Purpose |
|---|---|
| `spmm_sum(rowptr, col, val, mat)` | `A @ mat`; `val=None` is an unweighted adjacency |
| `spmm_mean(rowptr, col, val, mat)` | row-wise mean; empty rows are 0 |
| `spmm_max(rowptr, col, val, mat)` | `(values, argmax)`; empty rows are `(0, -1)` |
| `spmm(rowptr, col, val, mat, reduce)` | dispatch by name |
| `SparseAdj(rowptr, col, val)` | a CSR adjacency held once and applied many times |
| `SparseAdj.from_edge_index(edge_index, num_nodes, val)` | sort a `[2, E]` edge list into CSR |

`rowptr` is `[S+1]` int64, `col` is `[nnz]` int64, `mat` is `[K, C]` fp32.

## Method

One warp per output row. The row's nonzeros are walked cooperatively with
each lane holding a strided slice of the feature columns, so dense loads
coalesce across the warp and every output element reduces privately in a
register. Nothing is shared between warps, so there are no atomics: the
result is bitwise reproducible and independent of scheduling. `spmm_max`
resolves the argmax during the same sweep. `val=None` treats every stored
entry as 1 and skips the value load.

## Measured

fp32, median of 30, against `torch.sparse.mm` on a CSR tensor (cuSPARSE):

| nodes | nnz | C | avg degree | `torch.sparse.mm` | this | speedup |
|---|---|---|---|---|---|---|
| 10,000 | 200,000 | 32 | 20 | 0.028 ms | 0.023 ms | 1.23x |
| 10,000 | 200,000 | 128 | 20 | 0.035 ms | 0.032 ms | 1.10x |
| 100,000 | 1,000,000 | 32 | 10 | 0.081 ms | 0.038 ms | 2.11x |
| 100,000 | 4,000,000 | 64 | 40 | 0.301 ms | 0.245 ms | 1.23x |
| 1,000,000 | 10,000,000 | 16 | 10 | 0.894 ms | 0.494 ms | 1.81x |

The margin is largest on many nodes, low degree, modest channel count; at
high degree and wide features both are bandwidth-bound on the dense operand
and converge.

## Correctness

- Sum agrees with `torch.sparse.mm` to 1.5e-5 absolute across 10^4 to 10^6
  nodes (fp32 summation order); mean to 3.0e-7 against the row-normalized
  product.
- Max and its arg are checked against a per-row reference built from the CSR
  entries (a dense scatter silently drops duplicate pairs, which a real CSR
  can contain and this kernel sums).
- Unweighted adjacency matches an explicit all-ones value array; empty rows
  reduce to 0 for sum and mean and `(0, -1)` for max.
- Deterministic: repeated products are bitwise identical.
- `SparseAdj.from_edge_index` sorts an unsorted edge list into CSR and
  reproduces the same product.

## Requirements and limits

- NVIDIA GPU with compute capability 8.0+; fp32 only.
- One warp per row: power-law graphs with extreme hubs load-imbalance and
  want a segmented schedule this kernel does not implement.
- Forward only; the backward of a sum aggregation is the transposed product,
  composed from a second `SparseAdj` on the reversed edge list.

## References

Merrill and Garland, "Merge-based parallel sparse matrix-vector
multiplication" (2016); Bell and Garland, "Implementing sparse matrix-vector
multiplication on throughput-oriented processors" (2009).

## License

Apache-2.0.