作者：周楊傑、沈雯婷

開篇

近日，阿里雲機器學習平臺PAI和上海交通大學冷靜文老師團隊合作的論文《圖神經網路統一圖運算元抽象uGrapher》被ASPLOS 2023錄取。

為了解決當前圖神經網路中框架中不同的圖運算元在不同圖資料上靜態kernel的效能問題，uGrapher通過將所有圖運算元抽象為統一的中間表達形式，解耦圖運算元的計算和排程，並定義了在GPU上優化圖運算元的設計空間，以針動態變化的圖運算元和圖資料自適應的生成並行執行策略，為圖神經網路中的圖運算元提供高效能的計算支援。對比DGL [1], PyG[2], GNNAdvisor[3]，uGrapher平均可以取得3.5倍的效能提升。

背景

近年來，圖神經網路(Graph Neural Networks, GNNs)因其在非歐幾里德空間中對圖結構進行學習和推斷的強大能力而受到學術界和產業界的廣泛關注。GNNs將基於DNN的特徵變換和基於圖的操作結合起來，沿著圖結構傳播和聚合資訊。現有的GNN框架如DGL和PyTorch-Geometric（PyG）擴充套件了DNN 框架（如TensorFlow 和PyTorch），並引入了“訊息”這一概念，它是與每個邊相關聯的特徵向量的中間值。對於圖\(G=(V,E)\)上的任何操作，可以根據資料的屬性和資料移動的方向分為三個階段，即訊息建立、訊息聚合和特徵更新，公式化如下圖:

其中，\(u\)和\(v\)是頂點索引，\(e\)是\(u\)和\(v\)之間的邊索引；\(h_v\)是頂點\(v\)的特徵向量，\(m_e\)是邊\(e\)上的訊息。

uGrapher將需要遍歷輸入圖結構的操作符定義為圖運算元。圖運算元包括“訊息建立”、“訊息聚合”和“融合聚合”三類。其中，“融合聚合”是指當“訊息建立”是一個簡單的複製操作時，它可以與“訊息聚合”融合在一起，以避免冗餘訪問，DGL和PyG採用了這種融合優化。

以GAT模型為例，它包含幾個具有不同計算模式的圖操作符。第一個“訊息建立”操作非常輕量級，它將每個邊的源頂點和目標頂點的特徵相加作為訊息以計算注意力權重；第二個“融合聚合”操作首先從源頂點複製特徵，然後逐邊乘注意力權重，最後將變換後的邊上的訊息聚合為頂點的新特徵。第二個操作比第一個操作更加計算密集。

由於圖結構引起的不規則記憶體行為，再加上這些圖運算元中複雜的算術計算，圖神經網路中圖運算元的高效能運算成為重要挑戰。

現有的圖神經網路框架依靠手寫靜態kernel來實現圖運算元的計算。然而，隨著圖神經網路演算法演進，圖運算元的可變性和複雜性不斷增加，其計算也變得更加複雜；同時，不同分佈的圖資料作為輸入也給圖神經網路的計算帶來了特有的複雜性，這導致靜態運算元難以維持較好的效能。因此，本文探索瞭如何在變化的圖資料和圖模型上進行圖運算元的計算優化。

挑戰

（1）圖神經網路引入了圖運算元的複雜性和圖資料的變化性兩大特徵，導致了圖運算元計算優化上的難題。

下表根據輸入和輸出資料型別將DGL支援的160個圖操作符進行了分類。即使具有相同的輸入或者輸出資料型別，圖運算元也可以執行不同的計算模式。圖運算元的複雜性導致很難找到靜態的方式為所有圖運算元的計算操作提供高效能支援。

真實世界中的圖資料集有很大的差異。圖的規模，即頂點和邊的數量，圖的平衡程度，即鄰接矩陣行的非零值的標準差，以及特徵和類大小，這些屬性在不同的圖之間有顯著的差異。而這些差異影響了圖運算元的記憶體使用和計算複雜性。

（2）由於缺乏系統優化方法，現有GNN框架使用的底層CUDA kernel存在低效和缺乏靈活性的問題。

DGL 在支援上文中的訊息傳遞程式設計介面時呼叫靜態CUDA kernel，這些靜態kernel不能適應變化的計算場景。比如，在執行不平衡圖時，GPU 的低佔用率導致了硬體資源的浪費。當執行小圖時，GPU 效能通常受到並行性的限制，而執行大圖時，由於低區域性性，訪問頻寬成為瓶頸。同時，這些指標在不同圖運算元之間也會有所差異。

破局

uGrapher使用巢狀迴圈作為圖運算元的排程表達，並允許使用者定製輸入張量和不同階段的函式操作來表示不同的圖運算元運算。

下圖為面向圖神經網路中的圖運算元統一的抽象細節。

edge_op實現了邊上的訪存和計算的函式表示，而gather_op實現了邊到頂點的合併函式表示。另外還有三個輸入張量，可以為源頂點嵌入張量（Src_V），目的地頂點嵌入張量（Dst_V），邊嵌入張量（Edge），以及NULL的任意一種。 張量的資料型別還確定了迴圈計算中不同的定址模式（第10 至12 行)。

下面的公式正式定義了uGrapher的統一抽象，其中, \(\psi\)是edge_op 函式， \(\rho\)是gather_op 函式。該抽象捕捉了圖運算元的完整語義，包括其計算和記憶體移動模式。

根據圖運算元統一的抽象，uGrapher構建了運算元優化的設計空間，以實現高效能的圖運算元執行。

uGrapher使用區域性性、並行性和工作效率來描述GPU上圖運算元效能的指標。對巢狀迴圈應用tiling 或者blocking 技術可以提高圖運算元的區域性性；通過啟動更多的執行緒、warp 或執行緒塊可以提高圖運算元的並行性；工作效率用額外開銷的倒數表示，同一運算子的不同執行策略可能會引入額外的計算，例如地址計算等，共享頂點的邊平行計算可能需要原子指令。

現有圖處理系統中有兩種經典並行化策略：執行緒-頂點和執行緒-邊並行。前者降低了並行性，但提高了輸出資料的重用性和區域性性。後者降低了工作效率，因為可能需要原子更新操作。

由於GNN 中的頂點/邊特徵是向量，GNN增加了特徵維度的並行化策略，為warp-頂點和warp-邊，相對執行緒-頂點/邊策略，可以啟動更多的warp，從而增加並行性。然而，由於每個warp 的快取容量減少，這個策略也會損害區域性性。

因此，沒有一種單一的策略可以同時提高這三個指標，uGrapher通過上述統一的IR表達，設計了統一的高效能運算介面，來探索優化空間，進行效能的權衡。整體架構如下圖所示。

uGrapher提供的圖運算元統一高效能運算介面設計如下圖所示。

uGrapher 介面包含三個引數：graph_tensor，表示圖資料；op_info，用於傳遞edge_op、gather_op 和輸入張量的計算資訊；parallel_info，用於指定並行化策略。

uGrapher 的介面設計將運算元計算、圖資料和並行化策略分離，使得使用者可以通過手動，或者針對不同的運算元和圖結構提出自己的啟發式演算法來選擇執行策略。同時，當用戶未指定任何並行化策略時，uGrapher會利用LightGBM[4]訓練決策模型，選擇並行化空間中的最優策略來自動調整到最佳並行化策略，以在不同的GPU架構和圖資料集上為所有圖神經網路中的圖運算元提供專門和最優的計算排程。uGrapher實現了每種並行化策略的CUDA 核心模板，併為每種圖運算元保留裝置函式介面，並實現了端到端的程式碼生成，包含了運算元合併和裝置函式生成，以支援靈活和高效的算在實現。更多的細節請閱讀我們ASPLOS 2023的論文。

目前，阿里雲正在將uGrapher的關鍵設計整合進PAI自研的大規模圖神經網路框架GraphLearn中，從而為工業級別的圖神經網路應用帶來效能加速。

PAI長期招聘實習生，如果有對分散式深度學習訓練框架、分散式圖神經網路訓練框架、計算通訊優化等感興趣的同學，歡迎投遞簡歷到郵箱[email protected]或[email protected]

第五板塊：

• 論文標題：

uGrapher: High-Performance Graph Operator Computation via Unified Abstraction for Graph Neural Networks

• 論文作者：

周楊傑，冷靜文，宋曜旭，盧淑文，王勉, 李超，過敏意, 沈雯婷，李永，林偉等

• 論文pdf連結：

http://dl.acm.org/doi/10.1145/3575693.3575723

• 參考文獻：

[1] M. Wang, D. Zheng, Z. Ye, Q. Gan, M. Li, X. Song, J. Zhou, C. Ma, L. Yu, Y. Gai et al., “Deep graph library: A graph-centric, highly-performant package for graph neural networks,” arXiv preprint arXiv:1909.01315, 2019.

[2] M. Fey and J. E. Lenssen, “Fast graph representation learning with pytorch geometric,” arXiv preprint arXiv:1903.02428, 2019.

[3] Y. Wang, B. Feng, G. Li, S. Li, L. Deng, Y. Xie, and Y. Ding, “GNNAdvisor: An adaptive and efficient runtime system for GNN acceleration on GPUs,” in 15th USENIX Symposium on Operating Systems Design and Implementation (OSDI 21), 2021, pp. 515–531.

[4] Guolin Ke, Qi Meng, Thomas Finley, Taifeng Wang, Wei Chen, Weidong Ma, Qiwei Ye, and Tie-Yan Liu. 2017. Lightgbm: A highly efficient gradient boosting decision tree. Advances in neural information processing systems 30 (2017).