Source-linked AI summary
TensorIR: An Abstraction for Automatic Tensorized Program Optimization
Siyuan Feng, Bohan Hou, Hongyi Jin, Wuwei Lin, Junru Shao, Ruihang Lai, Zihao Ye, Lianmin Zheng, Cody Hao Yu, Yong Yu, Tianqi Chen
TL;DR
Machine-learning compilers have largely searched loop transformations without automatically handling tensorized programs, as deployment spans growing sets of models and devices. TensorIR introduces a block-based abstraction and automatic scheduling for tensorized programs; its end-to-end framework outperforms existing machine-learning compilation solutions by up to 7x and is competitive with heavily optimized platform-specific systems.
Problem
Earlier machine-learning compilation work searched loop-nest transformations but did not automatically handle tensorized programs, amid growing demands to deploy diverse models across devices.
Source: §1 · pp. 1–2Method
TensorIR uses blocks to isolate tensorized computations from loop nests and combines transformation primitives with a tensorization-aware automatic scheduler.
Source: §1 · p. 2Results
TensorIR outperforms existing machine-learning compilation solutions by up to 7x and delivers competitive performance against heavily optimized platform-specific systems.
Source: §1 · p. 2Takeaways & Limitations
TensorIR can automatically support emerging models such as Vision Transformer, which TensorRT does not yet support.
Source: §5.2 · p. 11The implementation uses a default order for fusion, with generalization to different fusion orders left for future work.
Source: §4.2 · p. 9
Abstract
from arXiv · showhide
Deploying deep learning models on various devices has become an important topic. The wave of hardware specialization brings a diverse set of acceleration primitives for multi-dimensional tensor computations. These new acceleration primitives, along with the emerging machine learning models, bring tremendous engineering challenges. In this paper, we present TensorIR, a compiler abstraction for optimizing programs with these tensor computation primitives. TensorIR generalizes the loop nest representation used in existing machine learning compilers to bring tensor computation as the first-class citizen. Finally, we build an end-to-end framework on top of our abstraction to automatically optimize deep learning models for given tensor computation primitives. Experimental results show that TensorIR compilation automatically uses the tensor computation primitives for given hardware backends and delivers performance that is competitive to state-of-art hand-optimized systems across platforms.
1 INTRODUCTION
Machine-learning deployment spans growing sets of models and devices, while specialized tensor instructions make automatic tensorized program optimization a challenge. TensorIR combines an abstraction, transformation primitives, and an automatic scheduler, and its end-to-end framework outperforms existing machine-learning compilation solutions by up to 7x while delivering competitive performance against heavily optimized platform-specific systems.
1 INTRODUCTION
- Growing model and device ranges make deployment challenging, while expert-built kernel libraries require substantial engineering and iteration to keep pace.Such libraries combine tensorized primitives with loops, threading, and caching.Source: §1 · pp. 1–2
- Modern accelerators and micro-kernels use tensor intrinsics on multidimensional regions, so optimization must combine hierarchical loops, multidimensional loads, and tensor operations.The paper calls this task tensorized program optimization; earlier machine-learning compilation work searched loop transformations but did not automatically handle tensorized programs.Source: §1 · pp. 1–2
- TensorIR’s block abstraction isolates tensorized computation from loop transformations while representing tensor intrinsics and hardware constraints uniformly.It also represents mixtures of irregular and tensor computations, expanding beyond normal tensor expressions [9].Source: §1 · p. 2
- TensorIR provides transformation primitives for a rich optimization search space and builds a tensorization-aware automatic scheduler.The scheduler performs tensorization jointly with other optimizations.Source: §1 · p. 2
- TensorIR outperforms existing machine-learning compilation solutions by up to 7x and delivers competitive performance against heavily optimized platform-specific systems.The paper reports this result for its end-to-end compilation framework.Source: §1 · p. 2
2 OVERVIEW
TensorIR brings the divide-and-conquer approach used by expert developers into machine-learning compiler design. Its block abstraction isolates tensorized computations from surrounding loops, while automatic search explores how to use hardware primitives and optimize the resulting program.
2 OVERVIEW
- Expert developers divide operators into tensorized subproblems and surrounding loops, then optimize the two separately.This divide-and-conquer approach can target multiple implementations of a tensorized computation.Source: §2 · p. 2, Figure 2 · p. 3
- Existing compilers use bottom-up loop-nest transformations or top-down decomposition, prompting the question of how to support divide and conquer in compiler design.The paper identifies the need to bring manual tensorized-program optimization insights into machine-learning compilers.Source: §2 · p. 2
- TensorIR adds blocks to loop nests, isolating an inner problem from the outer problem so both can be transformed independently using the block signature.Blocks can represent hardware tensor primitives or larger subproblems, and tensor computation is a first-class citizen.Source: §2 · pp. 2–3, Figure 3 · p. 3
- The compiler searches tensorization candidates and program transformations, automating decisions common in libraries and compilation pipelines to generate competitive vendor-specific solutions.The search space generalizes prior compiler approaches and optimization techniques from HPC and machine-learning engineering.Source: §2 · pp. 3–4
3 TENSORIR ABSTRACTION
TensorIR represents computation with multidimensional buffers, loop nests, and blocks, then applies transformations and validation to optimize programs. The framework also supports importing models for automatic optimization or writing customized operators in Python.
3 TENSORIR ABSTRACTION
- TensorIR programs combine multidimensional buffers, loop nests, and blocks, whose dependency information supports transformations over computation regions.The Python-AST dialect lets developers construct, inspect, modify, and transform TensorIR programs.Source: §3 · p. 4
3.1 Block
- Block signatures isolate tensorized computation from outer loops and record iterator constraints and buffer access dependencies for transformations.For reduction computations, an optional initialization statement can coordinate scheduling of initialization and update steps.Source: §3.1 · pp. 4–5
3.2 Scheduling Transformations
- Scheduling transformations tile, reorder, and relocate loops, while blockization isolates subregions as tensorization candidates and caching adds sub-blocks for shared-memory inputs.Standalone transformations make scheduling primitives easier to extend and allow developers to inspect programs at each stage.Source: §3.2 · pp. 5–6
3.3 Validation
- Validation checks iterator bindings, producer-consumer coverage, thread assignments, cooperative memory access, and intrinsic execution scope, while primitive-specific checks preserve transformation correctness.These checks filter invalid automatically generated programs and help users identify errors in manually crafted or imported programs.Source: §3.3 · pp. 6–7
3.4 Programming Effort
- Users can import TensorFlow or PyTorch models for automatic TensorIR generation and hardware-specific optimization, or write customized operators in Python.The framework also provides optimization transformations for custom TensorIR programs.Source: §3.4 · p. 7
4 AUTO-SCHEDULING TENSORIZED PROGRAMS
The auto-scheduler takes workload and hardware intrinsic descriptions, generates tensorization and program-sketch candidates, then searches them with an evolutionary procedure guided by a learned cost model.
4.1 Abstraction for Tensor Intrinsics
- TensorIntrin describes each instruction with one block for computation semantics and another for low-level implementation, while buffer specifications encode constraints checked during validation.The specifications include data type, storage scope, memory layout, and contiguity constraints.Source: §4.1 · p. 7
4.2 Tensorization Candidate Generation
- Tensorization matching first aligns C[.] += A[.] × B[.] patterns, then uses characteristic vectors to match iterator lists, with ReIndex normalizing buffer accesses.Iterators with matching characteristic vectors are fused and mapped to intrinsic iterators.Source: §4.2 · pp. 7, 9
- Candidates rewrite operands for intrinsic-compatible bodies, then pad to divisible shapes, tile inner loops, and blockize; memory-layout and threading constraints are checked during validation.Candidate generation may produce candidates that do not ultimately tensorize because some constraints depend on later transformations.Source: §4.2 · p. 8
4.3 Tensorized Program Sketch Generation
- Sketch rules fix parts of program structure while leaving choices such as tiling and caching open; AutoCopy exposes access needs so computation scheduling can proceed independently of memory scheduling.The data-movement scheduler can apply cache stages, tensor intrinsics, vectorization, cooperative fetching, or stride padding.Source: §4.3 · p. 9
4.4 Evolutionary Search
- Evolutionary search mutates sketch-induced programs, benchmarks promising candidates, and uses evaluation data to update a learned cost model.Search starts from random choices for the generated sketches.Source: §4.4 · p. 9
- A boosting-tree cost model uses memory-access, reuse, and loop-annotation features from both block signatures and block bodies; validation rejects invalid candidates during search.The validation step addresses candidates that violate tensor-intrinsic constraints or contain invalid loop nests.Source: §4.4 · pp. 9–10
5 EVALUATION
TensorIR is evaluated on individual operators and end-to-end models across NVIDIA GPU and ARM CPU platforms, against compiler and vendor-library baselines. The results report competitive performance across these settings, alongside faster tuning and support for emerging models.
5 EVALUATION
- The evaluation compares automatic optimization without external-library calls against compiler and vendor-specific baselines across operators, end-to-end execution, and hardware platforms.The NVIDIA single-operator setup uses an RTX 3080 with Tensor Cores and float16 inputs and accumulators.Source: §5 · p. 10, §5.1 · p. 10
5.1 Single Operator Evaluation
- On NVIDIA single operators, TensorIR improves performance by up to 7.5x over machine-learning compiler baselines and outperforms vendor-library baselines on five workloads.It reaches more than 75% throughput on C3D, GMM, and GRP; CUTLASS does not support DEP, GRP, or T2D.Source: Figure 10 · p. 10, §5.1 · p. 10, Figure 11 · p. 10
5.3 ARM CPU Evaluation
- TensorIR reaches performance close to ARMComputeLib on C2D and GMM, while using native acceleration to outperform TVM.The prose reports up to 12.5× speedup over TVM and 85%–105% of ARMComputeLib throughput.Source: §5.3 · p. 11, Figure 13 · p. 11
5.2 End-to-end Model Evaluation
- TensorIR outperforms PyTorch, TVM, and AMOS on NVIDIA end-to-end models, exceeds TensorRT on MobileNet V2, and supports Vision Transformer, which TensorRT does not yet support.Figure 12 summarizes the NVIDIA model comparisons.Source: Figure 12 · p. 11, §5.2 · p. 11
- TensorIR tunes end-to-end NVIDIA models up to 2x faster than TVM, with tensorization reducing profiling time and a divide-and-conquer search narrowing the optimization space.The framework can also reuse historical cost models and search records to avoid repeating searches for already tuned operators.Source: Table 1 · p. 11, §5.2 · p. 11
5.3 ARM CPU Evaluation
- On ARM end-to-end models, TensorIR outperforms PyTorch and TVM; the evaluation notes that PyTorch’s QNNPACK backend had not yet added sdot support.The authors connect this setting to framework maintenance costs as hardware changes.Source: §5.3 · pp. 11–12, Figure 14 · p. 12
6 RELATED WORKS
TensorIR draws on loop-based tensor compilers, nested polyhedral representations, and automatic scheduling work, then extends these ideas to tensorized programs. Its automatic scheduling focuses on tensorization for domain-specific hardware acceleration.
6 RELATED WORKS
- TensorIR generalizes divide-and-conquer optimization ideas from HPC to an abstraction for automatic tensorization.The passage names CUTLASS as an example of work targeting compute-intensive linear algebra operators.Source: §6 · p. 12
- TensorIR combines insights from scalar loop-nest and nested polyhedral approaches, extending their representations to tensorized programs while focusing on automatic tensorization across platforms.The paper describes TensorIR as orthogonal to sparse tensor algebra, recursive computation, and end-to-end compiler-chain efforts.Source: §6 · p. 12
- TensorIR’s scheduler draws on learned-cost-model and template-guided search, Triton’s tile-based templates, FlexTensor’s template generation, Halide’s Monte-Carlo tree search, and Ansor’s hierarchical search [48], generalizing these approaches to tensorized computation for domain-specific hardware acceleration.Source: §6 · p. 12
- TensorIR differs from prior tensorization work by jointly performing tensorization and other optimizations within a new abstraction.The passage contrasts it with approaches including AKG, UNIT [45], and AMOS.Source: §6 · p. 12
7 CONCLUSION
TensorIR provides an abstraction for automatic tensorized program optimization, with blocks isolating tensorized computations and transformation primitives supporting optimization. Its automatic scheduler performs tensorization jointly with other optimizations.
7 CONCLUSION
- TensorIR’s block abstraction isolates tensorized computations and supplies transformation primitives for program optimization.The conclusion presents this as a key part of the abstraction for automatic tensorized program optimization.Source: §7 · p. 12
- TensorIR’s automatic scheduler jointly performs tensorization and other optimizations to generate performant programs.Source: §7 · p. 12
- The authors hope TensorIR will encourage further studies of tensorized program optimization and create opportunities for hardware and software specialization.Source: §7 · p. 12