PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S. Moses, Sven Verdoolaege, Andrew Adams, Albert Cohen

arXiv:1802.04730v3cs.PLcs.LG

TL;DR

Deep-learning frameworks can require expensive custom operators and may miss optimizations tied to operator context, data shape, and hardware. TC addresses this with a mathematical tensor language, a polyhedral JIT compiler, and autotuned code caching; reported results include 1.4×–3.6× speedups over Caffe2 on time-domain convolutions, while important feature and optimization limitations remain.

  • Problem

    Existing computation-graph frameworks can require costly custom operators and miss optimizations across operators or for specific data sizes and shapes.

  • Method

    TC combines a mathematical tensor language with polyhedral JIT compilation, domain-specific optimizations, autotuning, and compilation caching to generate GPU kernels.

  • Results

    TC outperforms Caffe2 by 1.4× to 3.6× on time-domain convolutions, while another synthesized kernel runs 4× faster than a Caffe2 reference.

  • Takeaways & Limitations

    TC provides an end-to-end route from mathematical tensor expressions to optimized GPU kernels and remains competitive with vendor libraries on standard operators.

  • Takeaways & Limitations

    The initial release does not support recurrent definitions, and the implementation lacks register tiling and advanced promotion schemes.

Abstract

from arXiv · show

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]

1 Introduction

Existing computation-graph frameworks abstract tensor operations but can miss hardware- and shape-specific optimizations or require costly custom operators. Tensor Comprehensions addresses this gap with a mathematical tensor language and a polyhedral, autotuned compilation flow for optimized GPU kernels.

  • Motivation: Computation-graph frameworks rely on optimized individual-operator libraries, but unsupported computations require costly custom operators.These abstractions can fall short when computations do not fit existing library calls.
  • Motivation: Efficient execution depends on data size, tensor shape, neighboring computations, memory layout, and hardware features.Graph-level abstractions alone do not capture all refinements and lowering steps needed for accelerator performance.
  • Tensor Comprehensions: Tensor Comprehensions provides a concise, expressive language for tensor computations with shape inference, layout transformations, and specialization.Its syntax generalizes Einstein notation while supporting safety features such as static bound checking.
  • Compilation flow: The compilation flow lowers tensor comprehensions to GPU code using domain-specific polyhedral algorithms, kernel fusion, multilevel parallelism, and memory-hierarchy promotion.The system targets reduced launch and synchronization overhead while optimizing deeply nested tensor computations.
  • Autotuning: An autotuning framework uses JIT compilation and code caching to specialize generated kernels for non-standard sizes.It exposes optimization control from the machine-learning framework through the code generator.
  • Integration and scope: TC integrates with PyTorch and Caffe2 while initially focusing on CUDA code for NVIDIA GPUs.The authors state that the approach may apply to other heterogeneous nodes with shared or partitioned memory.

2 Related Work

Related systems motivate domain-specific optimization, but the paper positions TC as a more generic framework combining tensor-algebra descriptions with polyhedral scheduling and GPU-oriented code generation. Its distinctions include cross-operator optimization, specialized mappings, and independence from a specific computation-graph framework.

  • Active libraries: Active libraries generate specialized code, but isolated kernel tuning misses context-dependent opportunities and cannot feasibly cover all kernel combinations.This motivates optimization at a broader computational context than individual library kernels.
  • TC’s positioning: TC combines a domain-specific language with a generic code-generation framework grounded in loop-nest optimization and parallelization research.The framework targets existing and emerging machine-learning models rather than a single application domain.
  • TC’s positioning: TC automates hierarchical tiling, mapping, shifting, fusion, distribution, and interchange for parametric or fully instantiated tensor problems.The paper states these transformations are not accessible through Halide, Latte, or XLA tensor-operation representations.
  • Polyhedral compilation: TC’s compiler uses heuristics for long, non-uniform reuse patterns and deeply nested loops in deep-learning models.The paper contrasts these domain-specific heuristics with their absence in Halide and related variants.
  • Framework integration: Unlike XLA’s framework integration focus, TC remains independent of a specific computation-graph framework while preserving integration with production frameworks.The related-work discussion also distinguishes TC’s optimization and mapping design from XLA’s.

3 Tensor Comprehensions

Tensor Comprehensions provide a concise, Einstein-style notation for multidimensional tensor computations, with inferred indices and reductions. Their semantics support safe in-place updates, common neural-network kernels, layout transformations, and JIT specialization.

  • Notation and semantics: Tensor Comprehensions express multidimensional tensor computations using concise Einstein-style notation with implicitly defined indices and inferred ranges.Indices appearing only on the right-hand side become reduction dimensions, and iteration order does not affect the output.
  • Notation and semantics: TC preserves functional, full-tensor atomic semantics for in-place updates, so right-hand sides are fully read before left-hand-side assignment.The compiler checks causality and rejects unsafe cases such as in-place transposition with interfering liveness.
  • Kernel expressions: Common machine-learning kernels, including matrix-vector products, SGEMM, fully connected ReLU layers, convolutions, and max pooling, can be written in a few lines.Reduction initialization uses the +=! shorthand, while pointwise operations and tensor reuse express fused computations without temporary storage.
  • Range inference: Tensor comprehensions infer loop ranges from tensor accesses, using affine bounds and explicit where annotations when ranges are under-constrained or ambiguous.The inference procedure seeks maximal rectangular ranges that avoid out-of-bounds input accesses.
  • Scope: The first release does not support recurrent definitions required for recurrent neural networks.This bounds the initial language scope despite its support for convolutional and other tensor computations.
  • Data layout transformations: TC supports generalized layout transformations, including transpositions and implicit data tiling, with range checking to keep reshaped tensor accesses consistent.JIT partial evaluation can specialize affine indexing when stride values are constant.

4 High-Level Workflow Description

TC integrates into machine-learning computation graphs through an in-process API that can replace backend operators and compile tensor expressions. The workflow specializes tensor sizes, lowers through intermediate representations, searches optimization strategies with autotuning, and can emit reference or fallback implementations.

  • Framework integration: TC integrates in-process with computation-graph engines through a thin API that translates framework tensor objects and can override backend operators.A single TC may correspond to a DAG of operators, and users can define their own TC operators.
  • Compilation flow: The compiler starts from specialized tensor sizes and strides, lowers TC to a parametric Halide expression, then to a polyhedral representation.The evolved flow bypasses PENCIL by lowering Halide-IR directly to the polyhedral representation.
  • Autotuning: An autotuner searches scheduling and mapping strategies while serialized compilation and code caching support repeated compilation workflows.The autotuner interacts tightly with scheduling and mapping transformations to explore the optimization space.
  • Execution options: TC can generate a readable identity-mapped CUDA reference implementation for single-thread correctness checks, while CPU LLVM JIT support was not yet implemented.A future fallback path was planned to emit library calls backed by CUDNN.
  • Compilation flow: Early specialization was found beneficial for profitability decisions during polyhedral scheduling, even though the toolchain also supports parametric specifications.The workflow therefore uses specialization during compilation rather than relying exclusively on symbolic parameters.

5 Polyhedral JIT Compilation

The polyhedral JIT compiler lowers Tensor Comprehensions into GPU kernels through schedule-tree transformations, mapping, memory promotion, and safety checks. It targets deep-learning-specific parallelism, locality, layouts, and tensor shapes.

  • Representation and lowering: The compiler lowers tensor operations into a polyhedral representation that supports affine transformations, schedule trees, and target-specific GPU information.Schedule trees communicate execution order, thread-relative induction, synchronization, and data transfers to downstream optimization.
  • Scheduling: Schedule trees encode statement order and iteration structure, enabling transformations such as fusion, tiling, interchange, and distribution while preserving dependences.The canonical tree uses sequence, filter, and band nodes; context nodes record parameter assumptions such as tensor extents.
  • Safety and semantics: Polyhedral access relations infer tensor footprints and check that accessed elements fit declared tensor ranges, detecting out-of-bounds accesses.TC also supports in-place updates with full-tensor functional semantics and rejects programs when its syntactic causality check fails.
  • Scheduling: Loop tiling converts a permutable schedule band into outer tile loops and inner point loops, facilitating GPU mapping and temporal reuse.Figure 3.c illustrates the fused and tiled SGEMM schedule tree.
  • Memory promotion: Memory promotion caches constant-size, potentially strided array tiles in software-controlled local memories and supports explicit shared or private-memory transfers.The compiler also handles indirectly accessed arrays by caching indexed values locally when promotion conditions permit.

6 Autotuning and Caching

The autotuning system searches GPU kernel configurations by compiling and profiling many candidates, then stores the fastest known versions in a compilation cache. Cache keys specialize generated code to Tensor Comprehensions, input shapes, targets, architectures, and optimization choices.

  • Compilation cache: The compilation cache reuses autotuned kernels for similar input shapes, target architectures, and optimization options, avoiding repeated expensive compilation.Cache entries store generated CUDA or PTX code and the fastest known version for each key.
  • Compilation cache: Each cache key is a tuple of the Tensor Comprehension, input shapes, target, and architecture, while optimization choices determine the generated code.The cache can be serialized for persistence and queried before kernel optimization.
  • Search strategy: The autotuner initializes starting configurations, tuning dimensions, admissible values, and either genetic or random search before running for a prescribed time.It updates the cache with better versions as tuning proceeds.
  • Search strategy: Genetic search evaluates candidates by runtime-based fitness, combines genes from three probabilistically selected parents, and applies low-probability mutation.Mutation controls the exploration-versus-exploitation tradeoff.
  • Parallel autotuning: Autotuning evaluates hundreds to thousands of kernel versions, compiling candidates on CPU threads and profiling completed kernels on available GPUs in parallel.Profiling results update the autotuning database and generate subsequent candidates.
  • Tuned options: The search tunes tile, block, and grid sizes, unrolling bounds, fusion and schedule choices, and shared or private-memory usage affecting GPU occupancy.Tile and grid choices include powers of two and integer ceil divisors to reduce tail effects.

7 Examples And Performance Results

TC’s evaluation spans common kernels and production models on Maxwell and Pascal GPUs, comparing baseline and autotuned mappings with established framework and library references. Results show strong gains for several fused, specialized, and latency-bound workloads, while large matrix multiplication remains behind CUBLAS.

  • Kernel results: Up to 30% latency benefits arise when mappings avoid launching an extra block for problem sizes close to powers of two.This optimization targets latency-bound kernels by accepting a slightly off tile or mapping size.
  • Evaluation setup: The evaluation uses 1,000 GPU runs, reports p0, p50, and p90 timings, and compares TC variants with Caffe2 and ATen when available.Experiments use eight M40 or eight P100 GPUs per node; NVRTC and synchronization overhead are included in reported numbers.
  • Limitations: At the largest problem sizes, TC is 4.2× slower than Caffe2 with CUBLAS on Maxwell and 3.4× slower on Pascal.The paper attributes this gap partly to CUBLAS’s extensive hand tuning and TC’s lack of register tiling and advanced promotion schemes.
  • Kernel results: 3.5× speedup over CUBLAS is reached for Factorization Machines on Maxwell, and 3.7× on Pascal, at sizes (B, N, M, K) = (500, 26, 72, 26).The dedicated kernel takes 78µs versus 325µs for the cited CUBLAS kernel on Maxwell.
  • Kernel results: TC outperforms Caffe2’s CUDNN Winograd kernels by 1.4× to 3.6× on Maxwell and 1.9× to 8.8× on Pascal for time-domain convolutions.The comparison is notable despite TC’s current restriction to time-domain convolutions.
  • Production models: 4× speedup over Caffe2 is achieved for the synthesized 2LUT kernel on Maxwell, and 4.1× on Pascal, despite Caffe2 using additional parallel reduction support.A novel two-stage shared-memory loading scheme avoids latency dependencies that otherwise reduce overall performance by more than 5×.
  • Production models: TC achieves up to 1.5× speedup over CUBLAS for a single MLP on Pascal, while fused MLP layers are not beneficially combined by the compared frameworks.The result is obtained despite missing register optimization.
  • Production models: TC reaches up to 3.6× faster than Caffe2 for a low-latency binary classifier by emitting one CUDA kernel instead of 3 to 9 framework-dependent calls.The kernel runs in 32µs, while mean overhead is 25µs, partly attributable to NVRTC.

8 Perspectives

The work identifies opportunities to extend Tensor Comprehensions across architectures, data representations, graph partitioning, transformations, and differentiation. These directions aim to broaden automation, portability, and usefulness for machine-learning workloads.

  • Future work includes distributing best implementations and autotuning histories across architectures via protobuf.
  • The system could support more architectures and combine synthesized kernels with libraries of high-performance primitives.
  • Automated data-layout transformations could support tuning and vector types, including arbitrary-bit low-precision formats.
  • An automated DAG partitioning algorithm could use synthesized-kernel performance to guide partition decisions.
  • Further extensions include Halide-style tiling and model slicing, symbolic automatic differentiation, richer control flow, and sparse, vector, and mixed-precision types.
  • These opportunities are intended to accelerate machine-learning research while preserving performance and easing translation from mathematical specifications to implementations.

9 Conclusion

The conclusion presents Tensor Comprehensions as an end-to-end system that translates mathematical tensor programs into automatically generated GPU kernels. It combines polyhedral compilation, domain-specific optimization, autotuning, caching, and framework integration to address productivity and efficiency gaps.

  • Tensor Comprehensions supports a mathematics-like tensor language and an end-to-end flow to automatically generated GPU kernels.The system remains within polyhedral analysis while supporting tensor operations, affine transformations, code generation, autotuning, and a compilation cache.
  • TC gives domain experts more expressive control over computations and storage-computation tradeoffs while reducing dependence on highly tuned vendor libraries.The paper states that TC synthesizes solid baseline versions that lift bottlenecks in large training runs.

A Appendix

The appendix supplies technical material supporting the paper’s methods and documents the Tensor Comprehension language grammar in EBNF notation.

  • The supplementary material collects technical details intended to provide complete coverage of the proposed methods.
  • Figure 14 presents the Tensor Comprehension language grammar using EBNF notation.

A.1 ML Frameworks of Interest

The appendix reviews tensor-operation abstractions across major machine-learning frameworks and earlier systems. It contrasts execution models, intermediate representations, portability choices, and the degree of compiler or user control over optimization.

  • ML Frameworks of Interest: TensorFlow and most other machine-learning frameworks describe computations over n-dimensional tensors using operators.
  • ML Frameworks of Interest: Caffe2 resembles TensorFlow structurally but also permits opaque intermediate objects such as pre-packed matrices and specialized-hardware handles.
  • ML Frameworks of Interest: PyTorch differs from TensorFlow and Caffe2 by executing operations as specified rather than requiring a separately constructed static computation graph.
  • ML Frameworks of Interest: MXNet supports both declarative and imperative paradigms and includes NNVM, a graph-based intermediate representation for graph-rewriting transformations.
  • ML Frameworks of Interest: Lush introduced tensor-manipulation ideas including index iteration and stride-based index arithmetic, while leaving tiling and reordering to users.

A.2 ML Framework Interface API

The framework interface exposes tensor comprehensions through an execution engine that manages compilation and execution, while framework-specific wrappers support multiple ML systems and a common tensor interchange format.

  • Execution interface: Tensor comprehensions are loaded as strings into an execution engine that invokes JIT compilation and autotuning when code runs.The API version described is specific to Python and PyTorch, with analogous APIs for C++ and other ML frameworks.
  • Framework abstraction: Each framework-specific API is a small wrapper around a core framework-agnostic API.
  • Framework examples: TensorFlow training is organized into graph construction, optimization-node definition, and session-based execution on data.
  • Framework examples: The section illustrates comparable one-layer training code in TensorFlow, Caffe2, and PyTorch.

A.3 Background on Polyhedral Compilation

Polyhedral compilation represents tensor programs through iteration domains, dependences, and schedules, then generates parallel code while preserving semantics and optimizing memory movement.

  • Compilation workflow: A compilation cache reuses autotuned kernel results under similar input shapes, target architectures, and optimization options.
  • Polyhedral representation: An iteration domain is the convex polyhedron formed by loop iterations constrained by affine bounds.Statement instances are points whose coordinates are induction-variable values.
  • Polyhedral representation: Schedules determine the lexicographic execution order of iteration-domain points and can be represented as schedule trees.
  • Code generation: Code generation converts an iteration domain and schedule into loop nests for targets including CUDA, OpenMP, and OpenCL.
  • Dependence preservation: Dependence analysis identifies shared-element accesses involving writes and constrains transformations to preserve the relative order of dependent instances.
  • Dependence preservation: Schedule trees can enforce valid ordering, such as placing initialization before update statements in matrix multiplication.
  • Memory optimization: Tensor comprehensions avoid pointer arithmetic and subscript overflow, preventing memory aliasing and improving dependence-analysis precision.
  • Memory optimization: Memory promotion computes block footprints, groups overlapping references, copies array tiles into shared memory, and inserts synchronization around global-memory transfers.A 32 × 32 tile is obtained in the running matrix-multiplication example.

A.4 Detailed Results on the TC Examples

The TC examples show autotuned CUDA mappings, shared and private memory use, specialization, and fusion choices across matrix multiplication, convolutions, embeddings, and production-model components.

  • Experimental setup: The experiments report baseline mappings and best autotuned mappings for non-production TC examples, using an earlier TC and modified PPCG version.
  • Transposed batched matrix multiplication: For transposed batched matrix multiplication, autotuning selects specialized tiling, thread/block mappings, shared and private memory, and unrolling for (M, K, N) = (128, 32, 256) on Maxwell.
  • Transposed batched matrix multiplication: The generated transposed-multiplication CUDA improves on CUBLAS, although indirect expressions in register arrays remain and the autotuner misses potentially better mappings.
  • Grouped convolutions: Grouped-convolution specifications avoid non-affine indexing and mixed reduction/parallel semantics by using an explicit higher-dimensional formulation.
  • Grouped convolutions: Grouped-convolution mappings specialize thread layouts for small W and use shared memory with autotuned unrolling on Maxwell.
  • Grouped convolutions: The grouped-convolution implementation is not ideal because it omits registers and overprovisions threads, yet delivers much higher performance than a CUDNN implementation.
  • Production models: The production model is decomposed into manually grouped tensor comprehensions because TC lacks a generic operator-DAG abstraction and does not yet support TC calls from TC.
Loading 1802.04730v3…