A Novel Compiler Transformation for Fast Sparse Matrix Multiplication in GPUs

2025/06/18 by Hossein Albakri, Albakri, Hossein, Kazem Cheshmi +1
Computer Science · #Distributed and Parallel Computing Systems #FOS: Computer and information sciences #Parallel Computing and Optimization Techniques #Programming Languages (cs.PL) #Quantum Computing Algorithms and Architecture

paper · pdf · doi:10.48550/arxiv.2506.15174

openalex publication_date 2025/06/18 · openalex created_date 2025/10/19 · openalex updated_date 2026/07/28

Abstract

Sparse data structures are commonly used in neural networks to reduce the memory footprint. These data structures are compact but cause irregularities such as random memory accesses, which prevent efficient use of the memory hierarchy. GPUs are a common platform for machine learning practitioners, but running compact data structures on these devices often leads to slow-downs due to inefficient use of computing and memory resources. This paper proposes a new compiler transformation, enumerate-and-sparse-coarsen, that accelerates sparse matrix-matrix multiplication (SPMM) on GPU devices. The transformation increases data reuse in registers and caches while creating more balanced workloads for GPU computing resources. The transformation is tested on sparse neural networks in convolutional and transformer models. On an A100 GPU and across a columns of matrix B (bCols) in A × B = C from range of 32 to 128, the transformation yields a geometric mean speedup of 1.84× to 2.27× compared to cuBLAS and cuSPARSE baselines, respectively.

Citations

• SPLAT: A framework for optimised GPU code-generation for SParse reguLar ATtention
• Improving Locality in Sparse and Dense Matrix Multiplications
• JITSPMM: Just-in-Time Instruction Generation for Accelerated Sparse Matrix-Matrix Multiplication
• SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning
• Sparsity in Deep Learning: Pruning and growth for efficient inference\n and training in neural networks
• Sparse GPU Kernels for Deep Learning
• The State of Sparsity in Deep Neural Networks
• Design Principles for Sparse Matrix Multiplication on the GPU
• Learning Sparse Neural Networks through L₀ Regularization
• To prune, or not to prune: exploring the efficacy of pruning for model compression
• Variational Dropout and the Local Reparameterization Trick

Related

• Sparse GPU Kernels for Deep Learning
• Design Principles for Sparse Matrix Multiplication on the GPU
• GE-SpMM: General-purpose Sparse Matrix-Matrix Multiplication on GPUs for Graph Neural Networks
• PopSparse: Accelerated block sparse matrix multiplication on IPU
• Efficient Sparse-Dense Matrix-Matrix Multiplication on GPUs Using the Customized Sparse Storage Format