Getting a neural network to run fast on a particular chip means choosing among billions of ways to arrange its loops, and compilers do it by searching — trying candidates, timing them, mutating the good ones — which takes hours or days because a search that only knows whether a candidate was better is walking blind; this paper (Illinois) makes the whole space continuous and the whole latency estimate differentiable, so the compiler can compute which direction to move rather than sampling around, reaching 90 percent of its eventual best about three times sooner.
The argument in one glance
-
Problem
The same computation has millions of arrangements and they are not equally fast
A neural network layer is a small amount of mathematics wrapped in loops, and essentially all of its performance comes from how those loops are arranged for the particular chip it runs on — how they are split into tiles that fit a cache, which are unrolled, which are spread across parallel units. Doing this by hand needs deep expertise, has to be redone for every chip, and does not survive contact with the growing variety of hardware people want to deploy on.
-
What people did
Let the compiler search for the arrangement
Automatic tensor compilers define a space of schedules — sequences of transformations with tunable parameters — and search it. Because timing a candidate on real hardware is slow, a learned model predicts latency from features of the generated program, and an evolutionary algorithm or beam search proposes candidates the model scores. This works, and it is how the field currently produces high-performance kernels without human tuning.
-
Why it fell short
A search that only learns which candidate won is walking blind
Combinatorial search extracts almost nothing from each evaluation — it learns that this candidate scored higher than that one, and nothing about where to look next. In a space this large that means enormous numbers of candidates and tuning runs measured in hours or days. Some systems respond by shrinking the space with hand-written templates or aggressive pruning, which trades away the possibility of finding what was pruned.
-
This paper's idea
Make the space continuous so the compiler can differentiate it
This paper (Yifan Zhao, Hashim Sharif, Vikram Adve and Sasa Misailovic) observes that gradients carry far more information per evaluation than comparisons do. It relaxes the integer schedule parameters into real numbers, derives the program features as symbolic formulas in those parameters rather than measuring them, smooths the non-differentiable pieces, and turns legality rules into penalty terms — producing one differentiable function whose minimum is a fast program.
Why this problem is worth solving
The object being searched over deserves a precise description.
THE
mathematics of a tensor operation says what to compute, not how. A SCHEDULE is the sequence of
transformations that turns that specification into an actual loop nest — split this loop into
tiles of some size, reorder the resulting levels, unroll this one, vectorise that one, map
these levels onto blocks and threads of a graphics processor. Most of these transformations
carry integer parameters, and splitting a loop into several levels introduces one parameter
per split. The number of combinations is the product across every choice, which is
astronomical, and the overwhelming majority of points in it are either illegal — a tile size
that does not divide, a configuration that exceeds what the hardware provides — or simply
bad. So the space is enormous, sparse in useful points, and constrained in ways that must be
respected rather than discovered.
The second thing to establish is what the search is actually optimising against.
RUNNING
a candidate on the target hardware gives the true answer and is far too slow to do millions of
times, so tensor compilers learn a proxy. A program’s generated code is summarised by a vector
of features — how many additions it performs, how far apart repeated accesses to a buffer are,
how much parallel work each level carries — and a neural network trained on many measured
programs maps that vector to a predicted running time. The search then optimises the
prediction and only occasionally checks reality. This detail is what makes the whole idea in
this paper possible, because hardware is not differentiable but a neural network is. If the
features can also be made differentiable functions of the schedule parameters, the composition
of the two is a differentiable estimate of latency, and gradient descent applies.
What people did before, and where it broke
-
Hand-written kernel libraries. Every mainstream framework ships expert-tuned implementations of common operations. They are excellent, and for the operations they cover on the hardware they target, hard to beat.
-
The gap. They require expertise per operation and per chip, and they do not extend to the growing range of edge and accelerator platforms. On unusual operations or fused sequences, there is no library entry at all.
-
Template-based autotuning. Let a human write a parameterised schedule template per operation class, then search only the parameters. This shrinks the space to something searchable.
-
The gap. Someone must write and maintain a template for every operation, and the search can never find anything the template cannot express.
-
Search over an automatically generated space. The state of the art generates the space itself and explores it with an evolutionary algorithm guided by a learned cost model, which removes the templates and finds genuinely good schedules.
-
The gap. Speed. Each evaluation yields one comparison, so the search needs a great many of them, and tuning a network takes hours. On a constrained device, where tuning may run on the device itself, that is prohibitive.
-
Prune the space aggressively. Reduce tuning time by cutting the space down before searching it.
-
The gap. Whatever was pruned is gone. This trades the risk of a slow search for the certainty of an incomplete one.
How a compiler differentiates its own search space
From a tensor operation to a schedule found by gradient descent
-
-
new
-
new
-
new
-
new
-
new
The idea
A gradient says where to go; a comparison only says which was better. The motivating argument is about information per evaluation. An evolutionary search that scores two candidates learns their ordering and nothing about the shape of the landscape, so it must compensate with volume. A gradient at a point gives a direction in every dimension at once, which is why gradient methods dominate wherever they can be applied — training networks, searching architectures — and why they had not been applied here. The obstacle is not that schedules are unusual but that the space is discrete and the objective is a black box. The paper attacks both. Integer parameters are relaxed into real intervals, so an unroll factor that could be any whole number from 1 to N becomes any real number in that range. And the objective is made differentiable by composing two things that already are: a learned cost model, which is a neural network, and program features, which the paper makes into algebraic formulas rather than measurements. Composing a learned latency predictor with something the compiler can reason about symbolically is the same architecture as the cost-model work covered elsewhere on this site, which replaces hand-crafted features with a model that reads the schedule directly.
The features have to become formulas, and that is the contribution. Ordinarily a cost model is fed features extracted from a concrete generated program by counting things. That is useless for gradient descent, because counting is not a function you can differentiate against a tile size. So the analysis pass here runs on a symbolic program — one whose loop bounds and indices still contain the schedule variables — and emits each of its 82 features as an expression. The number of floating-point additions becomes a product of loop extents; the number of parallel blocks becomes a division by a tile size. Some expressions contain conditionals, since a feature may depend on whether a tile size exceeds one, and those are not differentiable, so a small set of rewrite rules substitutes smooth approximations. The paper notes there are fewer than ten such operators and that handling them was a one-time cost that does not grow as transformations are added. There is also a numerical problem worth naming: some features reach a hundred million or more, at which scale a small change makes no relative difference and the gradient vanishes. Taking logarithms of the features and substituting an exponential for each variable fixes it. That what the cost model is fed matters as much as the model itself is the lesson of the sparse co-optimisation work covered elsewhere on this site, which throws away hand-crafted descriptors and learns features from the raw problem instead.
Constraints become penalties, and the discrete answer comes at the end. Relaxation creates a problem the original space did not have — real-valued tile sizes are meaningless and many regions of the continuous space correspond to illegal programs. The rules that define legality are themselves differentiable expressions, so each becomes a term that is zero when satisfied and grows quadratically when not, which is the standard way of converting a constrained optimisation into an unconstrained one. The optimiser is then free to wander through illegal territory while being pushed back, and the discretisation happens only at the end: every point visited during descent is rounded to integers, the illegal roundings are discarded, and the survivors are ranked by the cost model with the best few actually timed on hardware. Those measurements also refine the cost model for later rounds. Optimising a relaxed problem and recovering a discrete answer afterwards is the opposite tactic to the equality-saturation work covered elsewhere on this site, which stays discrete and instead represents all equivalent programs at once so nothing has to be committed to early.
Turning one tile size into a derivative
How a schedule parameter becomes something to differentiate
Following a single tile-size variable from a symbolic schedule to a measured program.
Step 1 of 5
Apply the schedule with the parameter left as a symbol
The transformation says to split a loop into tiles, but the tile size is a variable rather than a number. Applying it produces a program whose loop bounds and array indices contain that variable.
This is the step that differs from every existing tuner, which instantiates a concrete schedule and then examines the result. Here the result is a family of programs described in one expression, and the search happens over that description.
Why it works
- A gradient carries more information than a comparison. Each evaluation yields a direction in every dimension rather than an ordering → far fewer candidates are needed to make the same progress → 90 percent of the eventual best performance arrives about 3.4 times sooner than with evolutionary search.
- Symbolic features make the objective differentiable end to end. Deriving each feature as a formula in the schedule variables rather than counting it in generated code → the composition with a neural cost model is differentiable by construction → gradient descent applies to a quantity that was previously only observable.
- Penalties keep a constrained space navigable. Expressing legality as terms that grow when violated → the optimiser can pass through illegal regions under pressure to return → paths to good areas are not blocked by thin bands of invalidity, as a hard barrier would.
- Rescaling prevents a flat landscape. Taking logarithms of features that reach hundreds of millions and substituting exponentials for the variables → relative changes stay visible wherever the optimiser is → the gradient does not vanish across most of the space.
- The relaxation never has to be trusted. Every visited point is rounded, checked against the original integer constraints, and the survivors are timed on real hardware → a continuous space that does not perfectly correspond to reality can still be a good proposal mechanism → correctness rests on the discrete check, not on the relaxation.
What it costs (trade-offs and scope)
- It finds good schedules sooner, not better ones. The paper is explicit that this approach and evolutionary search converge to similar final performance. The contribution is the shape of the curve, which matters when tuning time is bounded and not otherwise.
- Every non-differentiable operator needs a smooth replacement. Fewer than ten were required and they were written by hand. The authors argue this is a one-time cost independent of the transformations, and suggest it could be automated, but it is manual work in the artifact as it stands.
- It inherits a whole compiler. The space of schedule structures, the partitioning into subgraphs, the round-based tuning loop and the scheduler that picks which subgraph to work on all come from the system it is built on. What is replaced is the inner search.
- The cost model is still a proxy. Gradient descent optimises a prediction, so it can only be as good as that prediction, and the mechanism that keeps it honest — timing a handful of candidates and refitting — is inherited rather than improved.
- Everything is graphics processors. Three platforms are evaluated, spanning server, desktop, and edge, but all from one vendor. Nothing establishes how the feature formulas or the smoothing behave on a different kind of accelerator.
The evaluation
| Property | What was measured | The result | What it tells you |
|---|---|---|---|
| Against shipping frameworks | |||
| Against the search-based compiler | |||
| When the work per candidate is larger | |||
| Individual operators | |||
| Where the gain is largest |
Check yourself
Check your understanding
1. Why does the paper argue that gradient descent should beat evolutionary search here, given both use the same cost model?
2. What had to be made differentiable, beyond the cost model itself?
3. Why are legality rules expressed as penalty terms rather than as hard barriers the optimiser cannot cross?
4. Predict what happens to this approach's advantage when each candidate is more expensive to evaluate on hardware, as with larger batch sizes.
Sources
Felix: Optimizing Tensor Programs with Gradient Descent. Yifan Zhao, Hashim Sharif, Vikram Adve, Sasa Misailovic. ASPLOS 2024, Volume 3, pages 367 to 381. DOI 10.1145/3620666.3651348. Full text read from the authors’ copy hosted at the University of Illinois. The speedup milestones quoted here follow the evaluation section, which states 90 and 95 percent of peak performance; the abstract states the same two figures against 95 and 99 percent.