[论文解读] Minimum Cost Loop Nests for Contraction of a Sparse Tensor with a Tensor Network

Raghavendra Kanakagiri, Edgar Solomonik|arXiv (Cornell University)|Jul 11, 2023

Tensor decomposition and applicationsMathematics被引用 3

一句话总结

该论文提出了一种框架 SpTTN-Cyclops，可自动识别并执行稀疏张量收缩（SpTTN）中最具效率的循环嵌套结构，利用动态规划最小化缓冲区大小和缓存未命中等成本指标。该框架在张量分解和补全工作负载中，通过利用数据无关的稀疏模式，其性能优于通用库，并可与专用代码相媲美。

ABSTRACT

Sparse tensor decomposition and completion are common in numerous applications, ranging from machine learning to computational quantum chemistry. Typically, the main bottleneck in optimization of these models are contractions of a single large sparse tensor with a network of several dense matrices or tensors (SpTTN). Prior works on high-performance tensor decomposition and completion have focused on performance and scalability optimizations for specific SpTTN kernels. We present algorithms and a runtime system for identifying and executing the most efficient loop nest for any SpTTN kernel. We consider both enumeration of such loop nests for autotuning and efficient algorithms for finding the lowest cost loop-nest for simpler metrics, such as buffer size or cache miss models. Our runtime system identifies the best choice of loop nest without user guidance, and also provides a distributed-memory parallelization of SpTTN kernels. We evaluate our framework using both real-world and synthetic tensors. Our results demonstrate that our approach outperforms available generalized state-of-the-art libraries and matches the performance of specialized codes.

研究动机与目标

• 解决张量分解和补全中的性能瓶颈，其中稀疏张量网络（SpTTN）的张量收缩是主要计算内核。
• 在无需用户指导的情况下，自动选择 SpTTN 内核的最高效循环嵌套结构，克服可能的循环顺序组合爆炸问题。
• 通过通用化、自调优的框架集成现有高性能张量库，实现高性能、可扩展且可移植的 SpTTN 执行。
• 通过利用数据无关的稀疏性并高效使用 BLAS 内核，实现与手工优化的专用代码相当或更优的性能。

提出的方法

• 将循环嵌套表示为树结构，其中每个节点代表一个循环，其子节点为嵌套循环，从而实现结构化的成本建模。
• 应用动态规划高效搜索循环顺序的指数级空间，将复杂度从 O((m!)^N) 降低至 O(N^3 2^m m)。
• 基于循环嵌套树结构，为关键性能瓶颈（如缓冲区大小和缓存未命中）开发成本模型。
• 与 CTF 运行时系统集成，实现 SpTTN 内核的分布式内存并行化。
• 自动选择能最小化成本函数的循环顺序，同时支持高级别使用 BLAS-1 和 BLAS-2 内核以提升性能。
• 扩展 Cyclops 张量框架，支持单个稀疏张量与张量网络中多个稠密张量的同时收缩。

实验结果

研究问题

• RQ1 对于给定的 SpTTN 内核，什么样的循环嵌套结构能最小化内存占用和缓存未命中？
• RQ2 能否使用动态规划高效探索 SpTTN 内核的海量循环顺序排列空间，而无需穷举？
• RQ3 自动生成的循环嵌套结构在张量分解和补全过程中的性能，与手工优化的专用实现相比如何？
• RQ4 中间张量维度边界对循环嵌套效率和 BLAS 内核利用率有何影响？
• RQ5 通用框架是否能在保持跨多种 SpTTN 内核可移植性的同时，实现与专用代码相当的性能？

主要发现

• 在小尺寸、高度稀疏的张量（N=40，稀疏度 0.1%）上，SpTTN-Cyclops 相较 TACO 实现了 534 倍的加速，表明在低维场景下性能显著提升。
• 该框架在一系列真实世界和合成的 SpTTN 内核上，性能优于通用库如 TACO 和 SparseLNR。
• 尽管内存使用更高，但中间张量被限制为大小 T 和 S×T 的循环嵌套结构，由于更高效地利用了 BLAS-1 和 BLAS-2 内核，性能优于仅使用标量中间结果的结构。
• SpTTN-Cyclops 选择的循环顺序在随机采样的 25% 循环顺序中表现最佳，表明其自调优能力出色。
• 该框架在 MTTKRP、TTMc 和 TTTP 内核上的性能与专用代码相当，验证了其通用性与高效性。
• 使用动态规划将搜索空间复杂度从阶乘级降低至多项式级，使 SpTTN 循环嵌套的实用自调优成为可能。