Hector: An Efficient Programming and Compilation Framework for Implementing Relational Graph Neural Networks in GPU Architectures

GNNs are feed-forward neural networks that propagate and transform features using the connectivity of a graph. A widely used model is graph convolutional network (GCN) (kipf2016semi, ). Formally, a GCN layer is defined as $\overrightarrow{{h}^{(l+1)}}=\sigma\left(A^{*}\overrightarrow{{h}^{(l)}}W^{(l)}\right)$ , where $W^{(l)}$ denotes a trainable weight matrix of the $l$-th layer, $\sigma$ is an activation function and $\overrightarrow{{h}^{(l)}}$ is the $l$-th layer node representation. $A^{*}$ is the adjacency matrix normalized by node degrees. $\overrightarrow{h^{(0)}}$ denotes the node input features.

RGNNs extend GNNs to model different node and edge types for relational graph data. For example, extended from GCN, a relational graph convolutional network (RGCN) layer is defined as

, where $\mathcal{N}_{v}^{r}$ denotes neighbors of node $v$ in relation $r\in R$, ${h}_{n}^{(l)}$ is the $l$-th layer node representation of $n$. $W_{r}^{(l)}$ is the weight for relation $r$. $c_{v,r}$ is a problem-specific normalization factor. Figure 1 shows an example of how output features is produced in the message passing formulation equivalent to Formula 1: The forward propagation of an RGNN layer could be divided into ① the edge message generation stage and ② the node aggregation stage. For simplicity, we focus on the output feature $\overrightarrow{h_{z}^{(out)}}$ of node $z$: To obtain $\overrightarrow{h_{z}^{(out)}}$, ① a message $\overrightarrow{msg}$ is generated for each incoming edge, and ② the edge messages go through weighted aggregation and an activation function $\sigma$ to produce $\overrightarrow{h_{z}^{(out)}}$. Notably, to obtain the output feature of node $v$, the input feature of $v$ itself is applied to the $W_{0}^{(l)}$ and added to the transformed neighbor features. We call this a virtual self-loop because it could be seen as if each node now has a new edge to itself.

Relational graph attention network (RGAT) (busbridge2019relational, ) and heterogeneous graph transformer (HGT) (hgt, ) are shown in Figure 2. Attention is introduced in these more complex models: Attention is produced in the message generation stage together with edge messages. Similar to the normalization factor, it is a scalar that emphasizes the message associated with the same edge during the subsequent node aggregation stage. However, attention is learned, as is produced by operations among weights and features.

In non-graph neural networks, most linear operators, e.g., convolution, can be efficiently implemented with GEMM kernels. GEMM takes up most of the execution time due to its cubic complexity. While some operators can be optimized by transformations, e.g., Winograd for convolution layers (lavinFastAlgorithmsConvolutional2015, ), the operators are still computation-intensive after such computation reduction. GPUs are excellent at GEMM because the latter’s high computation complexity allows leveraging the massive parallel compute units on GPUs, while the input data could be sufficiently reused to allow the memory bandwidth to keep up with the computation throughput.

In contrast, GNNs spend a much larger portion of their execution time on memory-intensive, non-GEMM operations  (wangEmpiricalAnalysisPerformance2021, ; zhengNatureGraphNeural2021, ). One major source of memory-intensiveness is the sparsity of graphs: to be not bound by the memory bandwidth, Nvidia H100 GPU requires the data reuse of single-precision float to be at least 16 times. However, the average degree of a graph often falls below this threshold, e.g., the graph datasets in Table 3. The heterogeneity of RGNNs further exacerbates the issue due to lowered data reuse by the introduction of dedicated weights to different edge types and node types, as shown in Figure 2.

We use an edgewise typed linear layer as an example to walk through the various performance overheads in the existing computation stack, as summarized in Figure 4. Edgewise typed linear layer applies a typed linear operator on each edge to one of its vectors. The weight of the linear operator used in the computation depends on each edge’s type. For example, the edge message in an RGCN layer (Figure 1) or an RGAT layer (Figure 2), is produced by a typed linear layer.

A typed linear layer is typically implemented using batched matrix multiply (BMM) or segment matrix multiply (segment MM) (isratsegmm, ). For example, PyG FastRGCNConv implemented typed linear layers using BMM to unleash parallelism. However, a temporary tensor must be created from the weight tensor due to the lack of support for indirect addressing by PyTorch tensor APIs: the typed linear layer could be denoted as $Y[i,0,j]:=\sum_{k}(X[i,0,k]\times W[T[i],k,j])$ where $X[i,0,\cdot]$, $Y[i,0,\cdot]$ and $W[T[i],\cdot,\cdot]$ are input feature, output feature of node $i$ and the weight of node $i$’s type. The middle dimension of $X$ and $Y$ are needed to make the operation a matrix multiply. However, there is currently no support for specifying $T[i]$ as one of the arguments to an operator in PyTorch; one must create $W^{\prime}[i,k,j]:=W[T[i],k,j]$ before the typed linear layer.

Segment MM requires presorting features by types. After that, the node/edge feature tensor is in the form of segments of features of the same type: the segment MM kernel then applies the corresponding weight tensor of the type to each segment. If neither BMM nor segment MM can be employed, one may fall back to multiple matrix multiplies, leading to higher device API overhead and GPU under-utilization.

Another type of inefficiency is suboptimal math library calls. PyTorch has routines to handle various scenarios, e.g., a tensor is strided in memory layout or is NestedTensor, a pack of tensors. Consequently, Pytorch sometimes performs BMM by launching multiple general matrix-vector multiplies (GEMVs) kernels, which also leads to API overhead and GPU under-utilization.

Lastly, CUDA math libraries were initially developed for large inputs and may not be efficient for small inputs (nvidiaCublasgemmBatchedCuBLASDocuemntation, ).

To better illustrate the points, Figure 3 breaks down HGT and RGAT inference time on FB15k and MUTAG. Section 4.1 details the system configurations and datasets. This experiment measured Graphiler (xieGraphilerCompilerGraph, ), which executed compiled TorchScript code and delivered the best inference performance among the existing systems tested in this work. Figure 3 shows that indexing and copying take up a significant portion, and the portion of GEMM operations, i.e., MM vs. Other compute, varied with graphs. By profiling, we found that the CUDA API overhead is 22% the time of the critical path, which is the sum of the API overhead and kernel duration. This is partly due to a huge number of kernel launches caused by 1) libraries calling a series of kernels to fulfill an API invocation and 2) some operators calling separate sets of kernels for each types in the graph.

In contrast, Hector 1) lowers more of the logic to GEMM, and 2) assembles kernels with flexible access scheme to gather and scatter data on the fly to eliminate redundant data movement. Consequently, Hector does not replicate weights during computation. As shown, this strategy achieves better performance than using hand-optimized kernels with dedicated functions to data movement, e.g., in Graphiler.

To address the performance challenges in RGNN systems due to both RGNN’s inherent computation pattern and the system design, we propose the Hector IR and code generation framework. By the IR design that decouples and expresses the model semantic, data layout, and operator-specific schedules, Hector opens up these opportunities and the integration of all three aspects into the design space. Table 1 shows the feature comparison of Hector with existing systems.

Hector consists of a programming interface, a code generator, and Python modules. The code generator takes in the model definition and generates both CUDA kernels and host functions that configure and invoke the CUDA kernels.

Figure 5 uses an example to illustrate the workflow. The input is an excerpt of DGL code invoking a typed linear layer on the input node features. Applying the @hector.compile decorator triggers a transpiling pass to lower the code into Hector inter-operator level IR. In this example, the typed linear transformation typed_linear can be efficiently implemented as GEMM kernels. To this end, Hector lowers the transform to an operator instance derived from the GEMM template at the inter-operator level. After the analysis and optimizations at the inter-operator level, Hector further lowers the code to a detailed GEMM specification at the intra-operator level. The GEMM output $A$ collects edge data generated from the node data. The first input $B$ is the weight matrix $W$, and the second input $C$ is the collection of features of all the source nodes of the edges involved. The intra-operator level IR indicates that the GEMM operation should use the default tile width of 16 and be carried out without scatter, gather, or transpose applied to input or output matrices. Eventually, Hector generates a segment MM (Section 2.3) kernel, gemm_1. The Layout Choices section of Figure 5 shows the default layout choice. etype_ptr specifies the offsets of each segment of different type. row_idx is the source node index array in the COO format. The result tensor e["msg"] has the number of edges as the number of rows, and the number of the columns is the input dimension of the hidden layer. We detail in Section 3.2.2 an optimization technique, compact materialization, that is opened up by the decoupled layout choices from the inter-operator level IR.

The generated code is compiled into a shared library where host functions are exported through the pybind11 utilities. Hector falls back to existing routines in PyTorch when certain operators are not yet supported. During runtime, the precompiled functions are loaded and registered as subclasses of PyTorch autograd.Function.

The inter-operator level IR follows the Python grammar but involves some new constructs, as listed in Table 2. Listing LABEL:lst:ir_example illustrates how the attention calculation in a single-headed RGAT layer could be expressed using the inter-operator level IR. Lines 10-16 shows a code segment that generates attention values for all edges of graph g and then invoke the edge_softmax(g) function that spans lines 1 through 9. As shown in Listing LABEL:lst:ir_example, the message generation and aggregation stages are expressed as for-each edge loops starting from line 2, line 8 and line 10, and for-each node loop starting from line 4. To accumulate data from the incoming edges of each node n, the n.incoming_edges() iterator is used. Notably, the data layout that specifies how to access the input and output data per edge or node as well as the incoming edges associated with each node, is abstracted away in Listing LABEL:lst:ir_example.

The intra-operator level IR serves between the inter-operator level IR and the generated CUDA code. At this level, the IR should encode specifications to emit CUDA code and provide sufficient information specific to each operator invocation to the transform and lowering passes at the inter-operator level. The code transformation components at this level provide the methods to generate specialized CUDA code for the operators, to apply operator-specific schedules, and to return necessary information on operator selection and kernel fusion feasibility to the passes at the inter-operator level.

Hector’s code generator ultimately lowers the IR to two basic constructs, the GEMM template and the traversal template. Algorithms 1 and 2 illustrate the edge traversal template and the GEMM template. The node traversal template is similar to Algorithm 2, and we will revisit it in Section 3.4.1. For simplicity, function template specialization refers to routines specialized for the specific instances derived from the two templates and involve 1) function arguments, e.g., number of rows, etc., 2) special registers, e.g., threadIdx, and 3) loop variables.

Central to the code generator is the two-level IR. Inter-operator level IR optimizations address the opportunities brought in by heterogeneous relation types. These optimizations manipulate operators and their connections. A high-level IR abstracts away the low-level details that can complicate or even hinder the transformations. Intra-operator level IR optimizations reduce the data movement by generating access schemes in kernels rather than using specialized kernels and dedicated indexing/copying kernels. These optimizations manipulate low-level data access and schedule details, and thus are better supported by a low-level IR.

The two-level IR enables concerted but decoupled choices of intermediate data layout and compute schedules: For example, in Figure 5, the semantics of the model are decoupled from the layout choices. Hector implements the model semantics and layout choices in intra-operator level IR with specific access schemes. The next few paragraphs explain how the two-level IR design facilitates operator-specific optimizations, operator selection, and kernel fusion.

Similarly to PyTorch, Hector supports auto-differentiation by maintaining the backward propagation counterparts of the operators. Hector first emits the backward propagation via inter-operator level IR, and removes unused gradients and their computation. The lowering and code generation schemes are similar to those in forward propagation. However, additional processing is needed because the PyTorch auto-differentiation requires the backward propagation methods to be paired with the forward propagation methods in the autograd.Function definitions. To achieve this, Hector bookkeeps the kernel calls in each forward propagation method. For each forward propagation method, Hector puts all the corresponding backward propagation kernel calls in the body of the backward propagation method.

The code generation procedure emits code based on the CUDA kernel specifications detailed in the form of intra-operator IR. Kernel code generation is fairly straightforward and is implemented using a template-based approach. Hector then emits the host functions that configure grids and blocks, gets raw pointers from the libtorch at::Tensor references, and launches the corresponding kernel. The host functions are exported via pybind11 utilities.

The Hector performs a pass that scans all the functions generated to collect a list of preprocessing required for the input dataset, involving transposition, converting COO to CSR, etc. The code generator then emits the preprocessing code.

Linear operator reordering and compact materialization are specific to RGNNs. Linear operator reordering is specific to RGNNs because RGNNs typically require linear projections from different semantic spaces, introduced by the heterogeneity of node types and edge types, to a common space before further operations. Compact materialization is specific to RGNNs because of the additional tensor dimension brought in by different node types and edge types.

Some of the intra-operator IR optimizations could benefit ordinary GNNs, which can be treated as a special case of RGNNs whose relation type number is one. Intra-operator level IR allows specification of both data access schemes and schedules, thus allowing flexible code generation to accommodate different dense or sparse tensor layouts, a need that often arises from compact materialization. However, the ability to generate code for different data access schemes and schedules can be beneficial when compiling ordinary GNNs.

To assess performance, we measure the inference and training time of Hector and other systems on a single-GPU computer. Its hardware components include one Intel Core i9-9900K CPU, 128 GB dual-channel memory, and one Nvidia RTX 3090 GPU with 24GB memory. The operating system is Ubuntu 18.04.5, with kernel version 5.4.0-135. The CUDA and driver versions are 12.1 and 530.30.02, respectively. PyTorch and DGL versions are 2.0.1 and 1.1.1, respectively.

As shown in Table 3, we use public datasets from DGL (wang2019deep, ) and OGB (huOpenGraphBenchmark2021, ). We measure (1) inference and (2) training time on three RGNN models, RGCN (rgcn, ), RGAT (busbridge2019relational, ), and HGT (hgt, ), comparing with previous systems, involving DGL (wang2019deep, ), PyG (fey2019fast, ), Seastar (wuSeastarVertexcentricProgramming2021, ), Graphiler (xieGraphilerCompilerGraph, ), and HGL (guiHGLAcceleratingHeterogeneous, ). We ported Seastar artifacts to the same version of CUDA and Python packages as what Hector depends on because one of Seastar’s dependencies, dgl 0.4, used an API deprecated since CUDA 11.

For RGCN, RGAT, and HGT, excluding comments, Hector took in 51 lines in total and produced more than 3K lines of CUDA kernel code, 5K lines of other C++ code to define host functions, and 2K lines of Python code to define subclasses of Pytorch autograd.Function. The implementation also involves 2K lines of Python code providing common utilities.

To best align with the hyper-parameters prior work used in its evaluation, we set the input and output feature dimensions as 64 and the number of heads as 1. We measure the inference and training time of the single layer used. In training, to obtain a loss, we compute the negative log-likelihood loss by comparing the output with a precomputed random label tensor. For each case, we run the full graph inference and training for at least 10 epochs and average the elapsed time. To align with the existing system, nodes are presorted to enable segment MM for typed linear layers.

For the performance of DGL and PyG, we measure all public implementations of these models from DGL, PyG, and Graphiler artifacts. PyG provides two RGCN convolution layers: RGCNConv places nodes in segments of the same type but launches separate kernels for each of the node types, leading to device underutilization. FastRGCNConv replicates weights and uses bmm(). It is consistently faster than the RGCNConv implementation. Similarly, DGL’s built-in segmentMM-based RGCN layer is faster than other DGL implementations. For HGT, the DGL segmentMM-based HGTConv primitive generally has the best performance. In the cases where some variants encounter OOM errors, we choose the best among those that run without issues. Some cases are missing due to insufficient operator support, such as HGL on HGT and Graphiler on training. We do not measure HGL in inference because it is designed to optimize training.

Figure 8 shows that Hector’s best-optimized code consistently outperforms state-of-the-art systems. It achieves up to 9.9$\times$ speed-up in inference and up to 43.7$\times$ speed-up in training against the best of state-of-the-art systems. On geometric average, Hector gets 1.79$\times$, 8.56$\times$, 2.87$\times$ speed-up in inference via RGCN, RGAT, and HGT, respectively, and 2.59$\times$, 11.34$\times$, 8.02$\times$ speed-up in training RGCN, RGAT, and HGT, respectively. The performance advantage is larger in small graphs, demonstrating that generating a single kernel that performs the computation across multiple edge types boosts the performance on small graphs compared to existing systems that run many small kernels.

We see close performance achieved by Graphiler in RGCN and HGT inference. Graphiler leverages PyTorch utilities to produce TorchScript binaries before execution and utilizes edgewise parallelism for edgewise computation. Similarly to RGCNConv, it places node features into segments of the same type but runs separate kernels to perform a typed linear transformation. DGL and PyG under similar configurations achieve competitive performance. However, when it comes to RGAT, Graphiler suffers from performance degradation. Because Graphiler relies on pre-programmed fused kernels to deliver a significant portion of the performance boost (xieGraphilerCompilerGraph, ), we postulate that the degradation is due to the non-exhaustiveness of these pre-programmed kernels (xieGraphilerRepositoryGithub2023, ). This reflects the drawbacks of compiler design without a code generation mechanism. By contrast, with two-level IR and a code generator, Hector achieves better performance, showing that generating kernels with flexible access scheme that gather and scatter data on the fly eliminates redundant data movement and outperforms indexing/copying followed by hand-optimized GEMM and sparse kernels. Besides, it is challenging to extend Graphiler’s approach to training due to TorchScript’s limited auto-differentiation support. For example, dict object creation is not supported, but it is a common way to express nodewise data and edgewise data.

By comparing Hector with Seastar, which lowers all logic to sparse kernels, we realize that sparse kernel code generation alone is not efficient in RGNNs: it is better to lower to GEMM kernels as much as possible.

There are two reasons why Hector is more efficient in device memory usage. First, Hector only keeps a single copy of weights, as discussed in Section 3.2.2. Replicating weights also affects backward propagation because the gradient of each individual copy will be derived, occupying extra memory. Second, our compact materialization reduces memory and redundant computation, as explained in Section 4.3.

Notably, even without compact materialization or linear operator reordering, Hector still consistently outperforms existing systems, as Table 4 shows. In addition, the unoptimized Hector code triggers fewer OOMs than existing systems, with the only exception where the RGAT inference is run on mag and wikikg2. For comparison, we also show the statistics of the best optimized Hector code in Table 4.

Now, we study the effects of compact materialization and linear operator reordering. They are detailed in Sections 3.2.2 and 3.2.3. We investigate their effects on RGAT and HGT.

Table 5 shows the speed-up on top of Hector unoptimized code by these two optimizations. Due to compact materialization, when running RGAT on mag and wikikg2, Hector no longer triggers OOM errors. In addition, in some cases, the layout speeds up the execution as well due to the common subexpression elimination brought forth by the layout. Compact materialization is hardly possible without a code generation scheme or an IR design that decouples the model semantics, data layout, and operator-specific schedule. Besides, data layout choice, compact materialization in particular, allows further performance enhancement while prior work usually focuses on the improvement of schedule given a specific sparse matrix format. This is shown by the great speedups in the “C[ompact]” columns in Table 5.

To study how compact materialization reduces the memory footprint, We illustrate the Hector dram usage without compact materialization in Figure 10(b) and the portion of dram usage with compact materialization in Figure 10(a). For simplicity, we define the entity compaction ratio as the number of unique $(\text{source node},\text{edge type})$ pairs divided by the number of edges. Figure 10(b) shows that the memory use of inference and training is highly proportional to the number of edges of the datasets. Figure 10(a) shows that compact materialization significantly reduces DRAM usage in all datasets. The memory footprint ratio of compact materialization compared with the memory footprint of the unoptimized code correlates with the entity compaction ratio. The memory footprint ratio is higher than the entity compaction ratio, as the memory footprint consists of edgewise data, nodewise data, and weights, whereas the compaction applies to edgewise data only. Besides, in case the average degrees are larger, the memory footprint ratio reduces more significantly, getting closer to the entity compaction ratio.

To better understand the performance benefits of optimizations, Figure 9 studies two cases. The entity compaction ratio of AM and FB15k are 57% and 26%, respectively. On AM, the time GEMM instances take is greatly reduced. By comparison, in FB15k, compaction brings less performance improvement due to the less significant GEMM reduction.

In short, due to the data-dependent nature of computation in RGNNs, there is no one-size-fits-all optimization strategy. However, as shown in Table 5, Enabling compaction and reordering obtains fairly good performance consistently and is the best fixed strategy on average in all four scenarios, i.e., $\{\text{RGAT},\text{HGT}\}\times\{\text{training},\text{inference}\}$. If Hector presumably chooses the best configuration in every run, it could further get 1.06$\times$, 1.33$\times$, 1.02$\times$, and 1.08$\times$ speed-up in the four scenarios above, respectively. We leave autotuning to future work.

We show the average time of unoptimized Hector in Figure 11. We also further profile generated kernels when running Hector on RGAT on bgs and am, as shown in Figure 12.

One thing to note is the sublinear time increase in Figure 11: when the input and output dimension doubles, the amount of computation and memory accesses becomes close to 4$\times$ those of the original, but the time increase is typically lower than 2$\times$ of the original. The reason is increased computation throughput when the size increases, as corroborated by Figure 12. Moreover, we observed higher throughput when the graph scale increases, e.g., from bgs to am in Figure 12. Similarly, we witnessed the cuBLAS throughput increases steadily when we keep the right matrix size as (64, 64) and increase the number of rows of the left matrix from 1M (2¹⁷) to 8M (2²⁰). These suggest that an RGNN system should be memory-efficient in order to accommodate larger models and datasets to fully utilize the massive resources on GPUs. By eliminating unnecessary data copies, Hector achieves better memory efficiency than state-of-the-art systems.

The instruction per cycle (IPC) charts in Figure 12 indicate the traversal kernels are generally latency-bound: on RTX 3090, IPC is ideally 4 as each streaming multiprocessor (SM) has four schedulers. Backward propagation kernels have lower throughput due to worsened latency and increased memory bandwidth consumption by doubled memory accesses compared to forward propagation. In backward propagation, backward traversal kernels compute gradients using atomic updates, therefore hindering the throughput; GEMM kernels also on average have lower performance due to outer products that compute the delta of weights.