Optimizing Tensor Programs on Flexible Storage

TL;DR

This paper introduces a system that enables flexible, declarative storage formats for tensors and a cost-based optimizer to enhance performance, demonstrated by significant empirical improvements over existing systems.

Contribution

It presents a novel system allowing declarative specification of tensor storage formats and an optimizer to improve tensor program performance.

Findings

Significant performance improvements over state-of-the-art systems

Flexible storage formats tailored to tensor sparsity properties

Effective cost-based optimization for tensor memory layouts

Abstract

Tensor programs often need to process large tensors (vectors, matrices, or higher order tensors) that require a specialized storage format for their memory layout. Several such layouts have been proposed in the literature, such as the Coordinate Format, the Compressed Sparse Row format, and many others, that were especially designed to optimally store tensors with specific sparsity properties. However, existing tensor processing systems require specialized extensions in order to take advantage of every new storage format. In this paper we describe a system that allows users to define flexible storage formats in a declarative tensor query language, similar to the language used by the tensor program. The programmer only needs to write storage mappings, which describe, in a declarative way, how the tensors are laid out in main memory. Then, we describe a cost-based optimizer that optimizes the tensor program for the specific memory layout. We demonstrate empirically significant performance improvements compared to state-of-the-art tensor processing systems.