A tensor compiler generates millions of candidate program variants but can afford to actually TIME only a few thousand, so a learned latency predictor steers the search — yet prior cost models needed expert-crafted, per-architecture features and collapsed on new hardware; TLP instead reads the SCHEDULE, treating optimization primitives as words in a language and latency prediction as an NLP regression task, then adds per-platform multi-task heads — cutting search time 9.1×/3.0× versus the state of the art and standing up on a new platform with 7% of the data.
The argument in one glance
-
Problem
The compiler can imagine millions of programs but time only thousands
Deep-learning compilers like TVM don't hand-write fast kernels — they SEARCH for them: each network subgraph has millions (CPU) to billions (GPU) of computationally-equivalent implementations differing only in loop tiling, ordering, vectorization, and parallelization. Finding the fast ones requires knowing each variant's latency — but a single real measurement means compiling, loading, repeating runs with cache flushes, and monopolizing the machine: hundreds of milliseconds per program. Measuring the whole space is impossible.
-
What people did
Predict latency with a learned cost model
Auto-tuners (AutoTVM, Ansor, TIRAMISU) therefore train a model to predict latency and let it triage: tens of thousands of candidates are scored by the model, and only dozens of survivors touch real hardware. Deep-learning cost models beat all alternatives — TenSet even released a 51.57-million-program dataset across 6 platforms so models could be pre-trained.
-
Why it fell short
Hand-crafted features, and models that die on new hardware
Ansor's model needs 164 expert-picked features per innermost statement; TIRAMISU's needs 2,534 plus the syntax tree. Only hardware-architecture experts can design these, the features differ between CPUs and GPUs — even between Intel, AMD, and ARM CPUs — and they still miss things. Worse is cross-hardware unavailability: a model trained on one machine performs badly on another, and collecting a fresh multi-million-sample dataset per machine is brutal (the authors spent 50+ days gathering 8.65M samples on one laptop CPU).
-
This paper's idea
Read the schedule as a language; give each machine its own head
TLP never looks at the generated program. It featurizes the SCHEDULE PRIMITIVE SEQUENCE — the list of transformation commands (split, reorder, fuse, ...) that PRODUCED the program — treating primitives as words and latency prediction as NLP regression: one-hot the primitive type, keep numbers raw, tokenize names, feed 25×22 dense features to a single 8-head self-attention layer trained with a ranking loss. MTL-TLP adds multi-task learning: shared layers learn hardware-INDEPENDENT structure, per-platform heads learn hardware-DEPENDENT behavior. Search reaches state-of-the-art quality 9.1×/3.0× faster (CPU/GPU), and a new platform needs only ~500K samples — 7% of a full dataset.
Why this problem is worth solving
Every model that ships to production passes through a stage where
someone — or something — decides exactly how its
the low-level loop nest implementing one
neural-network subgraph (a fused convolution, a matrix multiply
chain). Equivalent tensor programs differ only in schedule — how
loops are tiled, ordered, fused, vectorized, parallelized — but
their latencies differ by orders of magnitude
a predictor that estimates a tensor program’s
latency without running it. The auto-tuner scores tens of thousands
of candidates with the model and sends only the predicted-best
dozens to real hardware — so the model’s ranking quality directly
bounds search quality and speed
The two problems TLP attacks are both bottlenecks of scale. Feature engineering doesn’t scale across architectures: every new target (Intel vs ARM vs AMD CPUs, let alone GPUs) needs experts to re-derive what matters, and the hand-picked features still underperform. And trained models don’t scale across machines: latency landscapes differ enough between platforms that a model is effectively stranded on the hardware whose data trained it — while gathering a fresh 8-million-sample dataset per machine costs weeks of dedicated compute. Any organization tuning models for a fleet of heterogeneous devices hits both walls immediately.
What people did before, and where it broke
- Template and search frameworks. AutoTVM searches within
hand-written templates; Ansor (the state-of-the-art auto-tuner in
TVM, and TLP’s host) generates schedules hierarchically —
applied in sequence, guided by a genetic algorithm plus a cost model.
the transformation commands an auto-tuner applies to derive a concrete tensor program from a subgraph: split a loop, reorder loops, fuse operators, vectorize, parallelize, cache-read/write, unroll. Ansor uses 11 types; a full sequence (up to ~54 primitives) IS the recipe that produces the program
- Hand-featured learned models. Ansor’s MLP consumes 164 features per innermost statement (arithmetic counts, memory access patterns, loop structure); TIRAMISU’s consumes 2,534 features propagated along the program’s syntax tree. Both work — and both encode expert judgments that differ per hardware class and still leave accuracy on the table.
- Datasets and transfer attempts. TenSet contributed the large-scale dataset (51.57M measured programs across 6 platforms) and an MLP pre-trained on it that accelerates Ansor search up to 10× — the strongest published baseline. For the cross-hardware problem, TenSet tried transfer learning (predicting the gap between platforms) and Moses tried model distillation; neither cracked the problem of standing up an accurate model on new hardware from little data.
- The structural obstacle. The obvious “learn from raw source” fix fails for tensor programs: the source is deeply nested, tree-structured loop code whose AST is awkward to encode, and most tokens are irrelevant boilerplate — features come out sparse.
How TLP works
TLP inside the auto-tuning loop
-
Ansor proposes schedules
For each computational subgraph, the auto-tuner generates candidate schedule primitive sequences (initial rules + genetic-algorithm mutations). Normally each sequence would be materialized into a tensor program for feature extraction — TLP skips that entirely.
-
Featurize the primitives, not the program new
Each primitive becomes a fixed layout: its TYPE as an 11-dimensional one-hot vector, numeric parameters kept as raw numbers, character parameters (loop-variable names, annotations) mapped to tokens — exactly how NLP tokenizes words. Concatenate in order; crop/pad to sequence length 25 × embedding size 22 = 550 dense features.
-
An NLP-shaped regressor new
Linear layers upsample each primitive embedding 22→256; ONE 8-head self-attention layer (a full transformer stack is unnecessary and overfits); two residual blocks; a scalar score. Trained with lambda rank loss — the model's job is ORDERING candidates, not absolute microseconds — which beats MSE (top-1 0.919 vs 0.913).
-
The tuning loop, accelerated
The genetic algorithm consults TLP thousands of times per round; because featurization needs no code generation, model time drops ~1.7–1.8× vs TenSet MLP (five GA rounds: ~20 s → ~6 s). Each round, the top 10 predicted candidates get real measurements; 200 rounds = 2,000 measurements per workload.
-
MTL-TLP: one body, many heads new
To serve a NEW platform cheaply: train one model with multi-task learning, where each task is a hardware platform. Shared shallow layers absorb hardware-INDEPENDENT structure (what makes schedules intrinsically good); per-task heads absorb hardware-DEPENDENT behavior. Recipe: 2–3 tasks, auxiliary tasks using ALL data from same-ISA platforms, target task with ≥500K local samples (~7% of a full dataset).
The idea
The recipe describes the dish better than a photo of it. The central insight is representational: a tensor program and the schedule-primitive sequence that generated it are informationally near-equivalent — same subgraph + same primitives = same program, and empirically only 1.04% of 8.65M programs share a primitive sequence — but they are wildly different as inputs to a neural network. The program is a nested tree bloated with boilerplate tokens; the primitive sequence is a short, regular, dense sequence of typed commands whose numeric parameters (tile sizes, unroll factors) are exactly the quantities that determine performance. By predicting from the recipe, TLP gets three wins at once: no expert feature design (the featurization is a mechanical one-hot + raw numbers + tokens), no CPU/GPU feature divergence (primitives look the same everywhere), and no code generation in the inner loop (1.7–1.8× faster model calls). The “language” framing is earned, not decorative — sequences of typed words with parameters, where order matters, is precisely the input shape self-attention was built for.
Rank, don’t measure. The model’s downstream job is picking which dozens of candidates deserve real measurement — an ordering problem. TLP embraces this with lambda rank loss, and the architecture stays deliberately small: one attention layer, not a transformer stack (the authors tried GPT- and BERT-style pretraining and found the parameter count a liability for 550-feature inputs). Even the feature cropping helps accuracy: shrinking from the maximal 54×40 layout to 25×22 removes mostly zeros, and denser features train better.
Split the model where the hardware splits the problem. The cross-hardware fix mirrors the compiler’s own architecture. Tensor compilers separate hardware-independent graph IR from hardware-dependent tensor IR; MTL-TLP separates shared layers (structure that transfers: which schedules are intrinsically well-formed) from per-platform heads (what THIS machine rewards). Head-to-head at 500K target samples, multi-task learning beats fine-tuning (top-1 0.833 vs 0.790) and beats self-supervised GPT/BERT pretraining outright — and the ISA experiment confirms the mental model: auxiliary data from same-ISA platforms (Intel helping Intel) transfers best, ARM helping Intel least.
Tuning ResNet-50, step by step
One workload through TLP-accelerated Ansor
Subgraph = one fused operator group of the network (ResNet-50 yields ~29). Round = one tuning iteration per subgraph: GA search → model scoring → 10 real measurements. Top-1/top-5 score = achieved-best latency as a fraction of dataset-optimal when trusting the model's first 1 or 5 picks. Numbers from the paper's CPU (i7-10510U) and GPU (Tesla T4) experiments.
Step 1 of 7
The network shatters into search problems
ResNet-50 is partitioned into computational subgraphs, each with its own schedule search space — millions of variants on CPU, billions on GPU. Timing them all at hundreds of milliseconds each would take years; the budget is 2,000 real measurements total.
search: genetic algorithm over primitive sequences ground truth: 10 measurements per subgraph-round
Why it works
- A schedule primitive sequence near-uniquely determines its tensor program (1.04% collisions) but is a short, dense, regular sequence rather than a sparse nested tree → it can be featurized mechanically, with zero architecture expertise, identically for CPUs and GPUs → the expert-feature-engineering bottleneck disappears rather than being optimized.
- The parameters that dominate performance (tile sizes, unroll factors, vectorization choices) appear literally as numbers in the primitives → the model sees causes, not symptoms → a single attention layer suffices where program-based models needed thousands of features plus AST machinery.
- The tuner needs a ranking, not a stopwatch → lambda rank loss optimizes exactly the decision being made (which candidates to measure) → higher top-k scores than MSE at identical architecture.
- Skipping code generation moves the model OFF the critical path’s most expensive step → 1.7–1.8× faster scoring compounds over the tens of thousands of GA evaluations per workload.
- Some schedule quality is intrinsic and some is machine-specific → multi-task learning factors the model the same way the compiler factors its IRs (shared layers ↔ graph IR, heads ↔ tensor IR) → auxiliary platforms’ abundant data teaches the shared part, so the target platform needs only enough data to calibrate its head — 7%, with same-ISA auxiliaries the measurably best teachers.
What it costs
- A 500K-sample toll per new platform. MTL-TLP shrinks the data requirement 14× but doesn’t eliminate it — below ~500K target samples accuracy drops off, and collecting 500K measurements is still days of machine time.
- Tied to primitive-based tuners. The representation exists only where programs are derived via schedule primitives (Ansor-style frameworks). Compilers that transform code directly — or search spaces expressed differently — can’t use TLP as-is.
- Framework-specific vocabulary. The one-hot type encoding and parameter layouts are per-auto-tuner; a new framework means a new (if mechanical) featurization, and sequences from different frameworks aren’t mutually intelligible.
- Rare collisions are mispriced. When distinct programs share a primitive sequence (~1%), TLP must assign one label; taking the optimum guarantees no good schedule is missed but can let a poor program masquerade as good, costing some tuning time.
- GPU memory up, CPU time down. Scoring moves work to the GPU: memory rises from 882MB to 1.6GB at batch 2048 — trivial for real rigs, but a real shift in resource profile.
- Ranking quality is dataset-bounded. TLP trains on TenSet-style datasets generated by Ansor itself; schedules outside that distribution (new primitive types, exotic hardware) inherit no guarantees.
The evaluation
Two evaluation modes. Dataset-based: top-1/top-5 scores on a held-out test set of five networks (ResNet-50, MobileNet-V2, ResNext-50, BERT-tiny, BERT-base) across TenSet’s platforms — where TLP’s architecture choices are ablated (attention vs LSTM, rank vs MSE loss, feature cropping, head counts, layer counts). Search-based — the gold standard: wall-clock time for TLP-integrated Ansor to reach reference tuning quality on a real CPU (Intel i7-10510U) and GPU (NVIDIA Tesla T4), 200 rounds × 10 measurements per workload.
Cost models for tensor program search
| Property | Ansor (built-in model, no pretraining) | TenSet MLP (hand features + 51.57M-program pretraining) | TLP / MTL-TLP |
|---|---|---|---|
| Feature source | |||
| Model & loss | MLP trained online during search | MLP, MSE-style regression | |
| Search time to match reference quality | Baseline: 2,000 measurements per workload | ||
| Standing up a new hardware platform | Tune from scratch every time | ||
| Scoring cost in the tuning loop | Feature extraction requires generating each tensor program | Same requirement |
Check yourself
Check your understanding
1. Prior cost models featurized the generated tensor program; TLP featurizes the schedule-primitive sequence that produced it. Predict why the recipe beats the artifact as a neural network input.
2. TLP CROPS its features from the maximal 54×40 layout down to 25×22 — discarding representable information — and accuracy IMPROVES. Why?
3. For standing up a cost model on new hardware with 500K local samples, the paper compares fine-tuning, GPT/BERT-style self-supervised pretraining, and multi-task learning. MTL wins (top-1 0.833 vs 0.790 for fine-tuning). What makes it the right fit?
4. Why does the whole enterprise hinge on the cost model rather than on a smarter search algorithm?
Sources
- arXiv preprint 2211.03578 (v2) — full text used for this explainer
- ACM DL: TLP: A Deep Learning-Based Cost Model for Tensor Program Tuning — ASPLOS 2023
- TLP code on GitHub