作者：周杨杰、沈雯婷
开篇
近日，阿里云机器学习 PAI 平台对于图神经网络对立高性能 IR 的论文《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 中，从而为工业级别的图神经网络应用带来性能减速。
第五板块：
- 论文题目：
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).