Molecular property prediction is a critical task in computer-aided drug discovery workflows. In this context, representing molecules as graph structures has become a common approach, leading to the development of graph-based neural networks (GNNs). A prominent type of GNN are the graph convolution network (GCN) that incorporates a number of graph convolution layers followed by fully connected layers in the form of multi layer perceptrons (MLP). Addressing the need for enhanced efficiency, this paper introduces an innovative approach that maps different computation types to specialized accelerators in a heterogeneous edge platform. The proposed system integrates an FPGA accelerator for the sparse operations present in the GCN layer and an Edge TPU for ordinary dense operations in the fully connected layer. The FPGA accelerator is described using High-Level Synthesis (HLS) and optimized to handle the sparsity found on the graph structure and node features while the Edge TPU deploys a systolic array of MAC units optimized for dense matrix multiplication. This hardware/software co-designed GCN+MLP architecture utilizes Pynq and TensorFlow Lite Runtime functions, executed on a multi-core ARM CPU available in an AMD/Xilinx Zynq ultrascale+ device, in conjunction with the Edge TPU and programmable logic. Our results demonstrate performance improvements, achieving speedups of up to 47x compared to traditional software implementations.
Funding Agencies|Wallenberg AI autonomous systems and software (WASP) program - Knut and Alice Wallenberg Foundation