POSTER: Optimizing Sparse Tensor Contraction with Revisiting Hash Table Design

Abstract

Sparse tensor contraction (SpTC) serves as an essential operation in high-performance applications. The high dimensionality of sparse tensors makes SpTC fundamentally challenging in aspects such as costly multidimensional index search, extensive intermediate output data, and indirect addressing. Previous state-of-the-art work addresses some of these challenges through hash-table implementation. In this paper, we propose a hash-table based and fully optimized SpTC by providing a more carefully designed customized hash table design, proposing an architecture-aware algorithm for hash table selection with size prediction, applying cross-stage optimizations to exploit shared information and avoid redundant operations. Evaluating on a set of tensors extracted from the real world, our method can achieve superior speedup and reduce the memory footprint substantially compared to the current state-of-the-art work.