SENSEi: Input-Sensitive Compilation for Accelerating GNNs

TL;DR

SENSEi is a system that optimizes GNN performance by selecting the best matrix re-association based on input attributes, achieving significant speedups on various hardware platforms.

Contribution

It introduces an input-sensitive compilation approach for GNNs that dynamically chooses optimal matrix re-associations for improved performance.

Findings

Achieves up to 2.012x speedup on GCNs and 16.274x on GATs.

Effectively generalizes to GNN variants including sampling-based methods.

Provides hardware-agnostic techniques that enhance diverse GNN systems.

Abstract

Over the years, many frameworks and optimization techniques have been proposed to accelerate graph neural networks (GNNs). Compared to the optimizations explored in these systems, we observe that different matrix re-associations of GNN computations lead to novel input-sensitive performance behavior. We leverage this observation to propose SENSEi, a system that exposes different sparse and dense matrix primitive compositions based on different matrix re-associations of GNN computations and selects the best among them based on input attributes. SENSEi executes in two stages: (1) an offline compilation stage that enumerates all valid re-associations leading to different sparse-dense matrix compositions and uses input-oblivious pruning techniques to prune away clearly unprofitable candidates and (2) an online runtime system that explores the remaining candidates and uses light-weight cost models to select the best re-association based on the input graph and the embedding sizes on a given hardware platform. On a wide range of configurations, SENSEi achieves speedups of up to $2.012\times$ and $1.85\times$ on graph convolutional networks and up to $6.294\times$ and $16.274\times$ on graph attention networks, on GPUs and CPUs respectively. We also show that its technique generalizes to GNN variants, including those that require sampling. Furthermore, we show that SENSEi's techniques are agnostic to the underlying GNN system, and can be used to yield synergistic improvements across a diverse set of implementations.