中文題目：Leda：利用分塊數據流在高帶寬內存FPGA上加速圖神經網絡中的稀疏矩陣-稠密矩陣乘

論文題目：Leda: Leveraging Tiling Dataflow to Accelerate SpMM on HBM-Equipped FPGAs for GNNs

錄用期刊/會議：43rd ACM/IEEE International Conference on Computer-Aided Design (ICCAD '24) (CCF-B類會議)

原文DOI：https://doi.org/10.1145/3676536.3676773

錄用/見刊時間：2024-10-27（錄用時間）

作者列表：

1）伊恩鑫 中國石油大學（北京）人工智能學院 計算機技術 碩21

2）白佳睿 中國石油大學（北京）人工智能學院 計算機科學與技術 本21

3）聶怡婕 中國石油大學（北京）人工智能學院 電子信息工程 本22

4）牛   丹 東南大學 自動化學院

5）金   洲 中國石油大學（北京）人工智能學院 計算機系教師

6）劉偉峰 中國石油大學（北京）人工智能學院 計算機系教師

摘要:

圖神經網絡（GNNs）在從圖結構中提取數據表征方面發揮著關鍵作用，推動了多個領域的發展。稀疏矩陣-稠密矩陣乘（SpMM）是GNN的核心運算之一。但是，由于圖矩陣的高度稀疏性和非零元的隨機分布，加速SpMM面臨著巨大挑戰。新興的高帶寬內存（HBM）的高并發能力為加速SpMM提供了新機遇。然而，由于負載不均衡和隨機內存訪問模式，在HBM FPGA上加速SpMM仍然困難重重。

在本文中，我們介紹了配備HBM的FPGA上的高性能SpMM加速器Leda。通過利用分塊稀疏格式來平衡負載并增強數據局部性。最小相似度重排算法顯著改善了寫后讀（RAW）依賴。此外，改進的外積數據流調度策略有效緩解了隨機內存訪問瓶頸。最后，我們提出了一種基于定制數據流的高度并行的硬件架構設計，并探索了輸入矩陣的重用性。實驗結果表明，Leda在幾何平均吞吐量和能效方面，分別比當前最先進的SpMM加速器Sextans和SDMA，以及K80 GPU高出1.27倍和1.36倍、1.85倍和2.00倍、1.95倍和5.23倍。

背景與動機:

目前，現有研究已在通用計算平臺（如CPU和GPU）上提出了SpMM的加速方案。然而，由于缺乏專用的內存結構和數據流設計，這些方案在訪存效率和計算資源占用率方面表現不佳。配備HBM的FPGA因其深度內存結構定制和高并發能力，已成為加速SpMM的極具吸引力的平臺。然而，充分利用HBM來加速SpMM仍然充滿挑戰，主要包括以下幾個方面：（1）現有的稀疏存儲格式和并行方法導致處理單元（PE）間的負載高度不均衡；（2）圖結構的冪律分布引入隨機內存訪問模式和較差的數據局部性；（3）RAW依賴導致的高延遲限制了計算資源占用率和吞吐量。

設計與實現:

一、分塊稀疏存儲格式

我們使用稀疏塊作為基本單位，確保PE間相對負載平衡；為了利用輸入稠密矩陣的重用性，我們用CSC格式來存儲稀疏塊；為了減少額外控制開銷，我們利用列主序COO格式存儲每個稀疏塊中的非零元信息；設置16位無符號掩碼數組，存儲非零元位置信息。

 

圖 1：分塊稀疏存儲格式 

二、最小相似度重排算法

我們提出了一種最小相似度重排算法。該算法利用掩碼數組計算稀疏塊中列間相似度，進而以列為單位進行非零元重排，同時結合插值操作，有效消除原始數據流中的RAW依賴。

 

圖 2：最小相似度重排算法

三、改進的外積數據流

我們以16×16固定分區大小的稀疏塊作為外積數據流的基本單元。由于中間積規模較小，片上緩沖區可完全容納，既避免了大量片外流量，又提高了輸入矩陣的重用性。

 

圖 3：改進的外積數據流

四、Leda的硬件架構設計

我們精心規劃了HBM通道的分配方案，以充分利用內存帶寬；專用計算核心陣列中的矩陣乘法單元（MMU）可以跳過稀疏矩陣中的空白結構，以減少冗余的片上內存寫入，靈活的重用隊列暫存可復用元素，以提高稠密矩陣重用性，設置多個PE組以提高計算并行性；累加器中的矩陣累加單元（MAU）平衡了傳輸與計算速度，并且支持快速隨機讀寫。

 

圖 4：Leda的硬件架構設計

實驗結果及分析:

一、Leda與FPGA上的SpMM加速器對比

Leda的幾何平均吞吐量分別比最新的SpMM加速器Sextans-hls和Sextans-tapa高出1.27倍和1.02倍。此外，幾何平均能源效率分別提高了1.36倍和1.10倍。

 

圖 5：Leda與Sextans-hls和Sextans-tapa的吞吐量與能效對比

二、Leda與K80 GPU對比

Leda與K80 GPU上的cuSPARSE和GE-SpMM相比，幾何平均吞吐量分別提高了1.95倍和1.58倍，幾何平均能效分別提高了5.23倍和4.23倍。

三、Leda與2060 GPU對比

Leda與2060 GPU上的cuSPARSE和GE-SpMM相比，幾何平均吞吐量分別提高了0.41倍和0.33倍（值得注意的是，2060 GPU的頻率遠高于Leda，幾乎是其8倍左右），幾何平均能效分別提高了1.49倍和1.18倍。

 

圖 6：Leda與cuSPARSE和GE-SpMM在K80 GPU和2060 GPU上的吞吐量與能效對比

四、預處理開銷分析

我們比較了Leda的預處理與CPU上單次SpMM的執行時間。在絕大多數矩陣上，預處理時間幾乎等同于單次SpMM的時間。只有少數大型圖矩陣的預處理時間略高于單次SpMM時間。不過，由于經過一次預處理后，在聚合階段可以進行多次SpMM迭代，因此Leda的預處理開銷可以忽略不計。

圖 7：12個圖矩陣上的預處理時間和單次SpMM執行時間的對比

結論:

本文中，我們在配備HBM的FPGA上提出了一種用于GNNs的高性能SpMM加速器Leda。采用分塊稀疏格式的定制數據流充分利用了HBM的優勢。最小相似度重排算法顯著改善了累加階段的RAW依賴。改進的外積數據流緩解了隨機訪存瓶頸。以MMU為中心的硬件架構設計進一步提高了并行性和數據重用性。實驗結果表明，與當前最先進的SpMM加速器和GPU實現相比，Leda在吞吐量和能效方面展現出了顯著的提升。

通訊作者簡介:

金洲，副教授，中國石油大學（北京）計算機系副教授，入選北京市科協青年人才托舉工程、校青年拔尖人才。主要從事集成電路設計自動化（EDA）、面向科學計算的DSA軟硬件協同設計等方面的研究工作。主持并參與國家自然科學基金青年項目、重點項目，科技部重點研發微納電子專項、高性能計算專項青年科學家項目，國家重點實驗室開放課題、企業橫向課題等。在DAC、TCAD、TODAES、SC、PPoPP、IPDPS、TCAS-II、ASP-DAC等重要國際會議和期刊上發表60余篇高水平學術論文。獲EDA2青年科技獎、SC23最佳論文獎、ISEDA23榮譽論文獎、IEEJ九州支部長獎等。

聯系方式：jinzhou@cup.edu.cn