SpTFS: Sparse Tensor Format Selection for MTTKRP via Deep Learning

Authors: Qingxiao Sun, Yi Liu, Ming Dun, Hailong Yang, and Zhongzhi Luan (Beihang University); Lin Gan and Guangwen Yang (Tsinghua University, China); and Depei Qian (Beihang University)

Abstract: Canonical polyadic decomposition (CPD) is one of the most common tensor computations adopted in many scientific applications. The major performance bottleneck of CPD is matricized tensor times Katri-Rao product (MTTKRP). To optimize the performance of MTTKRP, various sparse tensor formats have been proposed, such as COO and CSF. Due to the complex sparsity patterns of the tensors, however, no single format fits all tensors for optimal performance. To address this problem, we propose SpTFS, a framework that automatically predicts the optimal storage format for an input sparse tensor. Specifically, we propose tensor lowering and matrix representation techniques to transform the high-dimensional tensors into fix-sized matrices. In addition, we develop a customized convolutional neural network by incorporating additional feature layer to compensate the sparsity features lost during tensor transformation. The experiment results show that SpTFS achieves the prediction accuracy of 92.7% and 96% on average on CPU and GPU, respectively.




Back to Technical Papers Archive Listing