Cached Operator Reordering: A Unified View for Fast GNN Training

TL;DR

This paper introduces a unified framework for optimizing GNN training by reordering operators and caching, significantly improving speed and efficiency across different GNN layers and hardware platforms.

Contribution

It proposes adaptive operator reordering with caching strategies that enhance GNN training speed and memory efficiency, a novel approach compared to prior methods.

Findings

Up to 2.43x speedup for GCN training

Up to 1.94x speedup for GAT training

Memory savings and hardware-agnostic implementation

Abstract

Graph Neural Networks (GNNs) are a powerful tool for handling structured graph data and addressing tasks such as node classification, graph classification, and clustering. However, the sparse nature of GNN computation poses new challenges for performance optimization compared to traditional deep neural networks. We address these challenges by providing a unified view of GNN computation, I/O, and memory. By analyzing the computational graphs of the Graph Convolutional Network (GCN) and Graph Attention (GAT) layers -- two widely used GNN layers -- we propose alternative computation strategies. We present adaptive operator reordering with caching, which achieves a speedup of up to 2.43x for GCN compared to the current state-of-the-art. Furthermore, an exploration of different caching schemes for GAT yields a speedup of up to 1.94x. The proposed optimizations save memory, are easily implemented across various hardware platforms, and have the potential to alleviate performance bottlenecks in training large-scale GNN models.