[论文笔记] GNNAdvisor: An Adaptive and Efficient Runtime System for GNN Acceleration on GPUs

GNNAdvisor: An Adaptive and Efficient Runtime System for GNN Acceleration on GPUs

GNNAdvisor: GPU 上 GNN 加速的自适应高效运行时系统 [Paper] [Slides] [Code]
OSDI’21

摘要

提出了 GNNAdvisor, 一个用于加速 GPU 平台上各种 GNN 工作负载的自适应高效运行时系统.

• 从 GNN 模型和输入图中探索并识别了几个与性能相关的特征
• 实现了一种新颖且高效的 2D 工作负载管理
• 利用 GPU 内存层次结构进行加速
• 集成了一个轻量级分析模型, 用于有效的设计参数搜索

1 介绍

GNN 突出特点: 在聚合阶段的图操作(scatter-and-gather) 和在更新阶段的神经网络(NN)操作(矩阵乘法)交错执行.
最先进的 GNN 框架遵循一刀切(one-size-fits-all, 通用适配)的实现方案.

先前 GNN 系统的不足:

• 未能利用 GNN 的输入信息: 模型的层序列、聚合方法、结点嵌入维度, 以及输入图都具有多样性.
• 优化非 GNN 定制: 现有 GNN 框架未解决 GNN 和图处理之间的差异.
• 对输入适应性的运行时支持较差: GNN 的一些关键性能相关信息仅在运行时可用.

提出了 GNNAdvisor, 一个用于 GPU 上 GNN 加速的自适应高效运行时系统.
单 GPU 系统; 前端: Pytorch; 底层: C++/CUDA 构建, 使用 Pytorch Wrapper 与 Pytorch 集成.

• 引入了输入加载器和提取器(Loader&Extractor)来利用输入级信息来指导系统级优化
• 包含一个决策器(Decider), 其中包括用于自动运行时参数选择的分析模型和一个轻量级结点重编号例程.
• 集成了内核与运行时定制工具(Kernel&Runtime Crafter)来定制参数化 GNN 内核和 CUDA 运行时, 其中包括 2D 工作负载管理和一组 GNN 专用内存优化.

本文贡献:

• 第一个探索 GNN 输入属性, 并深入分析了其在指导基于 GPU 的 GNN 计算的系统优化方面的重要性.
• 提出了一组具有参数化的 GNN 定制的系统优化, 包括 2D 工作负载管理和 GPU 上的专用内存定制; 并结合分析建模和参数自动选择来简化设计空间探索.
• 实验证明 GNNAdvisor 相对于最先进的 GNN 执行框架的性能优势.

2 背景和相关工作

2.1 图神经网络

一次迭代中 GNN 的计算流程:

基于第 $k$ 层($k\le 0$)的嵌入信息计算第 $k+1$ 层结点 $v$ 特征向量的公式:
$$\begin{array}{rl}& a_v^{(k+1)}={\text{Aggregate}}^{(k+1)}(h_u^{(k)}|u\in \text{N}(v)\cup h_v^{(k)})\\ & h_v^{(k+1)}={\text{Update}}^{(k+1)}(a_v^{(k+1)})\end{array}$$

• $h_v^{(k)}$: 结点 $v$ 在第 $k$ 层的嵌入向量
• $a_v^{(k+1)}$: 结点 $v$ 收集邻居信息的聚合结果
• 不同的 GNN 的聚合方法以及聚合和更新的顺序可能有所不同, 更新函数通常由标准的 NN 操作组成.

2.2 图处理系统

扩展图处理系统支持 GNN 计算的挑战:

1. 图处理中的常见算法优化可能不会使 GNN 受益. 如依赖于结点边界(激活邻居)的优化, 而 GNN 始终是全部结点.
2. 图处理的系统优化技术只有经过精细的调整和校准后才会使 GNN 受益. GNN 的数据集并行性还包括嵌入维度.
3. 图系统中缺少 GNN 计算的一些基本函数. 基于神经网络计算的前向值传播和反向梯度传播的结点更新在图系统中不可用.

2.3 深度学习框架

缺乏对 GNN 中的非欧几里德数据(如图)的支持. 扩展 NN 框架以支持作为输入的 GNN 的挑战:

1. 缺乏有效的后端支持来实现高性能.
2. 主要计算内核是硬编码的, 没有设计灵活性.

3 GNN 应用的输入分析

3.1 GNN 模型信息

GNN 两类主流聚合方法:

1. 仅涉及邻结点嵌入的聚合(如 $sum$, $min$). 如在图卷积网络(Graph Convolutional Network, GCN)中.
   常见设计: 在聚合(从邻结点嵌入收集信息)之前减少更新阶段(即结点嵌入矩阵与权重矩阵相乘)的结点嵌入维数.
   优化策略: 提高内存局部性, 因为更多的结点嵌入可以缓存在快速内存(如 GPU 的 L1 缓存)中.
2. 对每个邻结点应用特殊边特征(如权重, 通过组合源结点和目标结点计算的边向量)的聚合. 如在图同构网络(Graph Isomorphism Network, GIN)中.
   优化策略: 改进计算并行化, 因为工作负载可以在更多并发线程之间共享以提高整体吞吐量.

3.2 图信息

结点度数和嵌入维数:
工作负载大小如果由邻结点数量(如结点度数较大)主导, 可以定制能同时处理更多邻居的设计, 以增加邻居间的计算并行性; 如果由结点嵌入大小(如高维结点嵌入)主导, 可以考虑沿结点嵌入维度提高计算并行性.

图社区(Graph Community):
描述了一小组结点倾向于保持"强"组内连接(许多边), 同时与图的其余部分保持"弱"连接(更少的边).
在结点嵌入大小很大的 GNN 计算中, 以结点为中心的加载会触发大量不必要的内存访问, 因为重复邻居加载的开销占主导地位而不会被每个结点的并行性所抵消.

在 GPU 上利用图社区的关键: 通过利用 L1 缓存来有效地利用线程之间的数据局部性.
首先捕获图社区, 然后将此类局部性从输入级别(结点 ID 相邻)映射到底层 GPU 内核(线程/wrap/线程块 ID 相邻).

4 2D 工作负载管理

GNN 工作负载在两个主要维度上增长: 邻居数量和嵌入维度大小.

4.1 粗粒度邻居划分

旨在解决: 结点间工作负载不平衡和冗余原子操作.

将结点的邻居分解为一组大小相等的邻居分组, 并将每个邻居分组(neighbor group, NG)的聚合工作负载作为调度的基本工作负载单位.
两个组件:

• 邻居划分模块: 构建在图加载器之上的轻量级模块, 将图 CSR 划分为相等大小的分组. 每个邻居分组只涵盖一个目标结点的邻居.
• 邻居划分图存储: 维护每个邻居分组基于元组的元数据, 包括其 ID, 在 CSR 表示中邻居的起止位置, 以及源结点.

优点:

1. 与更粗粒度的基于结点/顶点为中心的划分相比: 减轻工作负载单元的大小不规则性
2. 与更细粒度的以边为中心的划分相比: 避免管理许多微小工作负载单元所带来的开销
3. 引入性能相关的参数: 邻居分组大小(neighbor group size, $ngs$).

4.2 细粒度维度划分

来沿嵌入维度进一步分配邻居分组的工作负载, 以提高聚合性能.
每个线程独立管理一维的聚合; 维度大小大于工作线程数则需多次迭代.

依据:

• 可以适应更多样化的嵌入尺寸大小
• 引入了另一个与性能相关的参数: 工作线程数(dimension-worker, $dw$)

4.3 基于 warp 的线程对齐

每个 warp 独立管理来自一个邻居分组的聚合工作负载.

优点:

• 可以最小化线程间同步
• 减少了单个 warp 的工作负载, 且不同 warp 的工作负载更均匀
• (来自同一 warp 的连续)内存访问可以合并

5 专用内存优化

5.1 社区感知的结点重编号

通过通过轻量级结点重编号重新排列结点 ID, 以改进 GNN 聚合期间的时间/空间局部性.

何时应用:
对于形状更不规则的图, 而非邻接矩阵已近似块对角模式的图, 重排序可以带来显著的性能改进.

确定图重排序是否有益的度量标准: 平均边跨度(Average Edge Span, AES):
$$\text{AES}=\frac{1}{\#E}\sum _{(sr{c}_{id},tr{g}_{id})\in E}|sr{c}_{id}-tr{g}_{id}|$$
当 $\sqrt{AES}>\lfloor \frac{\sqrt{\#N}}{100}\rfloor$ 时结点重编号更可能提高运行时性能.

如何应用:
利用 Rabbit Reordering, 一种完全并行且低开销的图重排序技术.

5.2 warp 感知的内存定制

以 warp 为中心的共享内存优化技术. 根据块级 warp 的组织模式定制共享内存布局, 以显着减少原子操作和全局内存访问的数量.

1. 为每个邻居分组(warp)的目标结点保留一个共享内存空间(对于浮点嵌入为 $4\times Dim$ 字节), 以便在共享内存中缓存归约的中间结果.
2. 只指定一个 warp(称为 leader) 来将每个目标结点的中间结果复制到全局内存.

内存定制过程:

(注: 算法 11 行和 14 行的 $local\_cnt$ 笔者认为更可能是 $local\_cnt\times Dim$, 因为每个 warp 计算的中间结果大小为嵌入维度 $Dim$.)
每个 warp (在 $warpPtr$ 中维护) 具有三个属性: $nodeSharedAddr$ (邻居分组聚合结果的共享内存地址)、$nodeID$ (目标结点ID）、和 $leader$ (布尔标志, 指示当前 warp 是否为leader warp).

6 设计优化

分析建模:
每线程工作负载(workload per thread, WPT)和每线程块共享内存使用率(shared memory usage per block, SMEM):
$$\text{WPT}=ngs\times \frac{Dim}{dw},\text{SMEM}=\frac{tpb}{tpw}\times Dim\times FloatS$$

• $ngs$: 邻居分组大小
• $dw$: 维度 worker 数
• $Dim$: 结点嵌入维度数
• $FloatS$: 浮点数大小, 4 字节
• $tpb$: 每线程块线程数(thread-per-block), 用户选择
• $tpw$: 每 warp 线程数(thread-per-warp), GPU 决定

参数自动选择:

1. 基于 $tpw$ (硬件限制) 和 $Dim$ (输入属性) 确定 $dw$:
   $$dw=\{\begin{array}{l}tpwDim\le tpw\\ \frac{tpw}{2}Dim<tpw\end{array}$$
2. 基于选择的 $dw$ 和用户指定的 $tpb$ 确定 $ngs$ (根据上一小节公式).
   限制: $WPT\approx 1024$, $SMEM\le SMEMperBlock$ ($SMEMperBlock$: 每个线程块的共享内存大小, 由硬件决定).
   经验选择: $tpb$ 为 2 的幂次且小于等于 1024, $32\le tpb\le 128$ 可提升 SM warp 调度灵活性避免长尾效应.

7 评估

性能: Figure 8 ~ 11, Table 2

各种优化的性能提升: Figure 12
额外研究: Figure 13, Figure 14

---

笔者总结

本文的核心在于通过探索 GNN 输入信息, 提出了包括 2D 工作负载管理和 GPU 上的专用内存定制的一组 GNN 定制的系统优化, ; 并结合了分析建模和参数自动选择.