Sparse tensor & array compilation

TACO: The Tensor Algebra Compiler

A compiler that generates fast kernels for tensor algebra expressions over dense and sparse tensors from a high-level index notation.

C++ C CUDA

The problem

Hand-writing sparse tensor kernels is tedious and error-prone, and every new expression, format, or target needs its own implementation.

The idea

Separate what is computed (an index-notation expression) from how tensors are stored (a format), then generate a specialized kernel for any combination of the two.

How it fits the group's work

TACO is the foundation of much of the group’s work on representation-polymorphic computation. Its central abstraction — describing computation as index notation and storage as a composable format — is what lets a single expression be compiled to code for many different data structures. Later projects extend the same idea to new domains and machines: distributing it across clusters (DISTAL), mapping it to dataflow accelerators (the Sparse Abstract Machine), binding subexpressions to external libraries (Mosaic), and adding shape operators on sparse arrays (Burrito).

Key publications

All software