一句话:提出可互操作的稀疏张量代数编译器 Mosaic,通过调度 API 将子表达式绑定到外部函数并自动融合生成其余代码。
Mosaic 是一个稀疏张量代数编译器,支持用户通过调度 API 将张量子表达式绑定到其他张量代数库或编译器的外部函数。编译器用函数调用替换被绑定子表达式,其余代码由默认代码生成器生成;因默认融合,同一编译器中可同时利用融合与专用函数调用。评估显示,调用手写 CPU 与专用硬件函数在某些情况下相比融合代码最高加速 206×,而融合代码在另一些情况下相比调用外部函数最高加速 3.57×。Mosaic 还提供搜索系统,自动将表达式映射到已注册外部函数,绑定与搜索均经其验证;新增函数接口简单通用,现已集成 38 个外部函数,平均每个约 20 行代码。
We introduce Mosaic, a sparse tensor algebra compiler that can bind tensor expressions to external functions of other tensor algebra libraries and compilers. Users can extend Mosaic by adding new functions and bind a sub-expression to a function using a scheduling API. Mosaic substitutes the bound sub-expressions with calls to the external functions and automatically generates the remaining code using a default code generator. As the generated code is fused by default, users can productively leverage both fusion and calls to specialized functions within the same compiler. We demonstrate the benefits of our dual approach by showing that calling hand-written CPU and specialized hardware functions can provide speedups of up to 206× against fused code in some cases, while generating fused code can provide speedups of up to 3.57× against code that calls external functions in other cases. 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. Additionally, the interface for adding new external functions is simple and general. Currently, 38 external functions have been added to Mosaic, with each addition averaging 20 lines of code.
@article{paperbot1940,
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},
number = {PLDI},
year = {2023},
doi = {10.1145/3591236}
}