# Optimizing Sparse Matrix Multiplications for Graph Neural Networks

arXiv: https://arxiv.org/abs/2111.00352 · HTML: https://mlsanity.com/paper/2111.00352/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 46 references that are in the map.
- First posted on arXiv 30 Oct 2021.

## Builds on (most cited first, 16 of 16)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Modeling Relational Data with Graph Convolutional Networks](https://mlsanity.com/paper/1703.06103/) (2017)
- [Graph Neural Networks: A Review of Methods and Applications](https://mlsanity.com/paper/1812.08434/) (2018)
- [Deep Gaussian Embedding of Graphs: Unsupervised Inductive Learning via Ranking](https://mlsanity.com/paper/1707.03815/) (2017)
- [Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks](https://mlsanity.com/paper/1909.01315/) (2019)
- [A Survey on Network Embedding](https://mlsanity.com/paper/1711.08752/) (2017)
- [GNN-FiLM: Graph Neural Networks with Feature-wise Linear Modulation](https://mlsanity.com/paper/1906.12192/) (2019)
- [Cross-lingual Knowledge Graph Alignment via Graph Matching Neural Network](https://mlsanity.com/paper/1905.11605/) (2019)
- [Machine Learning in Compiler Optimisation](https://mlsanity.com/paper/1805.03441/) (2018)
- [Optimizing Streaming Parallelism on Heterogeneous Many-Core Architectures: A Machine Learning Based Approach](https://mlsanity.com/paper/2003.04294/) (2020)
- [Characterizing Scalability of Sparse Matrix-Vector Multiplications on Phytium FT-2000+ Many-cores](https://mlsanity.com/paper/1911.08779/) (2019)

## Built on by (most cited first, 6 of 6)

- [Parallel and Distributed Graph Neural Networks: An In-Depth Concurrency Analysis](https://mlsanity.com/paper/2205.09702/) (2022)
- [Dynamic GPU Energy Optimization for Machine Learning Training Workloads](https://mlsanity.com/paper/2201.01684/) (2022)
- [Two Heads Are Better Than One: Boosting Graph Sparse Training via Semantic and Topological Awareness](https://mlsanity.com/paper/2402.01242/) (2024)
- [PEARL: Preconditioner Enhancement through Actor-critic Reinforcement Learning](https://mlsanity.com/paper/2501.10750/) (2025)
- [SENSEi: Input-Sensitive Compilation for Accelerating GNNs](https://mlsanity.com/paper/2306.15155/) (2023)
- [iSpLib: A Library for Accelerating Graph Neural Networks using Auto-tuned Sparse Operations](https://mlsanity.com/paper/2403.14853/) (2024)

## Abstract (arXiv)

Graph neural networks (GNNs) are emerging as a powerful technique for modeling graph structures. Due to the sparsity of real-world graph data, GNN performance is limited by extensive sparse matrix multiplication (SpMM) operations involved in computation. While the right sparse matrix storage format varies across input data, existing deep learning frameworks employ a single, static storage format, leaving much room for improvement. This paper investigates how the choice of sparse matrix storage formats affect the GNN performance. We observe that choosing a suitable sparse matrix storage format can significantly improve the GNN training performance, but the right format depends on the input workloads and can change as the GNN iterates over the input graph. We then develop a predictive model to dynamically choose a sparse matrix storage format to be used by a GNN layer based on the input matrices. Our model is first trained offline using training matrix samples, and the trained model can be applied to any input matrix and GNN kernels with SpMM computation. We implement our approach on top of PyTorch and apply it to 5 representative GNN models running on a multi-core CPU using real-life and synthetic datasets. Experimental results show that our approach gives an average speedup of 1.17x (up to 3x) for GNN running time.