[Paper Review] Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions

Nicolas Vasilache, Oleksandr Zinenko+7 more|arXiv (Cornell University)|Feb 13, 2018

Computational Physics and Python Applications61 references252 citations

TL;DR

Introduces Tensor Comprehensions (TC), a domain-specific language and polyhedral JIT compiler that generate high-performance CUDA kernels for ML workloads, enabling cross-framework integration and operator fusion with autotuning.

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]

Motivation & Objective

• Provide a concise, math-friendly language to express tensor computations for ML.
• Lower tensor comprehensions to optimized GPU code with memory management and scheduling.”
• Enable cross-framework integration by linking TC with PyTorch and Caffe2.
• Offer automatic tuning via an autotuner to explore optimization opportunities.
• Show competitive performance and practical integration with existing ML frameworks.

Proposed method

• Define Tensor Comprehensions (TC), a notation close to Einstein summation for tensor operations.
• Develop a polyhedral Just-In-Time (JIT) compiler to lower TC to CUDA kernels with memory management and synchronization.
• Create a tailored polyhedral optimization flow with kernel fusion and size-specific specialization.
• Implement an autotuning framework leveraging JIT compilation and a compilation cache.
• Integrate TC with PyTorch and Caffe2 via an in-process interface and ATen asynchronous tensor library.

Experimental results

Research questions

• RQ1 Can TC express common and custom ML operators concisely and safely?
• RQ2 Does the TC-based flow achieve performance competitive with vendor libraries on GPUs for standard ML kernels?
• RQ3 Can a polyhedral-based optimizer effectively fuse operators and optimize for data size/shape in deep learning models?
• RQ4 How well can TC integrate with mainstream frameworks like PyTorch and Caffe2 and provide a production-oriented pathway?
• RQ5 What is the impact of autotuning and JIT compilation on achieving hardware-specific performance gains?

Key findings

• The TC flow yields up to 4x speedup over NVIDIA libraries on kernels relevant to ML workloads.
• End-to-end TC integration with Caffe2 and PyTorch demonstrates practical, production-oriented applicability.
• A specialized polyhedral compiler provides effective optimizations for long-dependency deep learning kernels.
• Autotuning and code caching enable domain-specific optimization for non-standard sizes and layouts.
• The framework supports safe in-place updates and simple, declarative definitions for common layers (e.g., SGEMM, conv2d, maxpool).
• The system remains framework-agnostic at the core while offering tight integration with production environments.