作者：周楊傑、沈雯婷

開篇

近日，阿里雲機器學習平台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鏈接：

https://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).