Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions (1802.04730v3)

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]

An Analysis of Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions

The paper presents Tensor Comprehensions (TC), which are described as a framework-agnostic and high-performance system for expressing deep learning operations succinctly while optimizing their execution on heterogeneous hardware, particularly focusing on NVIDIA GPUs. The key contributions revolve around designing a domain-specific language (DSL) tailored for representing tensor operations efficiently, enhanced by the capabilities of a polyhedral compiler that transforms these high-level expressions into optimized CUDA kernels.

Core Contributions and Approach

The authors identify significant performance overheads associated with custom operator development in popular machine learning frameworks such as TensorFlow, PyTorch, and others. These frameworks often rely on backend libraries like CUDNN for performance-critical operations, but when new operators are required, developers face a high engineering cost and potential performance deficits. Tensor Comprehensions address this by simplifying the expression of tensor operations while leveraging a Just-In-Time (JIT) compilation environment coupled with automated memory management and kernel optimization strategies using polyhedral techniques.

Key contributions of this work include:

1. Introduction of Tensor Comprehensions Language: This DSL provides a way to encode deep learning mathematical expressions analogous to Einstein notation with automatic shape and size inference. It ensures operations are defined without explicit looping constructs, making programs concise and less error-prone.
2. Polyhedral Compiler Technology: The authors employ a polyhedral compiler approach to automatically generate efficient GPU code, enabling optimizations such as loop tiling, fusion, and parallelization at various levels of the hardware memory hierarchy. This strategy facilitates achieving significant performance improvements over standard reference libraries.
3. Integrated Autotuning Framework: A vital aspect of this work is its autotuning framework. This element systematically searches through possible kernel configurations to identify and implement performance-optimized code structures. The autotuning process is buttressed by a compilation cache, enhancing efficiency by reusing optimized configurations.
4. Framework Integration: The paper demonstrates the seamless integration of TC with Caffe2 and PyTorch frameworks, showcasing the ability to enhance both production-oriented and research-focused machine learning tasks. This adaptability indicates a powerful flexibility within the Tensor Comprehensions system.

Experimental Evaluation

The results presented reveal TC’s potential to surpass highly tuned GPU libraries like CUBLAS in numerous scenarios, often through aggressive loop parallelization and data residency in the GPU memory hierarchy. In particular, Tensor Comprehensions achieve a four-fold speedup over NVIDIA libraries on specific kernels relevant to Facebook’s production models.

Implications and Future Directions

The implications of this work extend to both theoretical and practical domains of AI and machine learning. By presenting a system that can automatically derive optimized, low-level GPU code from high-level tensor specifications, the authors provide a substantial efficiency gain for machine learning developers.

Future research can extend this methodology to support a broader spectrum of processors and specialized hardware accelerators. There are also opportunities to refine the automation and tuning processes, potentially incorporating more advanced machine learning techniques for optimization suggestion. Furthermore, system-level extensions could address dynamic computing graphs or introduce support for sparsity and mixed-precision data types directly at the TC level.

In conclusion, Tensor Comprehensions present a compelling advancement in the efficiency and expressiveness of machine learning systems, promising to streamline both research explorations and the deployment of machine learning models at an industrial scale. The framework’s optimization techniques and integration strategy exemplify the potential of domain-specific languages in high-performance computing contexts.