A sparse iteration space transformation framework for sparse tensor algebra
Abstract
We address the problem of optimizing sparse tensor algebra in a compiler and show how to define standard loop transformations---split, collapse, and reorder---on sparse iteration spaces. The key idea is to track the transformation functions that map the original iteration space to derived iteration spaces. These functions are needed by the code generator to emit code that maps coordinates between iteration spaces at runtime, since the coordinates in the sparse data structures remain in the original iteration space. We further demonstrate that derived iteration spaces can tile both the universe of coordinates and the subset of nonzero coordinates: the former is analogous to tiling dense iteration spaces, while the latter tiles sparse iteration spaces into statically load-balanced blocks of nonzeros. Tiling the space of nonzeros lets the generated code efficiently exploit heterogeneous compute resources such as threads, vector units, and GPUs.
Document Details
- Document Type
- Pub Defense Publication
- Publication Date
- Nov 13, 2020
- Source ID
- 10.1145/3428226
Entities
People
- Amalee Wilson
- Changwan Hong
- Fredrik Kjølstad
- Ryan Senanayake
- Saman Amarasinghe
- Shoaib Kamil
- Stephen Chou
- Ziheng Wang
Organizations
- Adobe
- Defense Advanced Research Projects Agency
- Massachusetts Institute of Technology
- National Science Foundation
- Stanford University
- Toyota Physical and Chemical Research Institute
- United States Department of Energy