The Sparse Abstract Machine

TL;DR

The paper introduces the Sparse Abstract Machine (SAM), an abstract model for sparse tensor algebra targeting specialized hardware, with a compiler and simulator to demonstrate its flexibility and optimization potential.

Contribution

It presents SAM as a new abstract machine for sparse tensor algebra, along with Custard compiler and dataflow simulator, enabling hardware-aware optimizations and flexible tensor computations.

Findings

SAM effectively separates tensor formats from algorithms.

SAM's dataflow graphs support arbitrary iteration orders.

Experimental results show SAM's extensibility and hardware representation capabilities.

Abstract

We propose the Sparse Abstract Machine (SAM), an abstract machine model for targeting sparse tensor algebra to reconfigurable and fixed-function spatial dataflow accelerators. SAM defines a streaming dataflow abstraction with sparse primitives that encompass a large space of scheduled tensor algebra expressions. SAM dataflow graphs naturally separate tensor formats from algorithms and are expressive enough to incorporate arbitrary iteration orderings and many hardware-specific optimizations. We also present Custard, a compiler from a high-level language to SAM that demonstrates SAM's usefulness as an intermediate representation. We automatically bind from SAM to a streaming dataflow simulator. We evaluate the generality and extensibility of SAM, explore the performance space of sparse tensor algebra optimizations using SAM, and show SAM's ability to represent dataflow hardware.