[论文解读] Polyhedral Specification and Code Generation of Sparse Tensor Contraction with Co-Iteration

Tuowen Zhao, Tobi Popoola+3 more|arXiv (Cornell University)|Aug 25, 2022

Parallel Computing and Optimization Techniques被引用 6

一句话总结

本论文提出了一种用于稀疏张量收缩的多面体框架扩展，通过解耦布局空间与计算空间，实现了灵活的布局指定以及通过SMT合成的查找算法实现的高效联合迭代。与TACO相比，该方法在平均性能上实现了1.63倍的加速，使用优化后的查找算法时最高可达2.72倍加速，支持此前无法实现的布局形式，如BCSR和对角线计算。

ABSTRACT

This paper 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 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$ imes$ faster parallel performance than TACO on sparse-sparse co-iteration and describe how to improve that to 2.72$ imes$ 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.

研究动机与目标

• 为解决现有稀疏张量编译器在支持复杂布局（如BCSR和对角线计算）方面的局限性，这些局限性源于计算空间与布局空间的紧密耦合。
• 通过将逻辑坐标空间与物理布局迭代空间解耦，实现对多个稀疏张量的一般性联合迭代。
• 将稀疏多面体框架（SPF）扩展为支持形式化布局描述和联合迭代，利用未解释函数与SMT求解器。
• 生成针对架构优化的代码，支持布局变换与联合迭代模式，包括析取合并（未来工作）。
• 通过支持数据布局与计算的联合优化，提升性能可移植性，适用于包括SIMD和张量核心在内的多种硬件。

提出的方法

• 通过将布局建模为物理位置与逻辑坐标之间的数学关系，解耦迭代空间，从而实现布局与计算的独立优化。
• 使用未解释函数表示多面体模型中依赖运行时的索引数组，支持稀疏计算中的非仿射访问。
• 利用SMT求解器合成查找算法，将一个张量布局中的索引映射到另一个张量中对应的非零元素，从而实现联合迭代。
• 将对一个张量布局的多面体扫描与合成代码相结合，以定位其他张量中匹配的非零元素，避免显式使用合并网格。
• 通过避免布局层级与计算索引之间的一一对应关系，支持复杂布局（如分块格式，例如BCSR）和对角线计算。
• 通过组合布局映射与计算变换生成优化代码，支持与分块、重排和向量化等优化的正交组合。

实验结果

研究问题

• RQ1 多面体框架能否被扩展以支持稀疏张量布局的形式化描述，同时将其与计算迭代空间解耦？
• RQ2 如何在形式化、数学基础坚实的框架下高效表达并编译多个稀疏张量的联合迭代？
• RQ3 SMT合成的查找算法是否能在联合迭代性能上超越手写或基于模板的方法？
• RQ4 在多大程度上可以实现布局与计算的联合优化，以支持此前不被支持的格式（如BCSR或对角线操作）？
• RQ5 生成代码的性能与最先进的张量编译器（如TACO）相比，在多样化的稀疏张量工作负载下表现如何？

主要发现

• 所提出的框架在基准测试套件上的并行稀疏-稀疏联合迭代中，与TACO相比实现了1.63倍的平均加速。
• 通过切换到优化后的查找算法，框架在与TACO的对比中最高实现了2.72倍的平均加速，证明了算法合成的显著性能影响。
• 布局与计算空间的解耦使得框架能够支持此前无法表达的布局形式，如分块压缩稀疏行（BCSR），而TACO无法表达此类布局。
• 该框架支持对角线计算及其他需要在多个层级间共享循环索引的模式，这些在TACO的层级格式模型中是不可行的。
• 使用SMT合成的查找算法可在无需合并网格或预计算索引映射的情况下实现高效的联合迭代。
• 该方法支持布局变换与其他优化（如分块和向量化）的正交组合，从而在现代架构上实现更好的性能。