图神经网络设计中的算子融合策略

©作者 |刘曜齐单位|北京邮电大学硕士生来源|北邮GAMMA Lab本文主要的描述基于消息传递机制的图神经网络设计中应用的算子融合策略，带领读者了解有关算子融合的相关问题以及方法。引言目前，图神经网络（GNN）的程序编写主要依赖 GNN 框架，例如 PyG，DGL 等，并从中享受到高效的设计。现有的 GNN 设计大多基于消息传递机制，包括三大步骤：消息创建，消息聚合，特征更新。假定图...

PaperWeekly

951人浏览 · 2023-01-07 22:40:21

PaperWeekly · 2023-01-07 22:40:21 发布

©作者 | 刘曜齐

单位 | 北京邮电大学硕士生

来源 | 北邮GAMMA Lab

本文主要的描述基于消息传递机制的图神经网络设计中应用的算子融合策略，带领读者了解有关算子融合的相关问题以及方法。

引言

目前，图神经网络（GNN）的程序编写主要依赖 GNN 框架，例如 PyG，DGL 等，并从中享受到高效的设计。现有的 GNN 设计大多基于消息传递机制，包括三大步骤：消息创建，消息聚合，特征更新。

假定图的节点集为，边集为，每一个节点有一个特征表示，边有一个特征表示，其中和是特征的维度。消息传递的过程，便是利用节点和边上的原始特征创建消息，并扩展成为边上的特征；紧接着将边上的特征聚合成为目标节点的特征；最后利用聚合结果以及目标节点的原始特征，经过更新后得到最后的结果。

如图 1 所示，在消息创建的过程中，我们将节点上的特征扩展到了边上，从而特征的大小由变为，而一般的图上的边的数量要远大于节点的数量，所以这样会造成大量的时间和内存开销。因此，这些 GNN 框架会利用算子融合等方式，越过消息创建过程，从而减少创建的冗余信息，以达到 GNN 设计时的优化。

▲ 图1 消息创建过程中冗余信息产生示意图

方法

算子融合主要通过对计算图上存在数据依赖的“生产者-消费者”算子进行融合，即仿照生产者消费者关系，如果生产者算子的结果可以直接输送到消费者算子而不需要中间结果表示，则可以将两个算子融合。

目前主流的 GNN 框架有 DGL（Deep Graph Library）和 PyG（PyTorch-Geometric），它们都提出了各自的算子融合策略。一些研究人员在深入思考 GNN 的特性之后，提出了各自的优化框架，例如 Seastar 和 Graphiler。因此本章节我们会介绍 GNN 框架 DGL，PyG，Seastar，Graphiler 中采用的算子融合策略。

2.1 DGL

DGL [1] 不仅支持用户自定义函数（UDFs）来完成消息创建和消息聚合的过程，而且提供了丰富的消息创建和消息聚合的原语（primitive）。DGL 推荐使用这些原语因为它们利用的是高性能稀疏矩阵内核，相比于 UDF，可以获得更好的性能。如果用户在消息创建和消息聚合的过程中，都使用 DGL 提供的内置原语，那么 DGL 就会融合消息创建和消息聚合两个阶段，从而越过消息创建的过程，因此可以节约上述的内存开销。[2]

例如，DGL 中的消息创建原语 u_mul_e 和消息聚合原语 max，代表的含义分别为节点上的特征和边上的特征相乘，和将目标节点的邻居特征取最大的方式聚合，等价于以下的写法：

def udf_u_mul_e(edges):
    return {'m' : edges.src['h'] * edges.data['w']}
def udf_max(nodes):
    return {'h_max' : th.max(nodes.mailbox['m'], 1)[0]}

在 udf_u_mul_e 中，我们将 edges 的源节点的特征 h 放到了边上，因此会造成额外的内存开销，并且在特征放置的过程中也会有额外的时间。对于 udf_max 也是类似的道理。所以，必须在消息创建和消息聚合的过程中都使用 DGL 提供的内置原语，才能够取得底层上的优化。

2.2 PyG

PyG [3] 同样可以完成消息创建和消息聚合的内核融合，使用的方法为 message_and_aggregate 函数。这个函数只有在它被实现并且传播是基于 torch_sparse.SparseTensor 的情况下才会被调用。如果可行，那么时间开销和内存开销都可以有所降低。[4]

例如，PyG 中对于 GraphSAGE 的处理 [5]，其中包含了两种方式，其中一种为普通的消息创建和消息聚合，即定义消息函数：

def message(self, x_j: Tensor) -> Tensor:
    return x_j

创建消息后，使用内置 aggregate 函数完成消息的聚合，在这个过程中，用户将节点上的消息映射到边上，会增加时间成本以及内存存储的开销。而另外一种方法，就是在我们传入基于 torch_sparse.SparseTensor 的情况下，会直接调用以下的 message_and_aggregate 函数：

def message_and_aggregate(self, adj_t: SparseTensor,
                            x: OptPairTensor) -> Tensor:
adj_t = adj_t.set_value(None, layout=None)
return matmul(adj_t, x[0], reduce=self.aggr)

在这个函数中，直接使用 SparseTensor 的矩阵相乘，直接完成消息传递，避免了边上的消息创建的过程，不仅可以降低时间开销，而且可以减少内存消耗。

2.3 Seastar

Seastar [6]，全称 Source-Edge-Aggregation star，是一个采用以节点为中心（vertex-centric）编程的 GNN 框架，由于计算是以节点为中心的，边围绕着中心节点，呈现为星形，故命名 Seastar。其针对 GNN 的算法设计了一套全新的编程接口。

▲ 图2 Seastar实现GCN和GAT的代码示例

Seastar 根据每一个数据的特征，按照指定的规则将其标记为四种类型，分别为 S（source），D（destination），E（edge），或 P（parameter）。紧接着，Seastar 对算子进行类型标注，如果一个算子的输出是 S、D 或 E 类型，那么它就被定义为 S、D 或 E 类型，此外，还定义了 A 类型作为聚合类型运算符的特殊集合。

Seastar 描述了三种融合的机会。S-E，E-E 和 E-A。S-E 融合将 Source 阶段的生产者和 Edge 阶段的消费者合并为一个运算符。E-E 融合更为直接，因为它们需要相同的索引集来访问特征（例如，GAT 前向传播中的 LeakyRelu 和 Exp 操作可以按照 E-E 融合来进行）。E-A 融合是通过首先在 Edge 阶段进行计算，然后使用目的节点的 id 来进行聚合（例如，GAT 前向传播中的 Exp 和 AggSum 运算可以融合）。

2.4 Graphiler

Graphiler [7] 是一个 GNN 的编译栈，它通过将 UDF 转换为一个新的抽象称为消息传递数据流图（Message Passing Data Flow Graph, MP-DFG），从而进行优化，包括内核融合等操作，利用DGL提供的原语库以及手动编写的一小部分内核，便可以将这些算子融合，从而减少中间变量的生成。

和 Seastar 类似，Graphiler 将数据分为五类，分别为节点数据，边数据，节点类型数据，边类型数据和共享数据，并将运算符标记为四种类型，分别为 Broadcast，Recude，Norm，Dense，随后根据以下规则识别出特定的运算操作

Broadcast 运算符通过索引查找将数据从依赖类型转换为另一依赖类型。
Reduce 运算符是 Broadcast 相反的
Norm 运算符在边数据上工作，将两个步骤合二为一，首先聚合属于同一目标节点的边数据，然后将聚合的结果广播回给边，例如，softmax 操作。
Dense 运算符操作符独立于图结构，不产生数据移动。

Graphiler 的算子融合优化可以分为以下两种：

Broadcast-compute：如图3，HGT 代码的第 5 行由 Broadcast 运算符 edge.type 产生的中间边的数据被接下来的 batch_mm 消耗掉。假定的大小为，如果将两个运算内核融合，那么融合的内核可以直接从中读取数据，从而减小中间大小的内存消耗以及相应的内存访问时间。
- 原始模式：或或，其中是 Dense 运算符，和是 Broadcast 运算符。
- 替换模式：。

▲ 图3 HGT的UDF实现

Broadcast-compute-reduce：这个模式常见于许多基于消息传递的 GNNs。以 GAT 为例，首先将源节点的特征创建消息存储到边上，通过边上的注意力值进行缩放，并将它们聚合到新的节点特征中。专家级开发者通常用框架提供的原语来重写它，将这三个步骤融合为一个步骤，以避免实例化大小为的中间数据。
- 原始模式：或或，其中和分别是 Reduce 和 Dense运算符，和是 Broadcast 运算符。
- 替换模式：。

Graphiler部分实验结果展示

我们利用其他实验环境复现 Graphiler 的实验结果，以展现算子融合在同质图和异质图上带来的收益。

3.1 实验环境

CPU：Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
GPU：Tesla T4 (16GB version)
软件：torch v1.8.2+cu111，dgl-cu111 v0.6.1，pyg v2.0.1

3.2 同质图实验

注意: compiler 和 reorder 是 Graphiler 中其他的优化操作，详见Graphiler 原论文。

● 时间消耗（ms/infer）

● 内存消耗（MB）

3.3 异质图实验

● 时间消耗（ms/infer）

● 内存消耗（MB）

参考文献

[1] Wang M Y. Deep graph library: Towards efficient and scalable deep learning on graphs[C]//ICLR workshop on representation learning on graphs and manifolds. 2019.

[2] https://docs.dgl.ai/api/python/dgl.function.html

[3] Fey M, Lenssen J E. Fast graph representation learning with PyTorch Geometric[J]. arXiv preprint arXiv:1903.02428, 2019.

[4] https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html?highlight=message_and_aggregation#torch_geometric.nn.conv.MessagePassing.message_and_aggregate

[5] https://github.com/pyg-team/pytorch_geometric/blob/5b1051fca16937d97ecce05fa5fe272efde0fd2d/torch_geometric/nn/conv/sage_conv.py#L143

[6] Wu Y, Ma K, Cai Z, et al. Seastar: vertex-centric programming for graph neural networks[C]//Proceedings of the Sixteenth European Conference on Computer Systems. 2021: 359-375.

[7] Xie Z, Wang M, Ye Z, et al. Graphiler: Optimizing Graph Neural Networks with Message Passing Data Flow Graph[J]. Proceedings of Machine Learning and Systems, 2022, 4: 515-528.

更多阅读