High Performance Unstructured SpMM Computation Using Tensor Cores

TL;DR

This paper introduces SMaT, a GPU-accelerated library that leverages Tensor Cores for high-performance unstructured sparse matrix-matrix multiplication, significantly outperforming existing libraries.

Contribution

The paper presents SMaT, a novel SpMM library that effectively utilizes Tensor Cores for unstructured sparsity, with algorithmic optimizations and CUDA MMA API integration.

Findings

SMaT outperforms state-of-the-art libraries by up to 125x.

Achieves an average speedup of 2.6x over existing solutions.

Demonstrates applicability to scientific computing and large-model workloads.

Abstract

High-performance sparse matrix-matrix (SpMM) multiplication is paramount for science and industry, as the ever-increasing sizes of data prohibit using dense data structures. Yet, existing hardware, such as Tensor Cores (TC), is ill-suited for SpMM, as it imposes strict constraints on data structures that cannot be met by unstructured sparsity found in many applications. To address this, we introduce (S)parse (Ma)trix Matrix (T)ensor Core-accelerated (SMaT): a novel SpMM library that utilizes TCs for unstructured sparse matrices. Our block-sparse library leverages the low-level CUDA MMA (matrix-matrix-accumulate) API, maximizing the performance offered by modern GPUs. Algorithmic optimizations such as sparse matrix permutation further improve performance by minimizing the number of non-zero blocks. The evaluation on NVIDIA A100 shows that SMaT outperforms SotA libraries (DASP, cuSPARSE, and Magicube) by up to 125x (on average 2.6x). SMaT can be used to accelerate many workloads in scientific computing, large-model training, inference, and others.