一句话：提出将仿射下标表达式逐步重写为利用稀疏性的循环的编译方法，为稀疏卷积等生成渐近高效代码。
现有稀疏张量代数编译器主要处理简单下标表达式，难以为含仿射下标表达式的张量表达式生成渐近高效代码，因而无法很好支持稀疏卷积和池化等机器学习与图形学计算。本文提出一种编译方法，将复合下标表达式逐步重写为简单下标表达式，并生成利用输入稀疏张量稀疏模式的循环；由此生成内核的时间复杂度由存储元素数量而非张量形状界定。该方法可无缝集成到现有框架，并兼容 MLIR 稀疏编译器对不同稀疏格式组合的处理。实验表明，99.9% 稀疏度下较最先进的编译器方法获得 19.5 倍加速，约 80% 稀疏度时生成的稀疏内核开始优于稠密卷积实现。
This paper extends prior work on sparse tensor algebra compilers to generate asymptotically efficient code for tensor expressions with affine subscript expressions. Our technique enables compiler support for a wide range of sparse computations, including sparse convolutions and pooling that are widely used in ML and graphics applications. We propose an approach that gradually rewrites compound subscript expressions to simple subscript expressions with loops that exploit the sparsity pattern of the input sparse tensors. As a result, the time complexity of the generated kernels is bounded by the number of stored elements and not by the shape of the tensors. Our approach seamlessly integrates into existing frameworks and is compatible with recent advances in compilers for sparse computations, including the flexibility to efficiently handle arbitrary combinations of different sparse tensor formats. The implementation of our algorithm is open source and upstreamed to the MLIR sparse compiler. Experimental results show that our method achieves 19.5x speedup when compared with the state-of-the-art compiler-based method at 99.9% sparsity. The generated sparse kernels start to outperform dense convolution implementations at about 80% sparsity
@article{LiuRXLKB24,
title = {Compiler Support for Sparse Tensor Convolutions},
author = {Peiming Liu and Alexander J Root and Anlun Xu and Yinying Li and Fredrik Kjolstad and Aart J.C. Bik},
journal = {Proceedings of the ACM on Programming Languages},
volume = {8},
number = {OOPSLA2},
year = {2024},
doi = {10.1145/3689721}
}