一句话:提出 DASTAC 框架,将张量高层结构传播到低层代码生成,结合布局压缩与多面体优化实现并行张量代数。
张量代数在机器学习和科学计算等数据密集型负载中至关重要。随着数据复杂度增长,研究者常在高度专用的稠密张量代数与结构感知的高效稀疏张量代数之间面临两难。本文提出 DASTAC,一个将张量捕获的高层结构传播至低层代码生成的框架,融合自动数据布局压缩、多面体分析与仿射代码生成技术。该方法通过自动检测最佳数据布局降低内存占用,充分利用多面体优化,并借助 MLIR 实现进一步优化与并行化。大量实验表明,DASTAC 可与专家手工调优实现相竞争甚至显著超越,单线程获得 0.16x–44.83x 加速,多线程获得 1.37x–243.78x 加速。
Tensor algebra is a crucial component for data-intensive workloads such as machine learning and scientific computing. As the complexity of data grows, scientists often encounter a dilemma between the highly specialized dense tensor algebra and efficient structure-aware algorithms provided by sparse tensor algebra. In this paper, we introduce DASTAC, a framework to propagate the tensors’s captured high-level structure down to low-level code generation by incorporating techniques such as automatic data layout compression, polyhedral analysis, and affine code generation. Our methodology reduces memory footprint by automatically detecting the best data layout, heavily benefits from polyhedral optimizations, leverages further optimizations, and enables parallelization through MLIR. Through extensive experimentation, we show that DASTAC can compete if not significantly outperform specialized hand-tuned implementation by experts. We observe 0.16x 44.83x and 1.37x - 243.78x speed-up for single- and multi-threaded cases, respectively.
@article{paperbot3195,
title = {Compressed and Parallelized Structured Tensor Algebra},
author = {Mahdi Ghorbani and Emilien Bauer and Tobias Grosser and Amir Shaikhha},
journal = {Proceedings of the ACM on Programming Languages},
volume = {9},
number = {OOPSLA1},
year = {2025},
doi = {10.1145/3720506}
}