A sparse iteration space transformation framework for sparse tensor algebra
Name
3428226.pdf
Description
Published version
Size
1.49 MB
Format
Adobe PDF
Checksum (MD5)
ce1d3946fb43f7540e4c4c753702e862
Author(s)
Senanayake, Ryan
Hong, Changwan
Wang, Ziheng
Wilson, Amalee
Chou, Stephen
Kamil, Shoaib
Amarasinghe, Saman
Kjolstad, Fredrik
Date Issued
2020
Journal
Proceedings of the ACM on Programming Languages
Publisher
Association for Computing Machinery (ACM)
Version
Final published version
Abstract
© 2020 Owner/Author. 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. We implement these concepts by extending the sparse iteration theory implementation in the TACO system. The associated scheduling API can be used by performance engineers or it can be the target of an automatic scheduling system. We outline one heuristic autoscheduling system, but other systems are possible. Using the scheduling API, we show how to optimize mixed sparse-dense tensor algebra expressions on CPUs and GPUs. Our results show that the sparse transformations are sufficient to generate code with competitive performance to hand-optimized implementations from the literature, while generalizing to all of the tensor algebra.
MIT Department
Massachusetts Institute of Technology. Computer Science and Artificial Intelligence Laboratory
Massachusetts Institute of Technology. Department of Electrical Engineering and Computer Science
Terms of Use
Creative Commons Attribution 4.0 International license
Persistent DSpace Link
DOI of Published Version
https://doi.org/10.1145/3428226
