A sparse tensor program's speed depends on the sparsity pattern itself — no format is universally best, format and schedule interact (co-optimizing both gains 2.02× where either alone gains 1.1×), and the config tuned for one matrix can run another at 0.37× — so WACO learns to read the pattern and pick both together, extracting features with a sparse convolutional network that operates on the raw matrix (strided so the receptive field grows past distant nonzeros, where downsampling-based CNNs destroy the pattern), embedding format and schedule in one SuperSchedule template scored by a ranking-loss cost model, and searching the huge joint space as an approximate nearest-neighbor walk over a graph of valid schedules — beating the best of four baselines (Intel MKL, format selectors, TACO defaults, ASpT) by 1.14–1.43× across SpMV, SpMM, SDDMM, and MTTKRP, and even learning the compiler's own vectorization quirks along the way.
The argument in one glance
-
Problem
In sparse computing, the input's shape is not the story — its pattern is
Dense tensor programs are tuned by shape alone. Sparse programs — SpMV in solvers, SpMM in graph learning, SDDMM in attention, MTTKRP in tensor factorization — live or die by WHERE the nonzeros sit: skewed row distributions demand fine-grained load balancing, clustered nonzeros reward blocked formats, and decades of format proposals (CSR, BCSR, DIA, ELL…) each win somewhere and lose elsewhere. Worse, the format and the loop schedule interact: on the TSOPF matrix, tuning either alone gains ~1.1×, tuning both together 2.02×.
-
What people did
Mechanisms without a policy, and single-axis tuners
TACO gave the MECHANISM: a format abstraction (coordinate hierarchies with per-level Uncompressed/Compressed storage) that generalizes the format zoo, plus schedule transformations (split, reorder, parallelize) — it can generate code for tens of thousands of format/schedule combinations. Auto-tuners supply partial policies: Intel MKL's inspector-executor tunes computation on a FIXED format; format-selection predictors pick a format but keep default schedules; features come from hand-crafted statistics or CNNs over downsampled matrices.
-
Why it fell short
The features destroy the pattern, and nobody tunes both axes
Downsampling a million-scale matrix to 128×128 for a CNN collapses structurally different patterns into identical inputs — a matrix of dense blocks and a diffuse one can downsample to the same image — and hand-crafted statistics must be redesigned per format and miss what matters. Meanwhile the coupled format-schedule space goes unexplored: MKL fixes the format, format selectors fix the schedule, and the anti-transfer data shows why partial answers fail — applying sparsine's optimal config to TSOPF runs at 0.96×, and TSOPF's config on sparsine at 0.37×. The right choice is per-input, and it lives in the JOINT space.
-
This paper's idea
Learn to read the raw pattern; search the joint space as a graph walk
WACO (MIT + UIUC) supplies TACO's missing policy with three pieces. WACONet reads the RAW sparse matrix — no downsampling — via submanifold sparse convolutions (compute only where the filter center hits a nonzero, so activations stay sparse and cheap) made STRIDED so the receptive field grows even when nonzeros sit far apart; 14 layers' pooled outputs concatenate into a 128-d pattern feature (~50% better loss than downsampled CNNs). The SuperSchedule template encodes format AND schedule as one parameter vector — splits, loop orders, level orders, U/C level formats, parallelization — whose degenerate settings (split=1) recover every sub-schedule, so CSR and BCSR are points in one space. A cost model trained with pairwise RANKING loss on ~6M measured samples scores (pattern, SuperSchedule) pairs; search becomes an approximate nearest-neighbor walk over a prebuilt KNN graph of valid schedules, descending predicted cost — 93.9% of search time spent actually evaluating candidates, versus 3.9–8.1% for Bayesian/ensemble tuners. Across 726 held-out matrices: 1.43×/1.18×/1.14×/1.27× over the BEST of four baselines on SpMV/SpMM/SDDMM/MTTKRP.
Why this problem is worth solving
Sparse kernels are the inner loops of scientific solvers,
graph analytics, and increasingly of machine learning — and
their tuning problem is categorically harder than the dense
one because the input data is part of the program’s
performance identity.
the compiler framework
this work builds on — a sparse tensor is a coordinate
hierarchy (a tree of index levels), each level stored either
Uncompressed (dense interval) or Compressed (explicit
coordinates); level splitting, level reordering, and per-level
format choices generate the whole classic format zoo (CSR =
UC, BCSR = UCUU) and thousands more, while schedules (split,
reorder, parallelize) independently transform the loop
traversal. TACO can GENERATE code for any combination; it has
no policy for choosing one
The instrument problem is what makes this a learning paper.
A cost model for
the standard architecture of
search-based program tuning — a fast cost model predicts a
candidate’s performance so hardware measurement is reserved
for finalists, and a search strategy proposes candidates. For
sparse programs, the cost model’s hardest input is the
sparsity pattern itself: millions of nonzeros that must be
summarized without destroying the very structure (block
shapes, skew, symmetry, locality) that determines
performance
What people did before, and where it broke
- Fixed expert implementations (CSR + default schedules, ASpT). Excellent on the patterns they anticipate; the anti-transfer tables show what “one size” costs elsewhere.
- MKL inspector-executor. Auto-tunes the computation at runtime — on a fixed format, forfeiting the axis where the biggest wins live (and it tunes only a few routines).
- Format selectors (BestFormat-style predictors). Choose among a handful of candidate formats, keeping default schedules — half the joint space, and the coupling (a format is only fast under a concordant traversal order) goes unmodeled.
- Hand-crafted pattern features. Row/column/nonzero statistics, diagonal distances — incomplete by construction, and each new format demands new features; in the ablation, a statistics-only cost model trails every convolutional one by a wide margin.
- CNNs over downsampled matrices. Fixed-size inputs force downsampling; million-scale matrices lose their local structure entirely — two matrices with opposite block characters collapse to the same 128×128 image, and the model cannot recover what the preprocessing destroyed.
How WACO works
Read the pattern, embed the choice, walk the graph
-
WACONet: sparse convolution on the raw matrix new
Submanifold sparse convolution computes an output only where the filter's CENTER lands on a nonzero — activations stay exactly as sparse as the input through every layer, so a 10-million-nonzero matrix is processed as-is, no downsampling. The catch: with stride 1, information cannot propagate between distant nonzeros (the receptive field stalls). WACONet's fix is strided convolution (3×3, stride 2) at every layer, FORCING the receptive field to double repeatedly; 14 layers deep, 32 channels wide (GPU memory bounds), with all 14 pooled intermediate outputs concatenated into the 128-d feature to recover capacity lost to narrow channels.
-
SuperSchedule: one template for both axes new
Each index is split once (i→i1,i0; k→k1,k0), making the template a SUPERSET: split size 1 collapses a level, so all sub-schedules — and all formats — are settings of one parameter vector. Compute-schedule parameters (split sizes, loop order, parallelized index, thread count, OpenMP chunk size) fix the traversal; format-schedule parameters (per-tensor level orders, U/C level formats) fix the storage. CSR, BCSR, row- vs column-major, sparse-blocked variants: all points in one searchable space, with the format-schedule coupling representable because both live in the same vector.
-
The cost model, trained to rank new
The program embedder maps SuperSchedule parameters through learned lookup tables (categoricals) and permutation-matrix MLPs (orders) into a program embedding; concatenated with the WACONet feature, linear-ReLU layers emit a score. The loss is pairwise RANKING (hinge on sign-of-runtime-difference), not L1/L2: the search only needs the argmin, so the model learns orderings, which is an easier and better-conditioned target than absolute runtimes. Training data: 21,400 matrices (SuiteSparse, augmented) × 100 random SuperSchedules each, measured — ~6M tuples, two weeks on ten nodes.
-
ANNS: search as a graph walk new
Auto-scheduling is cast as nearest-neighbor search: dataset = SuperSchedules, query = the input matrix, distance = predicted cost. A KNN graph (HNSW) is prebuilt over PROGRAM EMBEDDINGS (l2 distance) of training-set SuperSchedules; at query time the walk descends predicted cost along graph edges. Two economies: the pattern feature is computed ONCE per matrix (only the final predictor re-runs per candidate), and there is no optimizer metadata — where Bayesian/ensemble tuners spend 96.1%/91.9% of wall time on surrogate bookkeeping, ANNS spends 93.9% actually evaluating candidates.
-
Why a graph, not gradients new
The cost model is differentiable, so gradient descent over the schedule encoding is tempting — but one-hot categoricals and permutation matrices make almost every gradient step an INVALID encoding. The KNN graph contains only valid SuperSchedules, so the walk is effectively discrete gradient descent constrained to the valid set — the same preserve-validity-by-construction instinct as constraint-evolving search, realized geometrically.
-
Finalist measurement
The top-10 SuperSchedules by predicted cost are compiled by TACO and measured; the fastest ships. Hardware touches ten candidates, not thousands — the cost model and ANNS carry the rest.
The idea
The feature extractor must respect the data structure it summarizes. Every prior featurization imposed a dense worldview on sparse data — statistics that flatten structure, or downsampling that erases it. WACONet’s architecture is the insight: submanifold convolution gives cost proportional to nonzeros (so the raw matrix is affordable), and strided layers solve sparse convolution’s known blind spot (isolated nonzeros never communicating) by forcing receptive-field growth. The ~50% loss improvement over downsampled CNNs — and the margin over MinkowskiNet, the 3D-vision architecture lacking the striding — measures exactly the information the other pipelines were destroying. It is the site’s recurring lesson about learned systems in a new key: architecture choices ARE domain knowledge.
Make the coupled space one space. Format and schedule interact — a format is fast only under a concordant traversal — so tuning them separately forfeits the coupling (the 2.02×-vs-1.1× motivating gap). SuperSchedule’s split-once superset trick puts both axes in a single parameter vector where degenerate settings recover every classic point: the search space contains CSR-with-default-schedule and exotic sparse-blocked layouts as neighbors in one geometry. The discoveries this enables are the paper’s best exhibits: a UUC “sparse block” format with a 16,384-wide compressed inner split that cuts LLC misses from 36% to 7% (2.5× on sparsine), and — most delicious — the model learning that the Intel compiler emits AVX instructions only for block sizes of 16 and up, so it deliberately chooses memory-wasting half-empty blocks to trigger vectorization. No analytical cost model contains icc’s heuristics; a learned one absorbed them from measurements.
Spend learning on ordering; spend search time on candidates. Two economy decisions compound. The ranking loss trains the model for the only property search consumes — which schedule is faster — sidestepping the harder absolute-runtime regression (kin to this site’s forecasting page learning only utilization). And casting search as ANNS strips away optimizer overhead: no surrogate training, no metadata, just a prebuilt graph of valid schedules walked downhill in predicted cost, with the expensive pattern feature amortized to one computation per matrix. The 93.9%-vs-3.9% useful-work ratio against Bayesian optimization is the measure of how much conventional black-box search spends on itself.
A new matrix meets WACO
sparsine arrives: 50,000 rows, 1.55M nonzeros, diffuse
SuperSchedule = one vector fixing format + schedule. KNN graph = prebuilt over program embeddings. Numbers from the paper's evaluation and speedup analysis.
Step 1 of 6
One pass of WACONet reads the pattern
The raw matrix — no downsampling — flows through 14 strided sparse-convolution layers. Early layers see local texture (no dense blocks here; nonzeros scattered); striding doubles the receptive field each layer until the global picture (uniform diffuse spread, no skew) is in view. The 14 pooled outputs concatenate into a 128-d feature. Computed once; reused for every candidate scored.
cost scales with nonzeros, not matrix area — a 10M-nonzero matrix fits downsampling would have made sparsine indistinguishable from a blocked matrix
Why it works
- Submanifold convolution’s cost tracks nonzeros → the raw matrix is affordable at any shape → no information is destroyed before learning begins, and striding restores the receptive-field growth sparsity would otherwise stall — the feature extractor finally matches the data structure.
- Format and schedule share one parameter vector → the coupled optimum is a searchable point, not a composition of per-axis compromises → discoveries like the UUC sparse block and vectorization-triggering block sizes live only in the joint space.
- Search consumes only orderings → ranking loss trains exactly that, an easier target than absolute runtime → the cost model’s capacity goes where the argmin needs it.
- The KNN graph holds only valid schedules → cost-descending walks are discrete gradient descent that cannot leave the space → validity by construction, with the pattern feature amortized to one computation and zero optimizer metadata.
- Measurements contain the whole stack’s behavior → the learned model absorbs compiler heuristics, cache dynamics, and load-balancing effects no analytical model enumerates → policy quality exceeds what its designers explicitly knew.
What it costs
- Training is an expedition. ~2M measured tuples per algorithm (two weeks, ten nodes), per hardware target — and the transfer study shows hardware-specificity: an Intel-trained model helps on AMD but leaves a fifth of the win behind. New algorithm, new dataset.
- Search overhead needs amortization. WACO’s tuning cost (~821 SpMV-calls-equivalent) exceeds MKL’s inspector (113) and BestFormat (277); one-shot kernels favor cheaper tuners, and WACO wins where sparse kernels actually live — repeated invocation in solvers and training loops.
- The template bounds the space. One split per index (more showed diminishing returns), a fixed parallelization menu, and the KNN graph built from training-set schedules — near-optimal within the SuperSchedule superset, silent about points outside it.
- CPU scope. Evaluation is CPU-only (TACO’s mature backend); GPUs’ different cost surfaces and formats are future work, as are dynamic sparsity patterns that change between invocations.
- Ten measurements still required. The top-1 prediction is not trusted alone; the top-10-measure protocol is honest about residual cost-model error.
The evaluation
726 held-out real matrices, four algorithms spanning the sparse-kernel canon, four baselines covering both fixed and auto-tuned prior art, ablations isolating the feature extractor (four alternatives) and the search strategy (two black-box tuners), a cross-hardware transfer study, and end-to-end amortization scenarios.
Who chooses the format and schedule, and from what evidence
| Property | Fixed expert implementations (FixedCSR / ASpT) | Single-axis auto-tuners (MKL inspector / BestFormat) | WACO (learned joint co-optimization) |
|---|---|---|---|
| Tuning space | |||
| Pattern evidence | None | ||
| Search machinery | — | Inspector heuristics; classifier inference | |
| Results (726 unseen matrices) | Baseline (best-of-four includes these) | Baseline (best-of-four includes these) | |
| Costs | None at use; permanent at mismatch | Cheap tuning (MKL: 113 SpMV-equivalents) |
Check yourself
Check your understanding
1. The format and schedule optimized for TSOPF run sparsine at 0.37× — worse than untuned CSR. Predict why sparse optimizations anti-transfer so violently when dense-program tunings mostly generalize.
2. Feeding a CNN a million-scale sparse matrix seems impossible, and downsampling destroys the pattern. How does WACONet thread the needle — and why does it also beat MinkowskiNet, an existing sparse-convolution architecture?
3. WACO trains its cost model with a pairwise ranking loss rather than regressing runtimes, and searches by walking a KNN graph rather than running Bayesian optimization. What do these two choices each buy?
4. WACO sometimes chooses block sizes of 16+ for blocks less than half full — deliberately storing zeros and inflating memory. The paper traces this to something no analytical cost model would contain. What, and what does it imply about learned cost models?
Sources
- ACM DL: WACO: Learning Workload-Aware Co-optimization of the Format and Schedule of a Sparse Tensor Program — ASPLOS 2023 (open access)