Recent years have seen considerable work on compiling sparse tensor algebra expressions. This paper addresses a shortcoming in that work, namely how to generate efficient code (in time and space) that scatters values into a sparse result tensor. We address this shortcoming through a compiler design that generates code that uses sparse intermediate tensors (sparse workspaces) as efficient adapters between compute code that scatters and result tensors that do not support random insertion. Our compiler automatically detects sparse scattering behavior in tensor expressions and inserts necessary intermediate workspace tensors. We present an algorithm template for workspace insertion that is the backbone of our code generation algorithm. Our algorithm template is modular by design, supporting sparse workspaces that span multiple user-defined implementations. Our evaluation shows that sparse workspaces can be up to 27.12$\times$ faster than the dense workspaces of prior work. On the other hand, dense workspaces can be up to 7.58$\times$ faster than the sparse workspaces generated by our compiler in other situations, which motivates our compiler design that supports both. Our compiler produces sequential code that is competitive with hand-optimized linear and tensor algebra libraries on the expressions they support, but that generalizes to any other expression. Sparse workspaces are also more memory efficient than dense workspaces as they compress away zeros. This compression can asymptotically decrease memory usage, enabling tensor computations on data that would otherwise run out of memory.
翻译:近年来,稀疏张量代数表达式的编译研究取得了显著进展。本文针对该领域的一个不足,即如何生成(在时间和空间上)高效的代码,将值分散到稀疏结果张量中。我们通过一种编译器设计来解决这一不足,该设计生成的代码使用稀疏中间张量(稀疏工作空间)作为高效适配器,连接执行散射操作的代码与不支持随机插入的结果张量。我们的编译器自动检测张量表达式中的稀疏散射行为,并插入必要的中间工作空间张量。我们提出了一种工作空间插入算法模板,作为代码生成算法的核心框架。该算法模板采用模块化设计,支持跨越多个用户自定义实现的稀疏工作空间。评估表明,稀疏工作空间的速度最高可达先前工作中稠密工作空间的27.12倍。另一方面,在某些情况下,稠密工作空间的速度比我们的编译器生成的稀疏工作空间快7.58倍,这促使我们采用同时支持两种工作空间的编译器设计。我们的编译器生成的顺序代码在所支持的表达式上可与手工优化的线性代数与张量代数库相媲美,并且可泛化到任何其他表达式。此外,稀疏工作空间通过压缩零元素比稠密工作空间具有更高的内存效率。这种压缩能渐进式降低内存使用量,使得原本会因内存不足而失败的张量计算得以执行。