GNNBuilder: An Automated Framework for Generic Graph Neural Network Accelerator Generation, Simulation, and Optimization

GNNBuilder aims to provide users with a streamlined process to design, implement, validate, and optimize GNN models, transitioning from standard PyTorch models to FPGA bitstreams. As depicted in Fig. 1, GNNBuilder consists of five components: Compiler front-end parses the native PyTorch GNN model definition, including the number of GNN layers, layer type, activation type, data precision, pooling type, aggregation type, and MLP definition. Code generator builds upon a library of pre-defined hardware accelerator templates that adopt the message passing mechanism, making it compatible with various GNN types. We generate High-Level Synthesis (HLS) code targeting FPGAs, supported by the Xilinx Vitis HLS tool [36]. Design space exploration and performance model enables automated DSE for accelerator generation, encompassing hardware parallelism, resource allocation, and quantization (data precision). Simulation and testbench facilitates transparent hardware-compatible simulation using automatically generated testbenches, ensuring the correctness of accelerator functionality. It also generates plain C++ code for ”true” quantization simulation, accurately reflecting on-FPGA quantization accuracy. Hardware synthesis and deployment automatically generates hardware synthesis scripts, synthesizes FPGA bitstreams, and produces host code for executing the bitstream. Table II presents the representative GNNs supported by our framework. Although these are examples, GNNBuilder can flexibly accommodate a wide range of customized GNN models, including residual and skip connections, arbitrary quantization, aggregation functions, graph attention, activation, global pooling, and MLP head. Such user-defined features can be naturally expressed using PyTorch, granting GNNBuilder exceptional extensibility.

Table III showcases the user APIs provided by GNNBuilder, and Listing 1 illustrates an example of the user interface for a customized GNN model.

To begin, a user defines a GNNModel instance, incorporating an MLP and a xxxConv_GNNB module (e.g., PNAConv_GNNB). GNNBuilder offers wrapper classes for each graph convolution layer, enabling the user to specify parallelism factors p_in and p_out. The higher-level GNNModel supports arguments for defining architecture parameters and separate parallelism factors for the GNN head (gnn_p_in, gnn_p_hidden, gnn_p_out) and the MLP head (p_in, p_hidden, p_out). The user can train and manipulate the GNNModel instance as a standard PyTorch module.

A user can then define a GNNBuilder Project instance. The Project class has several arguments to define build paths, the GNNModel model instance, the PyTorch Geometric dataset for the model task, max_nodes and max_edges, numerical precision, and average number of nodes, edges, and node in-degree for synthesis runtime estimation.

After creating a Project instance, the user can call the code generation functions to produce the model kernel HLS code, the kernel testbench code and data, the testbench makefile, and the Vitis HLS build script. Post code generation, the user can call build_and_run_testbench() to build and execute the testbench, and run_vitis_hls_synthesis() to execute the Vitis HLS synthesis process. These execution scripts also return data for the testbench runtime, mean absolute error (MAE), and synthesis latency / resource usage.

Inspired by GenGNN [34] and FlowGNN [35], we adopt an explicit message-passing architecture for graph convolution layers, allowing us to support GNN layers like PNA, which are not compatible with traditional SpMM accelerator approaches.

For each node, the operations illustrated in Figure 3 are performed. We first gather the node’s neighbor indices using the neighbor table and offset table. Then, we iterate through each neighbor index to load its associated embedding from the input node embedding table, transform the embedding with $\phi(\cdot)$, and aggregate it with a partial aggregation. After processing all neighbors, we finalize the partial aggregation, combine it with the current node embedding, and transform it with the apply function $\gamma(\cdot)$. The resulting computed embedding is then written to the output node embedding table.

The functions $\phi(\cdot)$, $\gamma(\cdot)$, and the aggregation(s) used depend on the specific layer being implemented. Kernels for GCN, GraphSAGE, GIN, and PNA layers are included in the initial GNNBuilder kernel library.

Developing custom kernels is possible by contributing the appropriate hardware kernel code for the layer of interest to GNNBuilder’s kernel template library, as well as providing a matching GNNConv class that links the kernel in GNNBuilder’s Python library. This can be accomplished through a pull request with minimal effort, and the rest of our framework remains agnostic to the specific types a user intends to add support for.

Graph Data: In the model kernel, buffers depend on the number of nodes (num_nodes) or edges (num_edges), with buffer sizes set to an upper bound determined by MAX_NODES and MAX_EDGES parameters in a GNNBuilder Project instance. Model kernels require input graphs in COOrdinate format matrix with an input node feature table. The COO matrix is a $\text{MAX\_EDGES}\times 2$ integer array, while the input node feature table is a $\text{MAX\_NODES}\times\text{input\_dim}$ fixed-point datatype array. In-degree and out-degree buffers have a size of MAX_NODES. Additionally, a neighbor table stores each node’s neighbors, and a neighbor offset table indexes each node’s block of neighbors, both sized MAX_EDGES and MAX_NODES, respectively.

Degree + Neighbor Table Computation: Before model computation, the degree table of the input graph must be calculated. Node degrees are used by various graph convolutions for normalization purposes. Since these values are only known at runtime, the in-degree and out-degree tables need to be computed on-the-fly in the accelerator for each input graph. The COO format of input graphs allows for computation within the bounds of num_edges. Subsequently, the neighbor table and neighbor offset table are computed simultaneously using two loops: one iterating over num_edges and the other over num_nodes.

Partial Aggregations: To efficiently aggregate neighbor embeddings with constant memory ($O(1)$ space complexity), GNNBuilder defines single-pass algorithms for aggregation that avoid the need for buffering all neighbor embeddings in an intermediate buffer, which would consume substantial BRAMs. GNNBuilder supports sum, min, max, mean, variance, and standard deviation aggregations. Each aggregation is associated with a data structure for storing partial and final aggregation data. For variance, Welford’s one-pass algorithm [37] is used to compute variance efficiently.

Linear Layer: GNNBuilder  implements tiled matrix multiplication for linear layers, enabling hardware parallelization. The parallelization factor for each linear layer is controlled by the BLOCK_SIZE_IN and BLOCK_SIZE_OUT template arguments for the linear kernel function. These arguments determine the partition factors for input, weight, and bias arrays, thus controlling the parallelism of the multiply-accumulate (MAC) operations.

Global Pooling: GNNBuilder supports sum, mean, and max global graph pooling, aggregating node embeddings across all nodes into a single embedding of the same size. Multiple pooling methods can be combined using concatenation.

Activations: GNNBuilder supports ReLU, Sigmoid, Tanh, and GELU [38] activations, implemented using fixed-point math functions from the Vitis HLS fixed-point math library.

GNNBuilderis built on a template-based compiler which facilitates code generation by generating C++ HLS code for the top-level model kernel and associated header directly from a PyTorch model. This enables conditional and loop control flows for template blocks, useful for features like skip-connections, double-buffer array selection, and mapping layer kernel calls in the correct order with accurate input/output size.

The parameterized structure of the GNNModel allows GNNBuilder to match appropriate function calls to corresponding kernels from the C++ header-only template library. This approach is extensible, enabling users to add support for other layers, aggregations, or activations by creating associated kernels in the template library and updating the Jinja template.

GNNBuilder allows designers to generate and build C++ testbenches for their models, facilitating rapid testing of fixed-point quantizations without synthesizing designs. The testbench code, model parameters, dataset graphs, true output, and PyTorch model outputs are exported as binary files. During runtime, the testbench reads these files, loads weights into the model kernel, evaluates the kernel on all dataset inputs, and compares the output to the PyTorch model outputs.

The testbench calculates verification metrics, such as mean absolute error between the PyTorch-generated model output and the kernel output, and averaged kernel runtime. These values are written to text files during runtime.

For fixed-point models, the testbench ensures accurate fixed-point representations of input graphs and model parameters. By utilizing the Vitis HLS fixed-point library [36], functional equivalence with hardware modeling is maintained. The floating-point data from PyTorch is cast to the user-specified fixed-point format in the testbench.

Using the run_vitis_hls_synthesis() function, users can build a synthesized accelerator (Verilog RTL code) and execute the implementation flow to generate Vivado IP blocks (.zip) or Vitis Kernels (.xo) for hardware deployment. This streamlines the workflow from software model to fully implemented design.

GNNBuilder supports implementing Vitis kernels on platforms like Alveo U50 and Alveo U280, including full bitstream generation (.xclbin) and a host code testbench for on-chip graph dataset evaluation. This testbench, similar to the C++ testbench, uses Xilinx’s runtime library, XRT, and OpenCL for FPGA interfacing from the host CPU.

We implemented our models on the Xilinx Alveo U280 FPGA accelerator at 300 MHz using Vitis HLS [36] and Vitis [39] tools from Xilinx. Our framework directly provides the generated HLS code to the synthesis tools, accompanied by suitable build scripts.

To assess the effectiveness of runtime modeling in DSE, we examine direct-fit models for latency and BRAM prediction, comparing them to Vitis HLS-reported post-synthesis values. We focus on BRAM usage for resource modeling, as it is the primary constraint-violating resource.

The direct-fit latency and BRAM models are random forest regressors, fitted on datasets of model configurations and their post-synthesis values. Empirical testing showed that random forests outperformed linear/polynomial models, support vector machines, and gradient boosting tree models in avoiding overfitting.

These direct-fit models necessitate a pre-synthesized design database. As the number of possible configurations is too large for brute-force exploration, sparsely sampling the design space enables fitted models to interpolate between sampled designs, providing accurate estimates for unseen configurations.

Instead of requiring users to run HLS builds for each design configuration or create datasets, we provide serialized trained versions of the direct-fit models described earlier making it feasible to performance brute force or random sampling of the model configuration space. Evaluating these direct-fit models takes milliseconds, compared to minutes for HLS synthesis. This can reduce performance prediction runtime enables users to develop intelligent co-design tools for real-time optimization, paving the way for train-time model sparsity, quantization, and neural architecture search, among other possibilities.

The accuracy of the runtime and BRAM models is assessed against a database of 400 synthesized designs, randomly sampled from a configuration space of model parameters (see Listing 2). For the fitted models, a random forest regressor with 10 estimators is used. The models are evaluated using the mean absolute percent error (MAPE) between the true post-synthesis metrics and predicted metrics. To examine overfitting, a 5-fold cross-validation (CV) is conducted, averaging the test MAPE for each fold to obtain the final cross-validation MAPE.

This section evaluates various model architecture configurations across multiple datasets, comparing the proposed hardware implementations:

• •

  PyG-CPU: A PyTorch Geometric CPU model

• •

  PyG-GPU: A PyTorch Geometric GPU model

• •

  CPP-CPU: A C++ floating-point CPU model

• •

  FPGA-Base: Proposed hardware model without hardware parallelism

• •

  FPGA-Parallel: Proposed hardware model with hardware parallelism

Performance is analyzed using a fixed GNN model with varying GNNConv layers (GCN, GraphSAGE, GIN, and PNA), across graph-level task datasets such as QM9, ESOL, FreeSolv, Lipophilicity, and HIV from MoleculeNet [1].

The CPU models are evaluated on an Intel Xeon Gold 6226R, while the GPU models are assessed on an NVIDIA RTX A6000. The hardware models (FPGA-Base and FPGA-Parallel) are implemented as described in Section VII-A.

Each baseline is evaluated on a batch size of 1, with the runtimes for CPU and GPU implementations computed by averaging the runtime of the first 1000 graphs of each dataset (or the complete dataset if it contains fewer than 1000 graphs). FPGA implementations’ runtimes are obtained from the worst-case estimate provided by Vitis HLS after synthesis. The architecture configuration in Listing 3 is used for all models.

The FPGA-Parallel implementations employ different parallelism factors for GCN, SAGE, and GIN models (gnn_p_in=1, gnn_p_hidden=16, gnn_p_out=8, p_in=8, p_hidden=8, p_out=1), while PNA models use gnn_p_hidden=8 and gnn_p_out=8. These models utilize <16, 10> bit fixed-point data representations. FPGA-Base implementations have parallel factors set to $1$ and implement node features using <32, 16> bit fixed-point types.

The results of fitting the latency and BRAM models on our database of generated designs are illustrated in Figure 4. The direct-fit latency model achieved a CV MAPE of approximately 36%, while the direct-fit BRAM model obtained a CV MAPE of approximately 17%. Figure 4 demonstrates that the direct-fit model consistently predicts the true value with few outliers. These findings indicate that directly fitting models on a design database, which sparsely samples the design configuration space, is an effective and straightforward approach for performance modeling in GNNBuilder, enabling rapid Design Space Exploration (DSE).

To exemplify the speed up of direct fit models over standard evaluation of the HLS tool, we also analyze the performance estimate compute time for all 400 model configurations used to train the direct fit models. We present the results in Figure 5, which can be viewed as a timeline of runs. All model calls for the direct fit models to finish in under a second, while all Vitis HLS synthesis runs finish in under two days. An average direct fit model call takes $1.7$ ms, while an average Vitis HLS synthesis run takes $9.4$ minutes. This difference is around 6 orders of magnitude emphasizing the the real-time performance estimation of direct fit models.

The performance results for the proposed accelerator hardware framework in comparison to other implementations are shown in Figure 6, Figure 7, and Table IV. The values in Table IV indicate the speedup factors of teh FPGA-Parallel implementation for the latency values averaged across datasets. For all cases, there is at least a 6x speedup in the parallelized FPGA implementation over the PyG CPU, PyG GPU, and C++ CPU implementations. Across all models, there is a geometric mean speedup of 6.33$\times$ over PyG-CPU and 6.87$\times$ over PyG-GPU. The resource usage also shows more room for BRAM and DSP utilization across models indicating there is more room for increased parallelism and higher speedups.