GNNIE: GNN inference engine with load-balancing and graph-specific caching

Graph neural networks (GNN) analysis engines are vital for real-world problems that use large graph models. Challenges for a GNN hardware platform include the ability to (a) host a variety of GNNs, (b) handle high sparsity in input vertex feature vectors and the graph adjacency matrix and the accompanying random memory access patterns, and (c) maintain load-balanced computation in the face of uneven workloads, induced by high sparsity and power-law vertex degree distributions. This paper proposes GNNIE, an accelerator designed to run a broad range of GNNs. It tackles workload imbalance by (i) splitting vertex feature operands into blocks, (ii) reordering and redistributing computations, (iii) using a novel flexible MAC architecture. It adopts a graph-specific, degree-aware caching policy that is well suited to real-world graph characteristics. The policy enhances on-chip data reuse and avoids random memory access to DRAM.

GNNIE achieves average speedups of 21233× over a CPU and 699× over a GPU over multiple datasets on graph attention networks (GATs), graph convolutional networks (GCNs), Graph-SAGE, GINConv, and DiffPool. Compared to prior approaches, GNNIE achieves an average speedup of 35× over HyGCN (which cannot implement GATs) for GCN, GraphSAGE, and GINConv, and, using 3.4× fewer processing units, an average speedup of 2.1× over AWB-GCN (which runs only GCNs).