TensorIR: An Abstraction for Automatic Tensorized Program Optimization

Modern chips compute whole matrix tiles in one instruction, but compilers still think in scalar loops — TensorIR adds a "block" construct that walls off a tensor computation behind a signature, so loops outside can be transformed automatically while the hardware instruction lives inside.

Why this problem is worth solving

Hardware acceleration for machine learning has climbed a ladder of granularity. Scalar units compute one multiply-add at a time. SIMD vector units compute a handful in lockstep. The current rung is different in kind: tensor intrinsics — single instructions (or expert-written micro-kernels) that consume and produce whole multi-dimensional tiles. An NVIDIA Tensor Core mma instruction performs a full 16×16×16 matrix multiply; Google’s TPU is built around such units; ARM CPUs expose dot-product and interleaved micro-kernels. Peak throughput on modern chips is these instructions — miss them and the machine idles.

Using them is the problem the paper names tensorized program optimization: given an operator (a convolution, an attention layer), find the program that tiles its loops, arranges its threads, stages its data through the memory hierarchy, and maps the innermost computation onto the hardware’s tile instruction — with the right layouts, memory scopes, and load/store companions that each intrinsic demands.

Today this is done by hand. Vendor libraries (cuDNN, MKL-DNN, ARM Compute Library) encode expert answers per operator per chip — enormous engineering effort that lags every new model architecture and every new backend. The alternative, automatic compilation, had a hole in the middle: existing compiler search worked below the tensor-instruction level, not around it.

What people did before, and where it broke

Machine-learning compilers came in two shapes. Bottom-up systems — Halide, TVM, Tiramisu, AKG — represent an operator as loop nests around a scalar body (C[y,x] += A[y,k] * B[k,x]) and search over loop transformations: split, reorder, fuse, cache, bind to threads. Their dependency analysis reads the scalar statements. Top-down systems — Fireiron, Stripe — decompose the problem hierarchically through nested structures. Both automated real optimization, and both stumble on tensor intrinsics: an opaque instruction has no scalar body to analyze, so the compiler cannot verify that transforming the loops around it preserves meaning, cannot automatically decide where it applies, and cannot co-derive the special data movements (load-matrix instructions, interleaved layouts, shared-memory staging) each intrinsic requires.

Meanwhile, watch what a human expert actually does with, say, a 4×4 matmul instruction: divide and conquer. Split the 64×64×64 multiplication into an outer 16×16×16 grid of 4×4×4 sub-problems; optimize the outer loops (tiling, threading, caching) as one concern; implement the inner 4×4×4 with the hardware instruction as another. The two concerns meet only at an interface: what region the inner part reads and writes. Prior compilers had no way to express that interface. TensorIR’s contribution is precisely to make it a language construct.

The idea

The block: a computation with a signature

A block wraps a computation over a sub-region of buffers and carries a signature with exactly the information the outside world needs:

• Block iterators with domains and kinds — spatial axes versus reduction axes — and their current bindings to outer loop variables.
• Access regions: which rectangles of which buffers the block reads and writes (for the 4×4 matmul block: reads a 4×4 tile of A and of B, writes a 4×4 tile of C).

The signature is an isolation boundary. Loop transformations outside the block — tiling, reordering, moving a producer to a consumer’s loop level — compute their legality from the signature alone, never inspecting the body. That is what makes an opaque body (a Tensor Core instruction) exactly as transformable as a scalar one. Two design subtleties carry weight. Iterator kinds encode constraints: a reduction axis must not be bound to a parallel loop (unless atomic). And dependencies are recorded against buffers, not between statements — an indirection that legalizes layout transformations and recomputation, both essential around intrinsics.

Blocks nest: a schedulable block (loops with sub-blocks as leaves) can contain an opaque sub-block, and vice versa — the outer search proceeds while the inner instruction stays sealed. Blockization is the transformation that manufactures this structure, carving an inner loop nest into a new signed sub-block — the compiler’s version of the expert’s “divide.”

Scheduling as program-to-program transformations

Prior compilers accumulate scheduling decisions in a separate tree lowered at the end; adding a primitive means touching the tree and its lowering rules. TensorIR instead makes every scheduling primitive a standalone transformation from one TensorIR program to another — tiling, thread binding, caching into shared memory, reduction init/update splitting, blockization. Any intermediate program can be printed, inspected, hand- edited, or mixed with automated rewriting. The rich space of prior compilers’ loop transformations is reproduced — now with tensorized bodies inside.

Validation: correctness as a filter

With search generating programs, correctness must be checked, not assumed. TensorIR validates at three levels: loop-nest validation pattern-matches quasi-affine bindings to confirm iterator independence and domain coverage (binding two data-parallel iterators as i and i*2 is caught; i/4 and i%4 is legal); producer-consumer validation confirms writers cover readers’ regions; threading validation checks thread-binding consistency, cooperative shared-memory coverage across a thread group, and execution scope (a Tensor Core intrinsic must sit at warp level). These checks serve users hand-writing programs and — critically — filter false candidates during automatic search.

The tensorization-aware auto-scheduler

Automation proceeds in three moves. Candidate generation matches an operator against registered intrinsics — first the expression pattern (C += A*B), then the indices, using characteristic vectors that fingerprint how each iterator touches each buffer; this is how a 2-D convolution’s seven tangled iterators get remapped (via ReIndex and layout- transformation blocks) onto a 16×16×16 matmul intrinsic’s three. Sketch generation tiles the program so an intrinsic-shaped sub-block exists, then inserts AutoCopy blocks — data-movement stages that know only their source, destination and scope, and get scheduled independently with their own vectorization and layout decisions, absorbing the load-matrix instructions and interleaved layouts real intrinsics demand. Evolutionary search with a learned cost model fills in the remaining choices (tile sizes, thread mappings), with validation discarding illegal programs.

semantics:  for y,x,k in grid(4,4,4): C[y,x] += A[x,k]*B[k,y]
impl:       accel.dot(...)  + scope/layout constraints

Why it works

Because the block signature carries exactly the dependency information that loop transformations need — iterator domains, kinds, and access regions — the outside of a computation can be optimized without reading its inside, making opaque hardware instructions as schedulable as scalar code. → Because intrinsics are declared as semantics-plus-implementation block pairs, matching them is program analysis rather than per-backend compiler engineering — one declaration unlocks automatic use everywhere the pattern fits, index remappings included. → Because data movement materializes as independently-scheduled AutoCopy blocks, the awkward companions of every intrinsic — special loads, layouts, memory scopes — enter the same search space as the compute. → Because validation is mechanical, the enlarged search space stays safe: evolutionary exploration can be aggressive precisely because illegal candidates are provably filtered. → Therefore automatic compilation reaches where only vendor libraries lived: up to 7.5× over prior ML compilers on Tensor Core operators, 88–100% of TensorRT end to end, and native dot-product acceleration on ARM that the framework baseline simply had not gotten around to supporting.

What it costs

• Someone must declare the intrinsics. The semantics/implementation pair with layout and scope constraints is per-backend expert work — drastically less than a kernel library, but not zero.
• Search time is real. Auto-scheduling tunes for minutes to hours per model (up to 2× faster than TVM’s tuner at better quality, per the paper’s measurements) — amortized in deployment, felt in iteration.
• Not universally fastest. Against relentlessly-tuned vendor code, TensorIR wins on some operators (up to 13.9×) and reaches 75–105% on others (dense GEMM, 3-D convolution, ARM Compute Library’s home turf) — parity with experts, not uniform dominance.
• The abstraction has edges. Blocks generalize loop-nest programs; computation that resists region-structured signatures (dynamic shapes, data-dependent control flow) needs further work — as does the sparse world, which the same group addressed by building SparseTIR on top of TensorIR (covered elsewhere on this site).
• Ecosystem migration. Adopting the abstraction means adopting its compiler stack (it became TVM’s core IR) — the cost of any IR-level idea.

The evaluation

Three platforms, three altitudes. Single operators on an NVIDIA RTX 3080 (FP16 Tensor Cores): nine workloads spanning convolutions (1-D/2-D/3-D, depthwise, dilated, grouped, transposed) and matrix multiplication. End-to-end models on the same GPU: ResNet-50, MobileNet V2, BERT-large, Vision Transformer. And a portability test: quantized models on an AWS Graviton2 ARM CPU with native dot-product instructions.