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

TL;DR

This paper presents Tensor Comprehensions (TC), a domain-specific language that expresses tensor computations with Einstein-like notation, coupled with a polyhedral Just-In-Time compiler that lowers TC to optimized CUDA kernels. The approach enables cross-operator optimizations, kernel fusion, and data-layout transformations, addressing the limitations of existing ML frameworks in handling custom operators and architecture-specific optimizations. An autotuning framework with a compilation cache provides practical, production-aware performance by exploring many scheduling and mapping configurations. Integrated with PyTorch and Caffe2, TC demonstrates strong baseline performance and notable improvements on several kernels and production-style models, illustrating its potential to accelerate ML research and deployment. The work emphasizes abstraction without regret by marrying high-level declarativity with low-level, architecture-aware optimizations through a modular polyhedral pipeline.

Abstract

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]

Figures (13)

• Figure 1: Tensor Comprehension for the BLAS
• Figure 2: The JIT compilation flow lowers TC to Halide-IR, then to Polyhedral-IR, followed by optimization, code generation and execution
• Figure 3: Optimization steps for from Figure \ref{['fig:tc_sgemm']}
• Figure 4: Multithreaded autotuning pipeline for kernels
• Figure 5: Wall-clock execution of kernels (in $\mu s$). Each kernel ran 1000 times. The top half of each table is Tesla M40 (Maxwell) and the bottom half is Tesla P100 (Pascal); N/A indicates the framework lacked an implementation