[论文解读] Stardust: Compiling Sparse Tensor Algebra to a Reconfigurable Dataflow Architecture

Olivia Hsu, Alexander Rucker+3 more|arXiv (Cornell University)|Nov 7, 2022

Parallel Computing and Optimization Techniques被引用 4

一句话总结

Stardust 是首个自动将稀疏张量代数表达式编译为可重构数据流架构（RDAs）的编译器，通过一种新颖的调度与数据表示语言，高效映射稀疏计算。它利用 Capstan RDA 上的并行模式重写与片上内存优化，实现了相较于 CPU 内核 138× 的加速，以及相较于 GPU 内核 41× 的加速。

ABSTRACT

We introduce Stardust, a compiler that compiles sparse tensor algebra to reconfigurable dataflow architectures (RDAs). Stardust introduces new user-provided data representation and scheduling language constructs for mapping to resource-constrained accelerated architectures. Stardust uses the information provided by these constructs to determine on-chip memory placement and to lower to the Capstan RDA through a parallel-patterns rewrite system that targets the Spatial programming model. The Stardust compiler is implemented as a new compilation path inside the TACO open-source system. Using cycle-accurate simulation, we demonstrate that Stardust can generate more Capstan tensor operations than its authors had implemented and that it results in 138$ imes$ better performance than generated CPU kernels and 41$ imes$ better performance than generated GPU kernels.

研究动机与目标

• 为解决在 CPU 和 GPU 无法实现的性能与能效瓶颈下，加速稀疏张量代数工作负载长尾部分的性能差距。
• 通过支持从稀疏张量表达式进行高级编译，使数据科学家能够便捷使用可重构数据流架构（RDAs）。
• 在算法、数据表示与调度之间实现关注点分离，以实现高效映射至稀疏加速器。
• 证明通过编译器优化的 RDA 映射可超越手写优化的 CPU 与 GPU 实现。

提出的方法

• 引入一种新型数据表示语言，用于表达在 RDA 片上物理内存中的加速器特定张量布局。
• 设计一种调度语言，用于将变换后的稀疏迭代空间映射至稀疏硬件组件（如数据合并器与计算单元）。
• 采用一种算法，将抽象数据结构描述绑定至特定片上内存，以最小化数据移动。
• 使用并行模式重写系统，将稀疏张量代数表达式降级为 RDA 的声明式稀疏 Spatial 编程模型。
• 扩展开源 TACO 编译器，增加一条针对 Capstan（一种可重构数据流架构）的新编译路径。
• 利用周期精确仿真，评估其相对于 CPU 和 GPU 内核的性能表现。

实验结果

研究问题

• RQ1 高级编译器能否将任意稀疏张量代数表达式映射至可重构数据流架构，并实现超越 CPU 和 GPU 的性能？
• RQ2 如何设计调度与数据表示抽象，以实现在稀疏 RDA 上的高效内存布局与数据流映射？
• RQ3 算法、数据表示与调度的分离是否能实现超越以往手写代码的新型优化映射？
• RQ4 编译器优化在多大程度上可提升 RDA 上稀疏张量内核的算术强度并减少冗余计算？

主要发现

• Stardust 能够成功编译比原始作者手写实现更广泛的 Capstan 张量操作。
• 在周期精确仿真中，系统相较于生成的 CPU 内核性能高出 138×，相较于生成的 GPU 内核性能高出 41×。
• 所提出的调度与数据表示语言能够合成出优于现有手写优化实现的优化算法。
• 算法、数据表示与调度的分离使得各部分可独立优化，并支持在 Capstan 上实现此前未被针对的新型稀疏张量操作。
• 通过并行模式重写实现的编译器降级管道，成功将高层级张量表达式转化为 RDA 的高效 Spatial 代码。