Polyhedral Specification and Code Generation of Sparse Tensor Contraction with Co-iteration
Polyhedral Specification and Code Generation of Sparse Tensor Contraction with Co-iteration
复制标题
DOI：
10.1145/3566054
复制
发表时间：
2022-08
影响因子：
1.6
通讯作者：
Tuowen Zhao;Tobi Popoola;Mary W. Hall;C. Olschanowsky;M. Strout
中科院分区：
文献类型：
--
作者：
Tuowen Zhao;Tobi Popoola;Mary W. Hall;C. Olschanowsky;M. Strout
本文提出了一个用于稀疏张量压缩计算的代码生成器。它利用了稀疏多面体框架(SPF)中循环嵌套计算的数学表示，该框架扩展了多面体模型以支持非仿射计算，例如在稀疏张量中出现的计算。将SPF扩展用于稀疏张量代码的布局描述、优化和代码生成：(1)我们提出了一种多面体布局规范，将布局和计算的迭代空间解耦；(2)我们通过SMT求解器将对一个稀疏张量布局的多面体扫描与代码合成相结合来发现其他张量中的相应元素，从而实现了稀疏张量的高效联合迭代。我们将生成的代码与最先进的张量编译器Taco生成的代码进行了比较。我们在稀疏-稀疏联合迭代上获得了比TACO平均快1.63倍的并行性能，并描述了如何通过切换FIND算法将其平均加速比提高到2.72倍。我们还证明了将布局和计算的迭代空间解耦使得能够支持额外的布局和计算组合。
This article presents a code generator for sparse tensor contraction computations. It leverages a mathematical representation of loop nest computations in the sparse polyhedral framework (SPF), which extends the polyhedral model to support non-affine computations, such as those that arise in sparse tensors. SPF is extended to perform layout specification, optimization, and code generation of sparse tensor code: (1) We develop a polyhedral layout specification that decouples iteration spaces for layout and computation; and (2) we develop efficient co-iteration of sparse tensors by combining polyhedra scanning over the layout of one sparse tensor with the synthesis of code to find corresponding elements in other tensors through an SMT solver. We compare the generated code with that produced by a state-of-the-art tensor compiler, TACO. We achieve on average 1.63× faster parallel performance than TACO on sparse-sparse co-iteration and describe how to improve that to 2.72× average speedup by switching the find algorithms. We also demonstrate that decoupling iteration spaces of layout and computation enables additional layout and computation combinations to be supported.