TACO: Distributed and Heterogeneous Sparse Compiler

Tensor Algebra Compiler (TACO) is a compiler that compiles sparse tensor operations. It been shown to be competitive with or surpass manually optimized high-performance sparse linear and tensor algebra libraries like Eigen, Intel MKL, and SPLATT, while generalizing to any sparse tensor algebra expression on many different types of data structures. The current version of TACO has limited single-GPU support and no multi-node support. Many sparse problems are larger than the memory size of a GPU and therefore require multi-GPU support. Further, other applications and libraries are already multi-node (PETSc, Trilinos, etc.) and require fast distributed sparse solvers and tensor operations. We plan to express distribution across GPUs by extending the TACO format language to describe tensor locales. The code generator will automatically generate distributed code when a tensor algebra expression is executed on a sparse tensor located on a GPU. That is, the code will conform to the shape of the data, which is distributed. The system will then generate code to compute operations on the sparse blocks once they are co-located in the same locale. We will develop an extension to tensor formats that describes the locale where a tensor is located, such as on a GPU, or in a GPU’s shared memory. We will then extend the TACO scheduling language, so that the user or an auto-scheduling system can control how the computation should be orchestrated. One example of an aspect to be controlled by the scheduling language is where the computation should occur, such as an ‘owner computes’ pattern. Given such scheduling commands, the system will automatically orchestrate the communication. In Phase I, we will extend TACO to support distributed memory execution focusing on single-node, multi-GPU targets. After showing the technical feasibility of distributed memory execution in Phase I, we will extend this support to multi-node, multi-GPU targets in Phase II. We will harden these features and make them commercially usable by non-experts. TACO fulfills an enormous unmet market demand: an easy to use system that provides the full power of operating on sparse tensors. Writing the computation part of the algorithm is straightforward for any mathematically inclined individual. But deriving performant schedules requires a significant knowledge of the algorithm, the iteration spaces, and the hardware specifications. Many sparse problems are larger than the memory size of a GPU and therefore require multi-GPU support. Further, other applications and libraries are already multi-node (PETSc, Trilinos, etc.) and require fast distributed sparse solvers and tensor operations. A multi-GPU, multi-node tensor algebra compiler will enable commercial development of codes that can operate on large sparse data sets that are not possible today.