Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

iSpLib: A Library for Accelerating Graph Neural Networks using Auto-tuned Sparse Operations

Md Saidul Hoque Anik, Pranav Badhe, Rohit Gampa, Ariful Azad

TL;DR

The paper tackles the challenge of optimizing sparse matrix operations in GNNs, which heavily depend on input sparsity and hardware. It presents iSpLib, a PyTorch-based C++ library with auto-tuned SpMM/SDDMM kernels, cache-enabled backpropagation, and a Python plug-in for seamless PyG integration. Key contributions include general-purpose auto-tuning for CPU platforms, efficient backprop caching, semiring support, and PyG patching utilities, yielding up to $27\times$ speedups on CPU over PyTorch 2.1 and PyG 2.4 in experiments. This work enables broad, portable acceleration of GNN training and inference on multi-core CPUs with minimal user changes, highlighting the potential of auto-tuned sparse kernels in practical GNN workloads.

Abstract

Core computations in Graph Neural Network (GNN) training and inference are often mapped to sparse matrix operations such as sparse-dense matrix multiplication (SpMM). These sparse operations are harder to optimize by manual tuning because their performance depends significantly on the sparsity of input graphs, GNN models, and computing platforms. To address this challenge, we present iSpLib, a PyTorch-based C++ library equipped with auto-tuned sparse operations. iSpLib expedites GNN training with a cache-enabled backpropagation that stores intermediate matrices in local caches. The library offers a user-friendly Python plug-in that allows users to take advantage of our optimized PyTorch operations out-of-the-box for any existing linear algebra-based PyTorch implementation of popular GNNs (Graph Convolution Network, GraphSAGE, Graph Inference Network, etc.) with only two lines of additional code. We demonstrate that iSpLib obtains up to 27x overall training speedup compared to the equivalent PyTorch 2.1.0 and PyTorch Geometric 2.4.0 implementations on the CPU. Our library is publicly available at https://github.com/HipGraph/iSpLib (https://doi.org/10.5281/zenodo.10806511).

iSpLib: A Library for Accelerating Graph Neural Networks using Auto-tuned Sparse Operations

TL;DR

The paper tackles the challenge of optimizing sparse matrix operations in GNNs, which heavily depend on input sparsity and hardware. It presents iSpLib, a PyTorch-based C++ library with auto-tuned SpMM/SDDMM kernels, cache-enabled backpropagation, and a Python plug-in for seamless PyG integration. Key contributions include general-purpose auto-tuning for CPU platforms, efficient backprop caching, semiring support, and PyG patching utilities, yielding up to speedups on CPU over PyTorch 2.1 and PyG 2.4 in experiments. This work enables broad, portable acceleration of GNN training and inference on multi-core CPUs with minimal user changes, highlighting the potential of auto-tuned sparse kernels in practical GNN workloads.

Abstract

Core computations in Graph Neural Network (GNN) training and inference are often mapped to sparse matrix operations such as sparse-dense matrix multiplication (SpMM). These sparse operations are harder to optimize by manual tuning because their performance depends significantly on the sparsity of input graphs, GNN models, and computing platforms. To address this challenge, we present iSpLib, a PyTorch-based C++ library equipped with auto-tuned sparse operations. iSpLib expedites GNN training with a cache-enabled backpropagation that stores intermediate matrices in local caches. The library offers a user-friendly Python plug-in that allows users to take advantage of our optimized PyTorch operations out-of-the-box for any existing linear algebra-based PyTorch implementation of popular GNNs (Graph Convolution Network, GraphSAGE, Graph Inference Network, etc.) with only two lines of additional code. We demonstrate that iSpLib obtains up to 27x overall training speedup compared to the equivalent PyTorch 2.1.0 and PyTorch Geometric 2.4.0 implementations on the CPU. Our library is publicly available at https://github.com/HipGraph/iSpLib (https://doi.org/10.5281/zenodo.10806511).
Paper Structure (13 sections, 3 figures, 1 table)

This paper contains 13 sections, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related Works
  3. Library Design
  4. Overview
  5. Auto-tuning Mechanism
  6. Efficient Backpropagation
  7. Semiring Support
  8. Matmul Interface and Dependency
  9. PyTorch Geometric Integration
  10. Experimental Setting
  11. Results
  12. Discussion
  13. Conclusion

Figures (3)

  • Figure 1: Overview of the iSpLib Library
  • Figure 2: Tuning Graph for various CPUs
  • Figure 3: Average per-epoch training time and speedup for iSpLib w.r.t. other frameworks [PT2: PyTorch 2.1, PT1: PyTorch < 2, PT2-Compile: PyTorch 2 torch.compile, PT2-MP: PyTorch 2 Message Passing paradigm]