# TensorIR: An Abstraction for Automatic Tensorized Program Optimization

arXiv: https://arxiv.org/abs/2207.04296 · HTML: https://mlsanity.com/paper/2207.04296/

## Key facts

- Cited by 38 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 49 references that are in the map.
- First posted on arXiv 9 Jul 2022.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications](https://mlsanity.com/paper/1704.04861/) (2017)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [cuDNN: Efficient Primitives for Deep Learning](https://mlsanity.com/paper/1410.0759/) (2014)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems](https://mlsanity.com/paper/1512.01274/) (2015)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code](https://mlsanity.com/paper/1804.10694/) (2018)
- [Cortex: A Compiler for Recursive Deep Learning Models](https://mlsanity.com/paper/2011.01383/) (2020)
- [UNIT: Unifying Tensorized Instruction Compilation](https://mlsanity.com/paper/2101.08458/) (2021)
- [Automatic Kernel Generation for Volta Tensor Cores](https://mlsanity.com/paper/2006.12645/) (2020)
- [A Hardware-Software Blueprint for Flexible Deep Learning Specialization](https://mlsanity.com/paper/1807.04188/) (2018)
- [Stripe: Tensor Compilation via the Nested Polyhedral Model](https://mlsanity.com/paper/1903.06498/) (2019)
- [Automatic Generation of Efficient Sparse Tensor Format Conversion Routines](https://mlsanity.com/paper/2001.02609/) (2020)

## Built on by (most cited first, 30 of 38)

- [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](https://mlsanity.com/paper/2306.00978/) (2023)
- [Towards Efficient Generative Large Language Model Serving: A Survey from Algorithms to Systems](https://mlsanity.com/paper/2312.15234/) (2023)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [Mirage: A Multi-Level Superoptimizer for Tensor Programs](https://mlsanity.com/paper/2405.05751/) (2024)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [Allo: A Programming Model for Composable Accelerator Design](https://mlsanity.com/paper/2404.04815/) (2024)
- [MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs](https://mlsanity.com/paper/2512.22219/) (2025)
- [Relax: Composable Abstractions for End-to-End Dynamic Machine Learning](https://mlsanity.com/paper/2311.02103/) (2023)
- [SmartMem: Layout Transformation Elimination and Adaptation for Efficient DNN Execution on Mobile](https://mlsanity.com/paper/2404.13528/) (2024)
- [Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Training](https://mlsanity.com/paper/2302.08005/) (2023)
- [Improving Parallel Program Performance with LLM Optimizers via Agent-System Interfaces](https://mlsanity.com/paper/2410.15625/) (2024)
- [QiMeng-TensorOp: Automatically Generating High-Performance Tensor Operators with Hardware Primitives](https://mlsanity.com/paper/2505.06302/) (2025)
- [CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolution](https://mlsanity.com/paper/2608.12629/) (2026)
- [Car-GS: Addressing Reflective and Transparent Surface Challenges in 3D Car Reconstruction](https://mlsanity.com/paper/2501.11020/) (2025)
- [Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs](https://mlsanity.com/paper/2510.08726/) (2025)
- [Event Tensor: A Unified Abstraction for Compiling Dynamic Megakernel](https://mlsanity.com/paper/2604.13327/) (2026)
- [PyGraph: Robust Compiler Support for CUDA Graphs in PyTorch](https://mlsanity.com/paper/2503.19779/) (2025)
- [Tilus: A Tile-Level GPGPU Programming Language for Low-Precision Computation](https://mlsanity.com/paper/2504.12984/) (2025)
- [LoRAFusion: Efficient LoRA Fine-Tuning for LLMs](https://mlsanity.com/paper/2510.00206/) (2025)
- [Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers](https://mlsanity.com/paper/2601.19092/) (2026)
- [Minuet: Accelerating 3D Sparse Convolutions on GPUs](https://mlsanity.com/paper/2401.06145/) (2023)
- [Hexcute: A Compiler Framework for Automating Layout Synthesis in GPU Programs](https://mlsanity.com/paper/2504.16214/) (2025)
- [vMCU: Coordinated Memory Management and Kernel Optimization for DNN Inference on MCUs](https://mlsanity.com/paper/2406.06542/) (2024)
- [Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference](https://mlsanity.com/paper/2605.11581/) (2026)
- [ALT: Boosting Deep Learning Performance by Breaking the Wall between Graph and Operator Level Optimizations](https://mlsanity.com/paper/2210.12415/) (2022)
- [PowerFusion: A Tensor Compiler with Explicit Data Movement Description and Instruction-level Graph IR](https://mlsanity.com/paper/2307.04995/) (2023)
- [GeoT: Tensor Centric Library for Graph Neural Network via Efficient Segment Reduction on GPU](https://mlsanity.com/paper/2404.03019/) (2024)
- [FTuner: A Fast Dynamic Shape Tensors Program Auto-Tuner for Deep Learning Compilers](https://mlsanity.com/paper/2407.21418/) (2024)
- [ELK: Exploring the Efficiency of Inter-core Connected AI Chips with Deep Learning Compiler Techniques](https://mlsanity.com/paper/2507.11506/) (2025)
- [PerfDojo: Automated ML Library Generation for Heterogeneous Architectures](https://mlsanity.com/paper/2511.03586/) (2025)

## Abstract (arXiv)

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.