FeatGraph: A Flexible and Efficient Backend for Graph Neural Network Systems

In recent years, there is a rise of interest in adopting deep learning to structural data such as graphs. Unlike the dense objects (e.g., images, videos, texts) handled by traditional deep learning models, graphs represent sparsely, irregularly connected links. Essentially, graphs are defined on a non-Euclidean domain equipped with vastly different distance measurements and geometric properties, imposing the demand for new neural network architectures.

GNNs are an emerging family of neural networks capable of learning a joint representation for each vertex/edge using both features and topological data. Recent studies [3, 25] have unified different GNN models into a message passing paradigm where each vertex computes a new representation by aggregating features (messages) from its neighbors. More formally, given a graph $\mathcal{G(\mathcal{V},\mathcal{E})}$, we denote the input feature tensor associated with vertex $v$ as $\mathbf{x}_{v}$, and that associated with the edge pointing from vertex $u$ to $v$ as $\mathbf{x}_{uv}$. To get the representation of a vertex and an edge, the message passing paradigm carries out the following computations:

Here $\phi$, $\bigoplus$, and $\psi$ are customizable or parameterized functions (e.g., neural network modules) for calculating messages, aggregating messages, and updating edge representations, respectively. Similar to convolutional neural networks (CNNs), a GNN model iteratively applies Equations (1) (2) to generate vertex and edge representations for higher layers.

There is a strong connection between Equations (1) (2) and sparse matrix operations. For example, given the vertex feature matrix $\mathbf{X_{V}}\in\mathbb{R}^{|\mathcal{V}|\times d}$ and the adjacency matrix $\mathbf{A}$, the vertex-wise computation in the graph convolutional network (GCN) [26], which copies source vertex features as messages and aggregates messages by taking the sum, is equivalent to SpMM (sparse-dense matrix multiplication) as follows.

For edge-wise computations, many GNN models [27, 28] calculate an attention weight on each edge. One popular formulation for calculating attention weight is by a dot product between the source and destination vertex features [29], that is, $\psi(\mathbf{x}_{u},\mathbf{x}_{v},\mathbf{x}_{uv})\triangleq\mathbf{x}_{u}\mathbf{x}_{v}^{T}$. Its tensorized implementation corresponds to SDDMM (sampled dense-dense matrix multiplication) [30], which multiplies two dense matrices, followed by an element-wise multiplication with a sparse mask matrix, to output a sparse matrix.

Hence, Equations (1) and (2) when implemented as tensor operations are generalized SpMM and SDDMM, respectively. They represent two distinct computation patterns in GNN workloads: reduction on each vertex and reduction on each edge. Moreover, according to the chain rule, the gradient computation of SpMM with respect to $\mathbf{A}$ requires a dot product between the gradients of source and destination vertex features, thus following the SDDMM pattern. Likewise, the gradient computation of SDDMM follows the SpMM pattern. Therefore, these two computation patterns are essential for both inference and training of GNNs. In particular, our benchmarking shows that generalized SpMM and SDDMM occupy $\sim 95\%$ of the total run time in training a 2-layer GNN model, using the existing solutions with sub-optimized sparse kernels.

Existing graph processing systems [15, 16] express computations on graphs with a vertex- and/or edge-centric programming paradigm, and they employ a scheduler to realize efficient graph traversal. For example, to ensure load balance on GPU, Gunrock [16] assigns the edges of a vertex to be processed by a thread, a warp, or a block, according to the number of the edges. Edge is the unit for scheduling—the computation on an edge is blackbox to the scheduler. The underlying assumption is that the computation on an edge is lightweight.

However, that assumption breaks in GNNs, which attach a multi-dimensional feature tensor to each vertex/edge, and consequently have more complex vertex-wise and edge-wise computations than traditional graph workloads. For example, MLP aggregation, as shown in Figure 1, performs a sequence of vector matrix multiplication and non-linear activation on each edge. Treating the computation on each edge as a blackbox, Gunrock fails to exploit the abundant parallelism in it.

To enable performant processing of GNN workloads, we need a system that: 1) makes vertex-wise and edge-wise computations whitebox to the scheduler; 2) co-optimizes graph traversal and feature dimension computation. FeatGraph achieves the goals by adopting a tensor compiler approach.

Computation-intensive workloads typically operate on tensors, i.e., multi-dimensional data. For example, traditional deep learning models perform matrix multiplication and convolution over dense tensors; GNNs deal with both dense and sparse tensors (the feature tensor is dense and the adjacency matrix is sparse). Previously, people rely on vendor-specific libraries (e.g., MKL, cuBLAS) to obtain high performance of tensor computations over the vendors’ own CPUs and GPUs. These libraries require heavy manual tuning by experienced engineers. As a result, they evolve slowly in contrast with the rapid emergence of new workloads.

An alternative solution is tensor compilation, which expresses the processing of tensors in its own intermediate representation (IR) [31, 24]. Tensor compilation separates the computation definition (i.e., what to compute) from the scheduling (i.e., how to compute) so as to focus on the scheduling part for performance optimization. A scheduling scheme can apply loop transformations, vectorization, thread binding, etc., to manipulate the tensor computation. Optimizing one computation kernel for different hardware architectures is essentially searching for different scheduling schemes.

FeatGraph adopts the tensor compilation approach to optimize the computation kernels in GNN workloads. However, existing tensor compilers [24, 31] mostly focus on computations over dense tensors, and there is little support for computations involving sparse tensors. FeatGraph extends TVM [24] to support the core sparse patterns in GNNs (i.e., generalized SpMM and SDDMM), and allows customizable feature dimension computations on each vertex/edge by the design of a two-granularity programming interface (Sec III-B).

Figure 2 depicts the software stack of FeatGraph. At the top level, users define, train, and evaluate GNN models in specialized frameworks such as DGL and PyG, which handle dataflow programming and automatic differentiation. FeatGraph serves as a backend for these frameworks, targeting the message passing computation that is core to GNN workloads. FeatGraph provides a flexible programming interface to express the diverse variants allowed by the message passing paradigm. Specifically, FeatGraph describes feature dimension computations on each vertex/edge with user-defined functions (UDFs), and triggers UDFs by SpMM or SDDMM sparse template. FeatGraph incorporates optimizations for graph traversal into the sparse templates, and allows users to specify optimizations for UDFs with a feature dimension schedule (FDS). FeatGraph combines templates with FDS, and leverages the TVM tensor compiler [24] to generate efficient kernels for both CPU and GPU. By decoupling these two levels of optimizations, FeatGraph significantly improves the productivity of developing new kernels for emerging GNN models.

There are two principles in the design of FeatGraph’s programming interface. First, the interface should closely follow the mathematical definition of GNNs as described in Section II-A. Second, it should facilitate optimizations.

To these ends, we propose to decompose a kernel specification into two parts: UDFs written in a tensor expression language adopted from TVM to describe fine-grained feature dimension computations on each vertex/edge, and the choice of coarse-grained sparse patterns. FeatGraph provides two kernel templates featgraph.spmm and featgraph.sddmm for the SpMM and SDDMM sparse patterns that directly map to the vertex-wise and edge-wise computations in the message passing paradigm, i.e., Equations (1) and (2).

More concretely, featgraph.spmm takes in five arguments: an adjacency matrix, a message function, an aggregation function, the target (CPU or GPU), and an FDS to specify optimizations of the message function. Figure 3(a) shows the code for GCN aggregation, i.e., the message aggregation in GCN model as described in Section II-A. Given the edge ID tuple (src, dst, eid), the user-defined message function msgfunc (line 6–8) slices out the src row from the vertex feature matrix XV, which is equivalent to using the source vertex feature as the message. The aggregation function is sum and any commutative reducer is allowed. Figure 3(b) shows a more complex message function, which adds the source and destination vertex features, and then multiplies with a weight matrix, followed by a ReLU activation (i.e., taking the max with 0).

FeatGraph can easily support the commonly used message functions in GNNs—specifically, all the builtin ones provided by DGL¹¹1https://docs.dgl.ai/api/python/function.html#message-functions, including copying vertex or edge feature tensor, element-wise operations between vertex and edge feature tensors, etc. In addition, FeatGraph can express more complex message functions such as the one in MLP aggregation.

FeatGraph inlines the message function into the SpMM template to generate a fused kernel. In contrast, existing GNN frameworks (e.g., DGL, PyG, NeuGraph) that rely on deep learning systems as backend have to materialize the messages on every edge, causing inefficiency in both performance and memory consumption.

featgraph.sddmm takes in four arguments: an adjacency matrix, an edge function, the target (CPU or GPU), and an FDS to specify optimizations of the edge function. Figure 4(a) shows the code for dot-product attention, where the user-defined edge function edgefunc (line 6–10) performs a dot product between the source and destination vertex feature vectors, and returns an attention weight as the new feature on the edge. Figure 4(b) shows a more complex edge function, which performs multiple dot products over the feature tensors.

This two-granularity programming interface simplifies implementing new GNN kernels and, more important, facilitates optimizations. By cleanly decomposing a kernel specification into coarse-grained sparse templates and fine-grained feature dimension computations on each vertex/edge in the form of UDFs, FeatGraph enables decoupled, two-level optimizations. Specifically, FeatGraph incorporates optimizations for graph traversal into the sparse templates and allows users to specify optimizations for UDFs with an FDS. Some FDS examples, both for CPU and for GPU, are shown in Figure 3(a) at line 11–15 and line 19–22. It is worth noting that when the FDS is missing, FeatGraph essentially degrades to traditional graph processing systems that are designed without special handling of feature dimension computation.

This subsection describes the optimizations for graph traversal, which are incorporated into the sparse templates, and the optimizations for feature dimension computation, which are specified by users with an FDS. We analyze the interplay between these two levels of optimizations, and show that by combining them, FeatGraph enables performant processing of GNN workloads. Throughout this subsection, we use the sample graph shown in Figure 5 to illustrate optimization techniques.

We implemented the SpMM and SDDMM templates as TVM IR templates. TVM is a domain-specific language and compiler for tensor computations and has been widely adopted to accelerate deep learning workloads [35, 36]. Because TVM does not support sparse representation and computation in its tensor expression language, we implemented and optimized SpMM and SDDMM templates by directly constructing and manipulating the IR (intermediate representation) using lower-level APIs. Feature dimension computations on each vertex/edge described by UDFs are dense and therefore easily supported. FeatGraph combines scheduling parameters from the sparse templates (e.g., number of graph partitions, number of CUDA blocks) and those from the FDS (e.g., feature dimension tiling factors) to create the design space. In this work we use naïve grid search to find the optimal parameters under a given input shape, and it is an interesting future direction to try more intelligent tuners [37, 38] for faster design space exploration. After performing optimizations for both the templates and UDFs, FeatGraph inlines UDFs into the templates to generate fused kernels.

We parallelize the kernels over multiple threads on CPU using the customized thread pool [35] in TVM runtime, which is lightweight and particularly efficient in handling the kind of embarrassingly parallel workloads. To avoid LLC contention after graph partitioning, we assign multiple threads to collectively work on one graph partition at a time instead of assigning each thread to a different partition.

In order to evaluate the performance of FeatGraph in end-to-end GNN training and inference, we integrated FeatGraph into DGL, a popular open-source GNN framework. DGL implemented a minimal Gunrock-like graph kernel interface named Minigun [39]. With Minigun, DGL provided a set of builtin message functions and edge functions to support common GNN workloads. For each of these builtin functions, we implemented a corresponding one with the programming infrastructure of FeatGraph, such as GCN aggregation and dot-product attention. To handle more complex cases such as MLP aggregation, the current solution in DGL is to calculate and materialize the messages on every edge using deep learning systems as backend. In contrast, FeatGraph generates fused kernels, thus both saving memory and improving efficiency. FeatGraph generates kernel codes for a specific graph topology (i.e., the adjacency matrix); since GNN training typically involves hundreds of epochs, the compilation cost is amortized and negligible.

The integration requires a small amount of effort (only $\sim$ 300 lines of Python code) because both FeatGraph and DGL follow the message passing paradigm in their programming interface design. The integration with DGL demonstrates that it is straightforward to have FeatGraph be the backend to accelerate GNN frameworks in general, including PyG, NeuGraph, etc.

Environment. For CPU evaluation, we conduct experiments on Amazon EC2 c5.9xlarge instance, which is a one-socket 18-core 3.0 GHz Intel Xeon Platinum 8124M machine with 25 MB LLC and 68 GB DRAM. For GPU evaluation, we conduct experiments on p3.2xlarge instance, which has a Tesla V100 GPU with 80 SMs; each SM has shared memory configurable up to 96 KB (the default size is 48 KB).

Datasets. Table II lists the datasets used for evaluation: ogbn-proteins represents proteins and their biological associations with vertices and edges—this dataset is from Open Graph Benchmark²²2https://ogb.stanford.edu/, a realistic benchmark suite for GNNs; reddit [40] is constructed from the Reddit online forum wherein vertices represent posts and edges are established if two posts are commented by a same use—this dataset is commonly used in GNN research for evaluating the accuracy of new models; rand-100K is a synthetic graph wherein 20K vertices have an average degree of 2000 and the remaining 80K vertices have an average degree of 100—this dataset is specifically aimed at studying the effect of hybrid partitioning on GPU performance.

Baselines. We compare FeatGraph with state-of-the-art graph processing systems, specifically Ligra on CPU and Gunrock on GPU. We also compare with vendor-provided sparse libraries, specifically MKL (2019.5) on CPU and cuSPARSE (10.1) on GPU whenever possible, as only a subset of GNN kernels are supported in these libraries. In all the experiments, we first do a warm-up run and then take the average time of 10 runs as the measurement.

We evaluate the performance gain of FeatGraph on three kernels: GCN aggregation, MLP aggregation, and dot-product attention. The kernels are performed on the full graph. We do the evaluation across a spectrum of feature lengths. For MLP aggregation, the feature length refers to d2 that is shown in Figure 3(b); d1 is fixed as 8.

Single-threaded CPU Performance. Table III shows that across all the evaluated datasets under different feature lengths, FeatGraph achieves 1.4$\times$–4.0$\times$ speedup over Ligra on GCN aggregation, 4.4$\times$–5.5$\times$ speedup on MLP aggregation, and 4.3$\times$–6.0$\times$ speedup on dot-product attention, using a single thread. Compared against MKL on GCN aggregation, FeatGraph is faster in 14 out of 15 cases and achieves higher speedup with a larger feature length. Specifically, when the feature length is 512, FeatGraph is 1.8$\times$ faster on ogbn-proteins, 2.4$\times$ faster on reddit, and 4.4$\times$ faster on rand-100K. MKL does not support MLP aggregation and dot-product attention.

Multi-threaded CPU Performance. Figure 10 shows that with 16 threads, for GCN aggregation on reddit, FeatGraph achieves 12.6$\times$ speedup over its single-threaded execution, which is slightly higher than Ligra (9.5$\times$) and MKL (9.8$\times$). Similar observation applies to other datasets and kernels. As a result, FeatGraph outperforms the others consistently in multi-threaded environment. FeatGraph scales well due to two factors: 1) its parallelization method avoids LLC contention by assigning multiple threads to collectively work on one graph partition at a time; 2) the thread pool in TVM runtime is lightweight [35].

GPU Performance. Table IV shows that FeatGraph is 24$\times$–206$\times$ faster than Gunrock on GCN aggregation, 18$\times$–96$\times$ faster on MLP aggregation, and 1.2$\times$–3.1$\times$ faster on dot-product attention. The extreme slowness of Gunrock on GCN aggregation and MLP aggregation is caused by two reasons: 1) Gunrock’s edge parallelization execution incurs huge overhead of atomic operations for vertex-wise reductions such as GCN aggregation and MLP aggregation; 2) Gunrock fails to exploit parallelism in feature dimension computation. FeatGraph is on par with cuSPARSE on GCN aggregation, being 10%–20% faster on ogbn-proteins and rand-100K while 10% slower on reddit. Notably, cuSPARSE does not support MLP aggregation and dot-product attention.³³3 The latest cuSPARSE supports dot-product attention via ConstrainedGeMM.

This subsection investigates the performance boost of each individual optimization technique described in Section III. For the sake of space, in each ablation analysis we only pick one dataset to show the optimization effects. Other datasets share similar observations.

Graph Partitioning and Feature Dimension Tiling. Figure 13 shows that feature dimension tiling combined with graph partitioning effectively boosts the performance of GCN aggregation on CPU. Specifically, when the feature length is 512, feature dimension tiling alone and graph partitioning alone bring 1.2$\times$ speedup and 1.7$\times$ speedup, respectively; combining two achieves 2.2$\times$ speedup.

Adaptive Parallelization Strategies on GPU. Figure 13 shows that tree reduction boosts the performance of dot-product attention by up to 2$\times$. The naïve parallelization strategy in Gunrock that assigns the entire dot product operation on each edge to one CUDA thread is less efficient in handling large feature lengths due to consuming too many registers per thread.

Hybrid Partitioning on GPU. Figure 13 shows the effect of hybrid partitioning on GCN aggregation tested on rand-100. FeatGraph gets 10%–20% performance boost by hybrid partitioning, and consequently outperforms cuSPARSE.

Sensitivity to Partitioning Factors. Figure 14 shows that the performance of FeatGraph is sensitive to partitioning factors for GCN aggregation on CPU. Specifically, on reddit, when the feature length is 128, the best performance is achieved with 16 graph partitions and 4 feature partitions. On the same graph, as the feature length increases, the optimal number of feature partitions increases proportionately, while the optimal number of graph partitions stays constant. Transferable tuning across graphs, i.e., using the optimal partitioning factors tuned on one graph to predict the optimal partitioning factors for a new graph, is more challenging and worth further study.

Sensitivity to GPU Parameters. Figure 15 shows that FeatGraph performs better with a larger number of CUDA blocks for GCN aggregation on GPU, because a larger number of CUDA blocks can better utilize the massive parallel compute capacity of GPU. In the evaluation, we set the number of CUDA blocks to the number of rows of the adjacency matrix.

Sensitivity to Graph Sparsity. Table V shows that FeatGraph achieves higher speedup over MKL as the graph sparsity decreases for GCN aggregation on CPU. This trend is because a denser graph has more data reuse, which FeatGraph is able to exploit by graph partitioning and feature dimension tiling.

We integrated FeatGraph into DGL and evaluated the performance of FeatGraph in end-to-end GNN training and inference on three models: a 2-layer graph convolutional network (GCN) [26] of hidden size 512, a 2-layer GraphSage [40] of hidden size 256, and a 2-layer graph attention network (GAT) [27] of hidden size 256. GCN uses sum aggregation and requires generalized SpMM computations in both forward and backward propagation; GraphSage follows a similar architecture as GCN but allows more flexible aggregation functions (e.g., max); GAT uses dot-product attention, thus requiring both generalized SpMM and SDDMM computations.

Accuracy. FeatGraph as a backend is for performance optimization without changing the semantics of GNN models. As a sanity check, we evaluate the accuracy of the three models on the task of vertex classification. The 233K vertices of the reddit dataset are split into 153K, 24K, and 56K for training, validation, and testing, respectively. We train the models for 200 epochs. The testing accuracy obtained by DGL using FeatGraph matches that obtained by DGL using its original backend Minigun—93.7% for GCN and 93.1% for GraphSage. The training of GAT does not converge due to gradient explosion, with either FeatGraph backend or Minigun backend.

Speedup. Table VI reports the training and inference time of one epoch for the three GNN models. The time of tuning partitioning factors is excluded, because it is amortized over multiple epochs—it is less than 1% of the time of training GCN for 200 epochs. Furthermore, the partitioning factors tuned on GCN are directly applied to GraphSage and GAT—the number of graph partitions is kept the same and the number of feature partitions is adjusted to the feature length. The results show that on CPU, FeatGraph speeds up both training and inference by more than 20$\times$ on all the three models; on GPU, FeatGraph speeds up training by more than 2$\times$, and inference by 1.4$\times$–7.1$\times$. The highest speedup is achieved on GAT, which has a more complex architecture than GCN and GraphSage.