Sparse tensor algebra is challenging to efficiently parallelize due to the irregular, data-dependent, and potentially skewed structure of sparse computation. We propose the first partitioning algorithm that provably load balances the computation of any sparse tensor algebra expression across parallel execution units. Our algorithm generalizes parallel merging algorithms to any number of operands, and to multi-dimensional, hierarchical sparse data structures. We implement our algorithm within an existing sparse tensor algebra compilation framework to automatically generate parallel sparse tensor algebra kernels that target multi-core CPUs and GPUs. We show that our generated code is competitive with hand-implemented parallelization strategies used by vendor libraries like Intel MKL and NVIDIA cuSPARSE (geo-means of $0.73$--$3.4\times$) and \textsc{Taco} (geo-means of $1.0$--$2.4\times$), and significantly outperforms general-purpose strategies for sparse tensor expressions where specialized algorithms have not been developed (geo-means of $2.0$--$6.4\times$).

翻译：稀疏张量代数由于计算具有不规则性、数据依赖性和潜在的偏斜结构，难以高效地并行化。我们提出了首个分区算法，能够可证明地实现任意稀疏张量代数表达式在并行执行单元上的负载均衡计算。该算法将并行归并算法推广至任意数量的操作数，并适用于多维层次化稀疏数据结构。我们在现有稀疏张量代数编译框架中实现了该算法，以自动生成面向多核CPU和GPU的并行稀疏张量代数内核。实验表明，我们生成的代码与Intel MKL和NVIDIA cuSPARSE等供应商库的手工实现并行化策略（几何均值$0.73$--$3.4\times$）及\textsc{Taco}（几何均值$1.0$--$2.4\times$）具有竞争力，并在尚未开发专用算法的稀疏张量表达式上显著优于通用策略（几何均值$2.0$--$6.4\times$）。