Abstract

We introduce Mosaic, a sparse tensor algebra compiler that binds tensor (sub-)expressions to external functions of other tensor algebra libraries and compilers. Users can extend Mosaic by adding new functions and can bind a sub-computation to a function using a scheduling API. Mosaic substitutes the bound (sub-)expressions with the specified function call and automatically fills in the rest of the unbound code using a default code generator. Mosaic also offers a search system that can automatically map an expression to a set of registered external functions. Both the explicit binding and automatic search are verified by Mosaic. We demonstrate the benefits of our approach by showing that calling hand-written CPU and specialized hardware functions can provide speedup of up to 206x (for an SDDMM expression) and 173x (for an SpMV expression), respectively, over a homogeneous compiler. Mosaic's external function interface is simple and general. Currently, 38 external functions have been added to Mosaic, with each addition averaging 20 lines of code.

Article

pdf

BibTeX

@article{bansal2023,
  title={Mosaic: An Interoperable Compiler for Tensor Algebra},
  author={Manya Bansal and Olivia Hsu and Kunle Olukotun and Fredrik Kjolstad},
  journal={Proceedings of the ACM on Programming Languages},
  volume={7},
  issue={PLDI},
  year={2023},
  month={June}
}