Minimum Cost Loop Nests for Contraction of a Sparse Tensor with a Tensor Network

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.

15 pages, 7 figures

References in corpus (10)

• The ITensor Software Library for Tensor Network Calculations
• Simulating quantum computation by contracting tensor networks
• Advances on Tensor Network Theory: Symmetries, Fermions, Entanglement, and Holography
• A Systematic Survey of General Sparse Matrix-Matrix Multiplication
• Faster identification of optimal contraction sequences for tensor networks
• Scalable Tucker Factorization for Sparse Tensors - Algorithms and Discoveries
• Software for Sparse Tensor Decomposition on Emerging Computing Architectures
• DISTAL: The Distributed Tensor Algebra Compiler
• Scalable Task-Based Algorithm for Multiplication of Block-Rank-Sparse Matrices
• SparseLNR: Accelerating Sparse Tensor Computations Using Loop Nest Restructuring

Cited by in corpus (1)

• Distributed-Memory Randomized Algorithms for Sparse Tensor CP Decomposition