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

By Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zach DeVito, William S. Moses, Sven Verdoolaege, Andrew Adams, Albert Cohen. arXiv.org, 2018

490 citations recorded by the source corpus.

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 offering both imperative and declarative styles, (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]

Subfields: Artificial Intelligence, Computational Mathematics, Hardware and Architecture

View Paper: https://www.semanticscholar.org/paper/716b8a99dd90ae046388bb2fda81602505b51b62

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mathematics

In Hardware and Architecture

Citations in the corpus

Learning to Optimize Tensor Programs (2018, methodology): “Tensor Comprehensions [40] combines both approaches, using black-box optimization to choose parameters of thread blocks and polyhedral optimization to generate internal loops.”
Fast geometric learning with symbolic matrices (2020, unknown): “Modern toolboxes accelerate a wide range of operations but are not geared towards geometric problems: most of them keep a focus on distributed learning [63, 64, 97, 108] or image processing and dense tensor manipulations [22, 58, 74, 105].”
HAWQV3: Dyadic Neural Network Quantization (2020, background): “: A number of frameworks [4, 8, 9, 19, 31, 41, 45, 49] have been developed for deep learning.”
Ansor : Generating High-Performance Tensor Programs for Deep Learning (2020, background): “Given the importance of DNNs’ performance, researchers and industry practitioners have turned to search-based compilation [2, 11, 32, 49, 59] for automated generation of tensor programs, i.”
Triton: an intermediate language and compiler for tiled neural network computations (2019, methodology): “• The polyhedral model [18] has been used by Tensor Comprehensions (TC) [43] and Diesel [14] to parameterize and automate the compilation of one or many DNN layers into LLVM-IR and CUDA-C programs.”
Machine Learning Systems are Stuck in a Rut (2019, methodology): “In popular frameworks like TensorFlow [23] and PyTorch [3], user programs are written in Python and call into operators10 that are written in terms of back end-specific languages and libraries such as C++, MKL [24], CUDA [25], and cuDNN [10], or sometimes lower-level but portable domain-specific lan”
Glow: Graph Lowering Compiler Techniques for Neural Networks (2018, methodology): “Finally, Tensor Comprehensions [18] provides a language for neural network developers to specify their networks such that a JIT compiler can algorithmically search for the most eﬃcient execution plan possible.”
Domain-specific Genetic Algorithm for Multi-tenant DNNAccelerator Scheduling (2021, methodology): “STOKE [92] and Tensor Comprehensions [113] use GA to search the space of DNN code optimization.”
TVM: An Automated End-to-End Optimizing Compiler for Deep Learning (2018, methodology): “We include the TensorComprehension (TC, commit: ef644ba) [37] a recently introduced auto-tuning framework as an additional baseline.”
The Deep Learning Compiler: A Comprehensive Survey (2020, background): “Tensor Comprehensions [26], Glow [27], nGraph [28] and”
Efficient Deep Learning: A Survey on Making Deep Learning Models Smaller, Faster, and Better (2021, methodology): “The alternatives for XLA in the PyTorch world seem to be the Glow [115] and TensorComprehension [137] compilers.”
Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code (2018, methodology): “Some of them are designed for speciﬁc domains (such as Tensor Comprehensions and PolyMage), while Pluto, PENCIL, and Polly are more general.”
Flex-Convolution - Million-Scale Point-Cloud Learning Beyond Grid-Worlds (2018, methodology): “We did some experiments with a quite recent polyhedral compiler optimization using TensorComprehension (TC) [27] to automatically tune a flexconvolution layer implementation.”
TVM: An Automated End-to-End Optimizing Compiler for Deep Learning (2018, methodology): “We include TensorComprehension (TC, commit: ef644ba) [42], a recently introduced auto-tuning framework, as an additional baseline.”
Neural Kernels Without Tangents (2020, methodology): “We implemented all the convolutional kernels in the tensor comprehensions framework (Vasilache et al., 2018) and executed them on V100 GPUs using Amazon Web Services (AWS) P3.”
Nimble: Lightweight and Parallel GPU Task Scheduling for Deep Learning (2020, methodology): “For example, DL compilers [7, 15, 18, 31, 38] have been proposed to generate optimized codes for target hardware.”
Supporting Very Large Models using Automatic Dataflow Graph Partitioning (2018, methodology): “Concurrent with our work, TVM [23] and TC [24] are two other Halide-inspired DSLs.”
Deep Learning Inference in Facebook Data Centers: Characterization, Performance Optimizations and Hardware Implications (2018, background): “Automatic generation of fused kernels is an active area of research and early productization efforts are underway [10, 42, 57, 67].”