GNN加速器顶层架构
此GNN加速器是为GraphSAGE算法设计的,但是它的设计也可以应用于其他类似的GNN算法加速。其顶层架构如下图所示。
该架构由以下模块组成:
图中的GNN内核是算法实现的核心部分(详情如下)。
RoCE-Lite是RDMA协议的轻量级版本,用于通过高速以太网进行远程存储访问,以支持海量节点的图计算。
400GE以太网控制器用于承载RoCE-Lite协议。
GDDR6存储器用于存储GNN处理过程中所需的高速访问数据(DDR4作为备用大容量存储器)。该存储器用于存储访问频率相对较低的数据,例如待预处理的图形数据。
PCIeGen5 ×16 接口提供高速主机接口,用于与服务器软件进行数据交互。
上述所有模块均通过具有高带宽的NoC实现互连。
GNN内核微架构
在开始讨论GNN内核的微架构之前,有必要先回顾一下GraphSAGE算法。其内层循环的聚合和合并(包括卷积)占据了该算法的大部分计算和存储访问。通过研究,我们得出这两个步骤的特点,具体如下。
表2:GNN算法中聚合和合并操作的对比(来源:https://arxiv.org/abs/1908.10834)
步骤 | 聚合操作 | 合并操作 |
存储访问方式 | 间接访问,不规则 | 直接访问,规则 |
数据重用 | 低 | 高 |
计算模式 | 动态,不规则 | 静态,规则 |
计算量 | 低 | 高 |
性能瓶颈 | 存储 | 计算 |
可以看出,聚合操作和合并操作在计算和存储访问模式上有着完全不同的需求。聚合操作涉及相邻节点的采样。然而,图形是一种非欧几里得数据类型——它的大小和维度是不确定且无序,矩阵稀疏,节点位置随机。因此,存储访问是不规则的,并且难以重复利用数据。
在合并操作中,输入数据是聚合结果(节点的低维表示)和权重矩阵。它的大小和维度是固定的,具有线性存储位置。因此对存储访问没有挑战,但是矩阵的计算量非常大。 基于上述分析,我们决定在GNN内核加速器设计中选择使用两种不同的硬件结构来分别处理聚合和合并操作(如下图示):
聚合器——通过单指令多数据(SIMD)处理器阵列,对图形相邻节点进行采样和聚合。单指令可以预定义为mean()平均值计算,或其他适用的聚合函数;多数据是指单次mean()均值计算中需要多个相邻节点的特征数据作为输入,这些数据来自子图采样器。SIMD处理器阵列通过调度器Agg Scheduler进行负载平衡。子图采样器通过NoC从GDDR6或DDR4读回的邻接矩阵和节点特征数据h0v分别缓存在邻接列表缓冲区(Adjacent ListBuffer)和节点特征缓冲区(Node Feature Buffer)。聚合的结果hkN(v)存储在聚合缓冲区(Aggregation Buffer)中。
合并器——通过脉动矩阵PE对聚合结果进行卷积运算。卷积核是Wk权重矩阵。卷积结果由ReLU激活函数进行非线性处理,同时也存储在PartialSum Buffer中,以用于下一轮迭代。
合并结果经过L2BN标准化处理后,即为最终的节点表示hkv。在一个典型的节点分类预测应用中,节点表示hkv可以通过一个全连接层(FC)来获取节点的分类标签。这个过程是传统的机器学习处理方法之一,在GraphSAGE文献资料中没有体现,这个功能也没有包含在这个架构中。