[论文解读] Sparse Tensor Algebra Optimizations with Workspaces

Fredrik Kjølstad, Willow Ahrens+2 more|arXiv (Cornell University)|Feb 28, 2018

Tensor decomposition and applications参考文献 24被引用 3

一句话总结

本文提出了一种新的稀疏张量代数中间语言——具体索引记法（concrete index notation），以及一种工作区优化技术，该技术通过使用密集临时张量（工作区）来加速稀疏计算。通过提升循环不变操作并消除昂贵的稀疏插入和合并循环，该方法在性能上达到与手写优化库（如MKL和SPLATT）相当的水平，从而实现了此前不支持的稀疏结果，例如MTTKRP等核函数。

ABSTRACT

This paper shows how to optimize sparse tensor algebraic expressions by introducing temporary tensors, called workspaces, into the resulting loop nests. We develop a new intermediate language for tensor operations called concrete index notation that extends tensor index notation. Concrete index notation expresses when and where sub-computations occur and what tensor they are stored into. We then describe the workspace optimization in this language, and how to compile it to sparse code by building on prior work in the literature. We demonstrate the importance of the optimization on several important sparse tensor kernels, including sparse matrix-matrix multiplication (SpMM), sparse tensor addition (SpAdd), and the matricized tensor times Khatri-Rao product (MTTKRP) used to factorize tensors. Our results show improvements over prior work on tensor algebra compilation and brings the performance of these kernels on par with state-of-the-art hand-optimized implementations. For example, SpMM was not supported by prior tensor algebra compilers, the performance of MTTKRP on the nell-2 data set improves by 35%, and MTTKRP can for the first time have sparse results.

研究动机与目标

• 通过启用利用临时密集张量（工作区）的优化，解决稀疏张量代数编译器中的性能差距。
• 将稀疏张量代码生成泛化，以支持具有稀疏结果的核函数（如MTTKRP和SpMM），这些在以往难以实现或效率低下。
• 通过消除计算过程中的条件合并循环和昂贵的稀疏插入操作，提升性能。
• 提供一种形式化且可组合的中间表示，使在生成稀疏代码前能够进行激进的优化。

提出的方法

• 提出具体索引记法，作为张量索引记法的扩展，用于指定循环顺序、子计算位置及工作区使用方式。
• 应用工作区优化技术，将张量表达式重写为在密集临时张量（工作区）中预计算子表达式，以降低运行时开销。
• 使用调度构造来请求工作区插入，支持循环不变代码移动，并避免重复的稀疏插入操作。
• 将具体索引记法降低为先前工作（Kjolstad et al., 2017）中的迭代图，从而支持高效稀疏代码的生成。
• 将该优化应用于关键核函数：SpMM、SpAdd 和 MTTKRP，展示了其通用性与性能提升。
• 利用工作区将多路合并循环替换为随机访问模式，减少条件分支并改善数据局部性。

实验结果

研究问题

• RQ1 如何通过使用密集临时张量（工作区）来优化稀疏张量代数，以提升性能？
• RQ2 该工作区优化能否在高层中间语言中表达并编译，以精确控制子计算发生的时间与位置？
• RQ3 该优化是否能够使此前不支持的稀疏结果（如MTTKRP和SpMM等张量核函数）成为可能？
• RQ4 优化后代码的性能与手写优化的前沿库（如MKL和SPLATT）相比如何？
• RQ5 该优化能否推广至更广泛的稀疏张量核函数类别，而不仅限于所研究的案例？

主要发现

• 与先前的张量代数编译相比，该工作区优化在nell-2数据集上将MTTKRP的性能提升了35%。
• SpMM此前无法被先前的张量代数编译器支持，现已被高效编译，性能与手写优化实现相当。
• 该方法首次实现了MTTKRP的稀疏结果，扩展了受支持张量操作的范围。
• 该优化减少了昂贵的稀疏插入和条件合并循环，代之以对密集工作区的高效随机访问。
• 生成的代码性能可与前沿库（如MKL、Eigen和SPLATT）相媲美，证明了该方法的有效性。