Mosaic: An Interoperable Compiler for Tensor Algebra

Speaker: Manya Bansal, Undergraduate Student, Stanford University
Date: May 10, 2023

Sparse tensor algebra is an important computational language in many domains including graph processing, machine learning and scientific computing. This has resulted in a proliferation of high-performance libraries for CPUs, GPUs and domain-specific hardware. At the same time, there has been a rise in compilers that produce fused code by default, breaking library boundaries. In this talk, I will present Mosaic, a system that enables users to productively leverage both fusion and calls to specialized functions within the same compiler. Mosaic can compose externally defined library functions to implement an arbitrary sparse tensor algebra expression, unlocking good performance where hand-optimized implementations or specialized hardware exist. It fills in the gaps that are not provided by the libraries, guaranteeing generality in both expressions and data structures, as well as fusion. I will discuss Mosaic's design and show that there exist computations that benefit from both fused and factorized code optimizations and that these computations can be scheduled through Mosaic.