SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning

One storage format can never fit a whole sparse matrix, and one compiler pass can never keep up with new hardware — SparseTIR splits matrices into per-region formats and compilation into stackable transformations, then searches over both.

Why this problem is worth solving

A sparse matrix is one that is mostly zeros, stored in a compressed form that keeps only the non-zero values plus index structures to find them. The storage choice is called the format: compressed sparse row (CSR) keeps a pointer per row into an array of column indices; ELLPACK (ELL) pads every row to the same length so the data is rectangular; block-compressed (BSR) stores small dense blocks. Every format is a bet about the sparsity’s shape, and the kernel code is welded to the bet — code written for CSR is useless for ELL.

Deep learning now makes all of these bets at once. Graph neural networks multiply feature matrices by adjacency matrices of real-world graphs — the Reddit graph in the paper has 233 thousand nodes and 115 million edges, with row lengths following a power law: a few celebrity nodes with enormous degree, a long tail of tiny rows. Sparse transformers mask attention with hand-designed block patterns. Pruning compresses BERT’s weights to 6–7% density in block or unstructured layouts. And the operators keep multiplying: relational graph convolutions (RGCN) over graphs with 45–96 edge types need a gather-matmul-scatter operator that, at publication, no sparse library implemented at all.

Meanwhile the hardware target moved: modern GPUs earn their throughput from Tensor Cores — matrix-multiply units that only engage if the compiler can rewrite loops into their instruction shapes (tensorization). A sparse kernel that cannot reach them leaves most of the chip’s arithmetic dark.

Writing each (operator × format × hardware) combination by hand does not scale. Compilation must — and the existing compilers had two structural ceilings.

What people did before, and where it broke

Sparse tensor compilers like TACO were the breakthrough that separated what to compute (a format-agnostic tensor expression) from how the data is stored (a per-tensor format annotation). But two design decisions, reasonable at birth, became the ceilings.

One format per matrix. Real matrices are regionally different. In a power-law graph, dense-ish celebrity rows would happily live as blocks while the tail wants narrow padded rows; a single CSR format forces one kernel to handle both, and on a GPU that means threads assigned to giant rows finish thousands of cycles after threads assigned to tiny ones — load imbalance by construction, or expensive runtime rebalancing.

Single-shot compilation. These compilers lower a high-level representation (a “provenance graph” of merged iterators) directly to code. Loops, memory regions, and branches never appear as manipulable objects — so loop-level optimizations cannot be expressed. Tensorization is loop-level pattern replacement; caching partial sums in registers is loop-level; two-stage reductions are loop-level. Dense tensor compilers (TVM, Halide) have spent a decade building exactly this machinery with multi-backend code generation — and single-shot sparse compilers could reuse none of it. Every hardware advance had to be re-implemented inside the sparse compiler, forever behind.

The idea

Language: axes, sparse buffers, sparse iterations

SparseTIR describes sparsity with axes, each carrying two independent attributes: dense or sparse (are the coordinates contiguous?) and fixed or variable (is the count per row constant?). A variable axis carries an indptr array (where each row starts); a sparse axis carries an indices array (which coordinates are present). Composing a few axes expresses CSR, BSR, ELL, diagonal, ragged, compressed-fiber and more. Sparse buffers bind values to a composition of axes — cleanly separating index structure (in the axes, shareable between tensors with the same layout) from values. Sparse iterations loop over axis compositions with a computation body, and — beyond TACO’s reach — allow arbitrary index expressions like A[i * 2 + ii, k], which is what makes block formats and convolution expressible.

Composable formats: decomposition as a transformation

The first key move: storing different regions of one matrix differently is not a new compiler — it is a program transformation on Stage I. Given a list of target formats, format decomposition rewrites one sparse computation into several (one per format), plus copy iterations that move data from the original format into the new ones. For the sparse structures of deep learning — a graph fixed across a whole training run, pruned weights fixed after pruning — those copies run once as preprocessing and are amortized to nothing.

The workhorse instance is the paper’s parameterized hyb(c, k) format for graphs: split columns into c partitions (so each partition touches only a slice of the dense operand — cache locality), then within each partition sort rows into buckets by length, padding each bucket’s rows to the next power of two. Each bucket is a little ELL matrix — perfectly rectangular, so each GPU thread block gets an identical amount of work: load balancing decided at compile time, no runtime scheduler. The padding wastes some arithmetic (16–29% zeros on the GNN graphs) and still wins, because regularity is worth more than FLOPs.

Composable transformations: three stages, two lowerings

Stage I thinks in coordinate space — actual matrix row/column numbers. Stage II thinks in position space — the index of a non-zero inside the compressed storage (if a row’s non-zero columns are 1, 3, 9, 10, then coordinate 9 has position 2). The sparse iteration lowering pass bridges them: it materializes the indptr/indices arrays as real buffers, emits one loop per axis with extents read from those buffers, and rewrites every buffer access from coordinates to positions — inserting a binary search only where an access genuinely cannot be translated directly.

Stage II is a dialect of TVM’s TensorIR, so the entire dense-compiler scheduling arsenal applies unchanged: split and bind loops to GPU threads, cache partial sums in registers, vectorize loads, apply rfactor for two-stage reductions — and tensorize block computations onto Tensor Cores. This is the payoff of not being single-shot: the state-of-the-art SDDMM library optimization (vectorized loads plus two-stage reduction) is expressible as two stock schedule primitives, with its constants exposed as tunable parameters. Finally, sparse buffer lowering flattens everything to 1-D arrays with computed offsets, leaving plain loop IR for TVM’s backends.

On top sits a tuner: formats (c, k, block sizes) and schedules (tile sizes, thread bindings, vector widths) form one joint search space, explored per sparse structure. Compilation takes time; the kernel is reused across thousands of training steps on the same graph.

with sp_iter([I, J, K], "SRS", "spmm") as [i, j, k]:
    C[i, k] += A[i, j] * B[j, k]

Why it works

Because a format is a bet about local sparsity shape, letting each region of a matrix make its own bet (decomposition into buckets) turns an irregular problem into several regular ones — and regular problems are what GPU scheduling can actually optimize, at compile time, with zero runtime balancing machinery. → Because compilation is a stack of transformations rather than one shot, each optimization is small, composable, and inspectable — and the stack can end inside an existing dense tensor compiler. → Because the final stage is TensorIR, every schedule primitive and backend that dense compilers have accumulated — thread binding, register caching, vectorization, rfactor, and above all Tensor Core tensorization — applies to sparse code without being reimplemented. → Because formats and schedules form one parameterized search space, and deep learning reuses a fixed sparse structure thousands of times, expensive tuning amortizes into cheap execution. → Therefore one system generates state-of-the-art kernels across four different sparsity regimes — graphs, attention masks, pruned weights, point clouds — including an operator (relational gather-matmul-scatter) that no library had implemented at all.

What it costs

• Padding buys regularity with FLOPs. The composable formats add 4–29% zero-arithmetic on the evaluated graphs. It wins anyway — but on a graph with centralized (non-skewed) degrees, like ogbn-proteins, the hybrid format’s benefit is largely eaten by its padding.
• Tuning time is real. The joint format-and-schedule search pays off only when the sparse structure is known at compile time and reused — the deep-learning common case, but not the general one. Dynamic structures would pay format-conversion costs at run time.
• Index construction is the user’s job. Each format-decomposition rule must be registered with code that computes the new formats’ index arrays; the paper wraps the ones it uses as APIs and leaves automatic inference to integration with other systems.
• Not universally fastest. For sparse convolutions at large channel counts (over 128), plain cuBLAS-backed implementations win — matrix multiplication grows quadratically with channels and vendor GEMM is better tuned than generated code there. Fixed-width formats also emit binary searches for some fused accesses.
• Multiple formats mean multiple kernels. Extra launches are mitigated, not eliminated, by horizontal fusion.

The evaluation

Two GPUs (Tesla V100 and RTX 3070), four workload families, geometric-mean speedups across feature sizes 32–512.