GNN加速器頂層架構(gòu)
此GNN加速器是為GraphSAGE算法設(shè)計的,但是它的設(shè)計也可以應(yīng)用于其他類似的GNN算法加速。其頂層架構(gòu)如下圖所示。
該架構(gòu)由以下模塊組成:
圖中的GNN內(nèi)核是算法實現(xiàn)的核心部分(詳情如下)。
RoCE-Lite是RDMA協(xié)議的輕量級版本,用于通過高速以太網(wǎng)進行遠程存儲訪問,以支持海量節(jié)點的圖計算。
400GE以太網(wǎng)控制器用于承載RoCE-Lite協(xié)議。
GDDR6存儲器用于存儲GNN處理過程中所需的高速訪問數(shù)據(jù)(DDR4作為備用大容量存儲器)。該存儲器用于存儲訪問頻率相對較低的數(shù)據(jù),例如待預(yù)處理的圖形數(shù)據(jù)。
PCIeGen5 ×16 接口提供高速主機接口,用于與服務(wù)器軟件進行數(shù)據(jù)交互。
上述所有模塊均通過具有高帶寬的NoC實現(xiàn)互連。
GNN內(nèi)核微架構(gòu)
在開始討論GNN內(nèi)核的微架構(gòu)之前,有必要先回顧一下GraphSAGE算法。其內(nèi)層循環(huán)的聚合和合并(包括卷積)占據(jù)了該算法的大部分計算和存儲訪問。通過研究,我們得出這兩個步驟的特點,具體如下。
表2:GNN算法中聚合和合并操作的對比(來源:https://arxiv.org/abs/1908.10834)
步驟 | 聚合操作 | 合并操作 |
存儲訪問方式 | 間接訪問,不規(guī)則 | 直接訪問,規(guī)則 |
數(shù)據(jù)重用 | 低 | 高 |
計算模式 | 動態(tài),不規(guī)則 | 靜態(tài),規(guī)則 |
計算量 | 低 | 高 |
性能瓶頸 | 存儲 | 計算 |
可以看出,聚合操作和合并操作在計算和存儲訪問模式上有著完全不同的需求。聚合操作涉及相鄰節(jié)點的采樣。然而,圖形是一種非歐幾里得數(shù)據(jù)類型——它的大小和維度是不確定且無序,矩陣稀疏,節(jié)點位置隨機。因此,存儲訪問是不規(guī)則的,并且難以重復(fù)利用數(shù)據(jù)。
在合并操作中,輸入數(shù)據(jù)是聚合結(jié)果(節(jié)點的低維表示)和權(quán)重矩陣。它的大小和維度是固定的,具有線性存儲位置。因此對存儲訪問沒有挑戰(zhàn),但是矩陣的計算量非常大。 基于上述分析,我們決定在GNN內(nèi)核加速器設(shè)計中選擇使用兩種不同的硬件結(jié)構(gòu)來分別處理聚合和合并操作(如下圖示):
聚合器——通過單指令多數(shù)據(jù)(SIMD)處理器陣列,對圖形相鄰節(jié)點進行采樣和聚合。單指令可以預(yù)定義為mean()平均值計算,或其他適用的聚合函數(shù);多數(shù)據(jù)是指單次mean()均值計算中需要多個相鄰節(jié)點的特征數(shù)據(jù)作為輸入,這些數(shù)據(jù)來自子圖采樣器。SIMD處理器陣列通過調(diào)度器Agg Scheduler進行負載平衡。子圖采樣器通過NoC從GDDR6或DDR4讀回的鄰接矩陣和節(jié)點特征數(shù)據(jù)h0v分別緩存在鄰接列表緩沖區(qū)(Adjacent ListBuffer)和節(jié)點特征緩沖區(qū)(Node Feature Buffer)。聚合的結(jié)果hkN(v)存儲在聚合緩沖區(qū)(Aggregation Buffer)中。
合并器——通過脈動矩陣PE對聚合結(jié)果進行卷積運算。卷積核是Wk權(quán)重矩陣。卷積結(jié)果由ReLU激活函數(shù)進行非線性處理,同時也存儲在PartialSum Buffer中,以用于下一輪迭代。
合并結(jié)果經(jīng)過L2BN標(biāo)準化處理后,即為最終的節(jié)點表示hkv。在一個典型的節(jié)點分類預(yù)測應(yīng)用中,節(jié)點表示hkv可以通過一個全連接層(FC)來獲取節(jié)點的分類標(biāo)簽。這個過程是傳統(tǒng)的機器學(xué)習(xí)處理方法之一,在GraphSAGE文獻資料中沒有體現(xiàn),這個功能也沒有包含在這個架構(gòu)中。