We address the problem of optimizing mixed sparse and dense tensor algebra in a compiler. We show that standard loop transformations, such as strip-mining, tiling, collapsing, parallelization and vectorization, can be applied to irregular loops over sparse iteration spaces. We also show how these transformations can be applied to the contiguous value arrays of sparse tensor data structures, which we call their position space, to unlock load-balanced tiling and parallelism. We have prototyped these concepts in the open-source TACO system, where they are exposed as a scheduling API similar to the Halide domain-specific language for dense computations. Using this scheduling API, we show how to optimize mixed sparse/dense tensor algebra expressions, how to generate load-balanced code by scheduling sparse tensor algebra in position space, and how to generate sparse tensor algebra GPU code. Our evaluation shows that our transformations let us generate good code that is competitive with many hand-optimized implementations from the literature.

中文翻译：

---

稀疏和密集张量代数的统一迭代空间转换框架

我们解决了在编译器中优化混合稀疏和密集张量代数的问题。我们展示了标准循环变换，例如条带挖掘、平铺、折叠、并行化和矢量化，可以应用于稀疏迭代空间上的不规则循环。我们还展示了如何将这些转换应用于稀疏张量数据结构的连续值数组，我们称之为位置空间，以解锁负载平衡的平铺和并行性。我们在开源 TACO 系统中对这些概念进行了原型设计，它们作为调度 API 公开，类似于用于密集计算的 Halide 领域特定语言。使用这个调度 API，我们展示了如何优化混合稀疏/密集张量代数表达式，如何通过在位置空间中调度稀疏张量代数来生成负载平衡的代码，以及如何生成稀疏张量代数 GPU 代码。我们的评估表明，我们的转换让我们生成了与文献中的许多手动优化实现相竞争的优秀代码。