paperbot · PL 论文追踪

全部论文 待读 收藏

P2 Mosaic: An Interoperable Compiler for Tensor Algebra

Mosaic:面向张量代数的可互操作编译器
PLDI 7(PLDI)2023匹配度 2/10 · 该论文属于稀疏张量代数编译与代码生成,虽与机器学习系统底层编译相关,但未涉及用户核心关注的静态分析、LLM/Agent、浮点程序分析或贝叶斯程序分析。其调度与外部函数绑定机制也不属于以 Agentic 程序或形式化验证为中心的研究方向。引用 16 inbox
Manya Bansal, Olivia Hsu, Kunle Olukotun, Fredrik Kjolstad
crossref: 2026-09-10

AI 速览

一句话:提出可互操作的稀疏张量代数编译器 Mosaic,通过调度 API 将子表达式绑定到外部函数并自动融合生成其余代码。

Mosaic 是一个稀疏张量代数编译器,支持用户通过调度 API 将张量子表达式绑定到其他张量代数库或编译器的外部函数。编译器用函数调用替换被绑定子表达式,其余代码由默认代码生成器生成;因默认融合,同一编译器中可同时利用融合与专用函数调用。评估显示,调用手写 CPU 与专用硬件函数在某些情况下相比融合代码最高加速 206×,而融合代码在另一些情况下相比调用外部函数最高加速 3.57×。Mosaic 还提供搜索系统,自动将表达式映射到已注册外部函数,绑定与搜索均经其验证;新增函数接口简单通用,现已集成 38 个外部函数,平均每个约 20 行代码。

问题
张量代数编译器通常只能二选一:要么生成自身的融合代码,要么调用外部专用库或硬件函数,难以在同一编译流程中兼顾两者。
方法
Mosaic 提供调度 API,允许把张量子表达式绑定到外部函数,编译器以调用替换该子表达式并用默认代码生成器生成其余部分;同时提供自动搜索已注册外部函数的机制,并对显式绑定与自动搜索进行验证。
评估
通过调用手写 CPU 函数与专用硬件函数进行实验:外部函数调用相比融合代码最高加速 206×,融合代码相比外部函数调用最高加速 3.57×;并统计了 38 个已集成外部函数的添加成本(平均约 20 行代码)。
原文摘要(Abstract)

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.

链接与引用

DOI 原文 · PDF(开放获取) ·

BibTeX
@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}
}