一句话：提出基于递归循环嵌套重构的稀疏张量计算自动调度器，可生成渐近更优的调度与代码。
稀疏张量代数在量子计算、物理模拟、化学和机器学习中广泛应用，但通用稀疏张量代数编译器难以为稀疏张量收缩生成渐近最优代码。本文通过 kernel fission 与 fusion 实现递归循环嵌套重构，提出能降低渐近时间复杂度和内存占用的通用变换；并设计了基于偏序集（poset）代价模型的自动调度器，利用时间和辅助内存复杂度剪枝调度搜索空间。该调度器还借助 SMT 求解器在编译期用户约束下，以尽可能少的候选调度近似较优调度的 Pareto 前沿。实验表明，相比 TACO 生成的代码，所选自调度在多个真实张量计算上取得数量级加速。
Automated code generation and performance enhancements for sparse tensor algebra have become essential in many real-world applications, such as quantum computing, physical simulations, computational chemistry, and machine learning. General sparse tensor algebra compilers are not always versatile enough to generate asymptotically optimal code for sparse tensor contractions. This paper shows how to generate asymptotically better schedules for complex sparse tensor expressions using kernel fission and fusion. We present generalized loop restructuring transformations to reduce asymptotic time complexity and memory footprint. Furthermore, we present an auto-scheduler that uses a partially ordered set (poset)-based cost model that uses both time and auxiliary memory complexities to prune the search space of schedules. In addition, we highlight the use of Satisfiability Module Theory (SMT) solvers in sparse auto-schedulers to approximate the Pareto frontier of better schedules to the smallest number of possible schedules, with user-defined constraints available at compile-time. Finally, we show that our auto-scheduler can select better-performing schedules and generate code for them. Our results show that the auto-scheduler provided schedules achieve orders of-magnitude speedup compared to the code generated by the Tensor Algebra Compiler (TACO) for several computations on different real-world tensors.
DOI 原文 · arXiv · PDF（开放获取） · DBLP
@article{abs-2311-09549,
title = {SparseAuto: An Auto-scheduler for Sparse Tensor Computations using Recursive Loop Nest Restructuring},
author = {Adhitha Dias and Logan Anderson and Kirshanthan Sundararajah and Artem Pelenitsyn and Milind Kulkarni},
journal = {Proceedings of the ACM on Programming Languages},
volume = {8},
number = {OOPSLA2},
year = {2024},
doi = {10.1145/3689730}
}