Format Abstraction for Sparse Tensor Algebra Compilers

Abstract

This paper shows how to build a sparse tensor algebra compiler that is agnostic to tensor formats (data layouts). We develop an interface that describes formats in terms of their capabilities and properties, and show how to build a modular code generator where new formats can be added as plugins. We then describe six implementations of the interface that compose to form the dense, CSR/CSF, COO, DIA, ELL, and HASH tensor formats and countless variants thereof. With these implementations at hand, our code generator can generate code to compute any tensor algebra expression on any combination of the aforementioned formats. To demonstrate our technique, we have implemented it in the taco tensor algebra compiler. Our modular code generator design makes it simple to add support for new tensor formats, and the performance of the generated code is competitive with hand-optimized implementations. Furthermore, by extending taco to support a wider range of formats specialized for different application and data characteristics, we can improve end-user application performance. For example, if input data is provided in the COO format, our technique allows computing a single matrix-vector multiplication directly with the data in COO, which is up to 3.6$\times$ faster than by first converting the data to CSR.

Presented at OOPSLA 2018

References in corpus (4)

• TensorFlow: A system for large-scale machine learning
• Format Abstraction for Sparse Tensor Algebra Compilers
• High-Performance Tensor Contraction without Transposition
• A Unified Optimization Approach for Sparse Tensor Operations on GPUs

Cited by in corpus (20)

• Format Abstraction for Sparse Tensor Algebra Compilers
• Hardware Acceleration of Sparse and Irregular Tensor Computations of ML Models: A Survey and Insights
• Compiler Support for Sparse Tensor Computations in MLIR
• The Sparse Abstract Machine
• Capstan: A Vector RDA for Sparsity
• ALTO: Adaptive Linearized Storage of Sparse Tensors
• Automatic Generation of Efficient Sparse Tensor Format Conversion Routines
• TeAAL: A Declarative Framework for Modeling Sparse Tensor Accelerators
• SparseLNR: Accelerating Sparse Tensor Computations Using Loop Nest Restructuring
• Mechanised Hypersafety Proofs about Structured Data: Extended Version
• Distributed-Memory DMRG via Sparse and Dense Parallel Tensor Contractions
• UniSparse: An Intermediate Language for General Sparse Format Customization
• The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding
• A Unified Iteration Space Transformation Framework for Sparse and Dense Tensor Algebra
• An Attempt to Generate Code for Symmetric Tensor Computations
• The Continuous Tensor Abstraction: Where Indices are Real
• Dynamic Sparse Tensor Algebra Compilation
• GPTPU: Accelerating Applications using Edge Tensor Processing Units
• Extending Sparse Tensor Accelerators to Support Multiple Compression Formats
• A Parallel Sparse Tensor Benchmark Suite on CPUs and GPUs