# Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions

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

## Key facts

- Cited by 85 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 13 Feb 2018.

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

- [Efficient Deep Learning: A Survey on Making Deep Learning Models Smaller, Faster, and Better](https://mlsanity.com/paper/2106.08962/) (2021)
- [Full Stack Optimization of Transformer Inference: a Survey](https://mlsanity.com/paper/2302.14017/) (2023)
- [DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion](https://mlsanity.com/paper/2108.13342/) (2021)
- [TensorIR: An Abstraction for Automatic Tensorized Program Optimization](https://mlsanity.com/paper/2207.04296/) (2022)
- [ThunderKittens: Simple, Fast, and Adorable AI Kernels](https://mlsanity.com/paper/2410.20399/) (2024)
- [A Deep Learning Based Cost Model for Automatic Code Optimization](https://mlsanity.com/paper/2104.04955/) (2021)
- [Tensor Program Optimization with Probabilistic Programs](https://mlsanity.com/paper/2205.13603/) (2022)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [Mind Mappings: Enabling Efficient Algorithm-Accelerator Mapping Space Search](https://mlsanity.com/paper/2103.01489/) (2021)
- [CoSA: Scheduling by Constrained Optimization for Spatial Accelerators](https://mlsanity.com/paper/2105.01898/) (2021)
- [The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding](https://mlsanity.com/paper/2110.10221/) (2021)
- [Tensor and Matrix Low-Rank Value-Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/2201.09736/) (2022)
- [TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs](https://mlsanity.com/paper/2308.13490/) (2023)
- [Allo: A Programming Model for Composable Accelerator Design](https://mlsanity.com/paper/2404.04815/) (2024)
- [MAGMA: An Optimization Framework for Mapping Multiple DNNs on Multiple Accelerator Cores](https://mlsanity.com/paper/2104.13997/) (2021)
- [TLP: A Deep Learning-based Cost Model for Tensor Program Tuning](https://mlsanity.com/paper/2211.03578/) (2022)
- [UNIT: Unifying Tensorized Instruction Compilation](https://mlsanity.com/paper/2101.08458/) (2021)
- [HASCO: Towards Agile HArdware and Software CO-design for Tensor Computation](https://mlsanity.com/paper/2105.01585/) (2021)
- [Tensor Processing Primitives: A Programming Abstraction for Efficiency and Portability in Deep Learning & HPC Workloads](https://mlsanity.com/paper/2104.05755/) (2021)
- [Operator Fusion in XLA: Analysis and Evaluation](https://mlsanity.com/paper/2301.13062/) (2023)
- [GraphRARE: Reinforcement Learning Enhanced Graph Neural Network with Relative Entropy](https://mlsanity.com/paper/2312.09708/) (2023)
- [Neural Architecture Search as Program Transformation Exploration](https://mlsanity.com/paper/2102.06599/) (2021)
- [Horizontally Fused Training Array: An Effective Hardware Utilization Squeezer for Training Novel Deep Learning Models](https://mlsanity.com/paper/2102.02344/) (2021)
- [Decompiling x86 Deep Neural Network Executables](https://mlsanity.com/paper/2210.01075/) (2022)
- [oneDNN Graph Compiler: A Hybrid Approach for High-Performance Deep Learning Compilation](https://mlsanity.com/paper/2301.01333/) (2023)
- [LOOPer: A Learned Automatic Code Optimizer For Polyhedral Compilers](https://mlsanity.com/paper/2403.11522/) (2024)
- [Functional Collection Programming with Semi-Ring Dictionaries](https://mlsanity.com/paper/2103.06376/) (2021)
- [A Reinforcement Learning Environment for Polyhedral Optimizations](https://mlsanity.com/paper/2104.13732/) (2021)
- [Bring Your Own Codegen to Deep Learning Compiler](https://mlsanity.com/paper/2105.03215/) (2021)
- [Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures](https://mlsanity.com/paper/2304.12576/) (2023)

## Abstract (arXiv)

Deep learning models with convolutional and recurrent networks are now ubiquitous and analyze massive amounts of audio, image, video, text and graph data, with applications in automatic translation, speech-to-text, scene understanding, ranking user preferences, ad placement, etc. Competing frameworks for building these networks such as TensorFlow, Chainer, CNTK, Torch/PyTorch, Caffe1/2, MXNet and Theano, explore different tradeoffs between usability and expressiveness, research or production orientation and supported hardware. They operate on a DAG of computational operators, wrapping high-performance libraries such as CUDNN (for NVIDIA GPUs) or NNPACK (for various CPUs), and automate memory allocation, synchronization, distribution. Custom operators are needed where the computation does not fit existing high-performance library calls, usually at a high engineering cost. This is frequently required when new operators are invented by researchers: such operators suffer a severe performance penalty, which limits the pace of innovation. Furthermore, even if there is an existing runtime call these frameworks can use, it often doesn't offer optimal performance for a user's particular network architecture and dataset, missing optimizations between operators as well as optimizations that can be done knowing the size and shape of data. Our contributions include (1) a language close to the mathematics of deep learning called Tensor Comprehensions, (2) a polyhedral Just-In-Time compiler to convert a mathematical description of a deep learning DAG into a CUDA kernel with delegated memory management and synchronization, also providing optimizations such as operator fusion and specialization for specific sizes, (3) a compilation cache populated by an autotuner. [Abstract cutoff]