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文獻資料中沒有體現,這個功能也沒有包含在這個架構中。