一句话：用稀疏卷积网络与近似最近邻搜索协同优化稀疏张量程序的存储格式与调度。
稀疏张量程序的性能同时取决于数据格式与计算调度，而两者的选择又与具体稀疏模式强耦合，传统方法难以在巨大搜索空间中高效找到最优组合。本文提出 WACO，核心是设计一个轻量级代价模型：它利用稀疏卷积网络从稀疏模式中学习有意义的特征，并通过专门设计的调度模板刻画格式与调度之间的耦合关系，从而较准确地预测运行时。在协同优化的庞大搜索空间上，WACO 采用近似最近邻搜索的策略，为给定稀疏模式快速且准确地检索最佳格式与调度。作者在 CPU 上针对 SpMV、SpMM、SDDMM 与 MTTKRP 四种算法、726 种稀疏模式进行了评测，结果显示 WACO 优于 Intel MKL、BestFormat、默认调度的 TACO 和 ASpT 四个强基线；与其中最好的基线相比，四个算法分别取得 1.43×、1.18×、1.14× 和 1.27× 的平均加速。
In this paper, we present WACO, a novel method of co-optimizing the format and the schedule of a given sparsity pattern in a sparse tensor program. A core challenge in this paper is the design of a lightweight cost model that accurately predicts the runtime of a sparse tensor program by considering the sparsity pattern, the format, and the schedule. The key idea in addressing this is exploiting a sparse convolutional network to learn meaningful features of the sparsity pattern and embedding a coupled behavior between the format and the schedule using a specially designed schedule template. In addition, within the enormous search space of co-optimization, our novel search strategy, an approximate nearest neighbor search, efficiently and accurately retrieves the best format and schedule for a given sparsity pattern. We evaluated WACO for four different algorithms (SpMV, SpMM, SDDMM, and MTTKRP) on a CPU using 726 different sparsity patterns. Our experimental results showed that WACO outperformed four state-of-the-art baselines, Intel MKL, BestFormat, TACO with a default schedule, and ASpT. Compared to the best of four baselines, WACO achieved 1.43×, 1.18×, 1.14×, and 1.27× average speedups on SpMV, SpMM, SDDMM, and MTTKRP, respectively.
@inproceedings{paperbot10015,
title = {WACO: Learning Workload-Aware Co-optimization of the Format and Schedule of a Sparse Tensor Program},
author = {Jaeyeon Won and Charith Mendis and Joel S. Emer and Saman Amarasinghe},
booktitle = {ASPLOS},
year = {2023},
doi = {10.1145/3575693.3575742}
}