Aggregation engine for Graph Neural Networks

Graph Neural Networks (GNNs) are a class of deep learning methods intended to analyze graph data. GNNs include two different phases: the Aggregation phase, in which each node gathers information about its neighbors, and the Combination phase, which usually acts as a Neural Network on the output of the first phase. While the Combination possesses many of the same characteristics as other kinds of NNs with regard to the dataflow and can be optimized accordingly, the Aggregation phase presents some distinctive properties that prevent efficient mapping on traditional NN processors, and requires novel dedicated hardware and software schemes. In this work, an Aggregation Engine is designed based on a 2-D square mesh Network-on-Chip of SIMD cores. In order to have fast execution and efficient resource utilization, it is necessary to partition the input graph optimally among the different PEs at compile time. Such partitioning has been tested with different objectives, and the results have been compared for five distinct input graphs with different mesh sizes and other design parameters combinations. The best-suited partitioning objective proved to be weighted min-cut, and the most appropriate mesh size a 6-by-6 mesh. The experiments also highlighted the network congestion as the main factor limiting the scalability of the design, and degrading performance for high feature sizes. The 6-by-6 mesh has finally been implemented on an FPGA platform in order to report logic utilization and critical path. The resulting design achieves an average speedup of 74.38% over a sequential execution.