# Accelerating SpMM Kernel with Cache-First Edge Sampling for Graph Neural Networks

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 31 references that are in the map.
- First posted on arXiv 21 Apr 2021.

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

- [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)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Relational inductive biases, deep learning, and graph networks](https://mlsanity.com/paper/1806.01261/) (2018)
- [Deeper Insights into Graph Convolutional Networks for Semi-Supervised Learning](https://mlsanity.com/paper/1801.07606/) (2018)
- [Predict then Propagate: Graph Neural Networks meet Personalized PageRank](https://mlsanity.com/paper/1810.05997/) (2018)
- [Representation Learning on Graphs with Jumping Knowledge Networks](https://mlsanity.com/paper/1806.03536/) (2018)
- [Simple and Deep Graph Convolutional Networks](https://mlsanity.com/paper/2007.02133/) (2020)
- [DropEdge: Towards Deep Graph Convolutional Networks on Node Classification](https://mlsanity.com/paper/1907.10903/) (2019)
- [GraphSAINT: Graph Sampling Based Inductive Learning Method](https://mlsanity.com/paper/1907.04931/) (2019)
- [FastGCN: Fast Learning with Graph Convolutional Networks via Importance Sampling](https://mlsanity.com/paper/1801.10247/) (2018)
- [DeepGCNs: Can GCNs Go as Deep as CNNs?](https://mlsanity.com/paper/1904.03751/) (2019)
- [Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks](https://mlsanity.com/paper/1909.01315/) (2019)
- [Modeling polypharmacy side effects with graph convolutional networks](https://mlsanity.com/paper/1802.00543/) (2018)
- [Learning to Represent Programs with Graphs](https://mlsanity.com/paper/1711.00740/) (2017)
- [Adaptive Sampling Towards Fast Graph Representation Learning](https://mlsanity.com/paper/1809.05343/) (2018)
- [Sparse GPU Kernels for Deep Learning](https://mlsanity.com/paper/2006.10901/) (2020)
- [Bayesian Graph Neural Networks with Adaptive Connection Sampling](https://mlsanity.com/paper/2006.04064/) (2020)
- [GE-SpMM: General-purpose Sparse Matrix-Matrix Multiplication on GPUs for Graph Neural Networks](https://mlsanity.com/paper/2007.03179/) (2020)
- [FeatGraph: A Flexible and Efficient Backend for Graph Neural Network Systems](https://mlsanity.com/paper/2008.11359/) (2020)
- [Large Graph Convolutional Network Training with GPU-Oriented Data Communication Architecture](https://mlsanity.com/paper/2103.03330/) (2021)
- [Design Principles for Sparse Matrix Multiplication on the GPU](https://mlsanity.com/paper/1803.08601/) (2018)

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

- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [Characterizing and Taming Resolution in Convolutional Neural Networks](https://mlsanity.com/paper/2110.14819/) (2021)

## Abstract (arXiv)

Graph neural networks (GNNs), an emerging deep learning model class, can extract meaningful representations from highly expressive graph-structured data and are therefore gaining popularity for wider ranges of applications. However, current GNNs suffer from the poor performance of their sparse-dense matrix multiplication (SpMM) operator, even when using powerful GPUs. Our analysis shows that 95% of the inference time could be spent on SpMM when running popular GNN models on NVIDIA's advanced V100 GPU. Such SpMM performance bottleneck hinders GNNs' applicability to large-scale problems or the development of more sophisticated GNN models. To address this inference time bottleneck, we introduce ES-SpMM, a cache-first edge sampling mechanism and codesigned SpMM kernel. ES-SpMM uses edge sampling to downsize the graph to fit into GPU's shared memory. It thus reduces the computation cost and improves SpMM's cache locality. To evaluate ES-SpMM's performance, we integrated it with a popular GNN framework, DGL, and tested it using representative GNN models and datasets. Our results show that ES-SpMM outperforms the highly optimized cuSPARSE SpMM kernel by up to 4.35x with no accuracy loss and by 45.3x with less than a 1% accuracy loss.