A sparse iteration space transformation framework for sparse tensor algebra
A sparse iteration space transformation framework for sparse tensor algebra
复制标题
DOI：
10.1145/3428226
复制
发表时间：
2020-11
影响因子：
--
通讯作者：
Ryan Senanayake;Changwan Hong;Ziheng Wang;Amalee Wilson;Stephen Chou;Shoaib Kamil;Saman P. Amarasinghe;Fredrik Kjolstad
中科院分区：
文献类型：
--
作者：
Ryan Senanayake;Changwan Hong;Ziheng Wang;Amalee Wilson;Stephen Chou;Shoaib Kamil;Saman P. Amarasinghe;Fredrik Kjolstad
我们解决了在编译器中优化稀疏张量代数的问题，并展示了如何在稀疏迭代空间上定义标准循环变换-分裂，折叠和重新排序。其关键思想是跟踪将原始迭代空间映射到派生迭代空间的转换函数。代码生成器需要这些函数来发出在运行时在迭代空间之间映射坐标的代码，因为稀疏数据结构中的坐标保持在原始迭代空间中。我们进一步证明，派生的迭代空间可以瓷砖宇宙的坐标和非零坐标的子集：前者是类似于瓷砖密集的迭代空间，而后者瓷砖稀疏的迭代空间到静态负载平衡块的非零。对非零空间进行分块，可以让生成的代码有效地利用线程、向量单元和GPU等异构计算资源。我们通过扩展稀疏迭代理论在TACO系统中的实现来实现这些概念。相关的调度API可供性能工程师使用，也可作为自动调度系统的目标。我们概述了一个启发式自动调度系统，但其他系统是可能的。使用调度API，我们展示了如何在CPU和GPU上优化混合稀疏-密集张量代数表达式。我们的研究结果表明，稀疏变换足以生成具有竞争力的性能，从文献中手工优化实现的代码，同时推广到所有的张量代数。
We address the problem of optimizing sparse tensor algebra in a compiler and show how to define standard loop transformations---split, collapse, and reorder---on sparse iteration spaces. The key idea is to track the transformation functions that map the original iteration space to derived iteration spaces. These functions are needed by the code generator to emit code that maps coordinates between iteration spaces at runtime, since the coordinates in the sparse data structures remain in the original iteration space. We further demonstrate that derived iteration spaces can tile both the universe of coordinates and the subset of nonzero coordinates: the former is analogous to tiling dense iteration spaces, while the latter tiles sparse iteration spaces into statically load-balanced blocks of nonzeros. Tiling the space of nonzeros lets the generated code efficiently exploit heterogeneous compute resources such as threads, vector units, and GPUs. We implement these concepts by extending the sparse iteration theory implementation in the TACO system. The associated scheduling API can be used by performance engineers or it can be the target of an automatic scheduling system. We outline one heuristic autoscheduling system, but other systems are possible. Using the scheduling API, we show how to optimize mixed sparse-dense tensor algebra expressions on CPUs and GPUs. Our results show that the sparse transformations are sufficient to generate code with competitive performance to hand-optimized implementations from the literature, while generalizing to all of the tensor algebra.