TL;DR
This paper introduces a format-agnostic approach to sparse tensor algebra compilation, enabling flexible support for multiple data layouts and improving performance through modular code generation.
Contribution
It presents a modular interface for tensor formats, allowing easy addition of new formats and integration into a tensor algebra compiler, demonstrated in the Taco system.
Findings
Code generated is competitive with hand-optimized implementations.
Supporting multiple formats can significantly improve performance, e.g., COO format speeds up matrix-vector multiplication.
The approach simplifies extending tensor algebra compilers to new data formats.
Abstract
This paper shows how to build a sparse tensor algebra compiler that is agnostic to tensor formats (data layouts). We develop an interface that describes formats in terms of their capabilities and properties, and show how to build a modular code generator where new formats can be added as plugins. We then describe six implementations of the interface that compose to form the dense, CSR/CSF, COO, DIA, ELL, and HASH tensor formats and countless variants thereof. With these implementations at hand, our code generator can generate code to compute any tensor algebra expression on any combination of the aforementioned formats. To demonstrate our technique, we have implemented it in the taco tensor algebra compiler. Our modular code generator design makes it simple to add support for new tensor formats, and the performance of the generated code is competitive with hand-optimized…
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Code & Models
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
