[论文笔记] uGrapher: High-Performance Graph Operator Computation via Unified Abstraction for GNN

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

uGraper: 通过图神经网络的统一抽象实现高性能图算子计算 [Paper] [Presentation]
ASPLOS’23

摘要

提出了 uGrapher, 一个为不同图算子和数据集实现通用高性能的统一接口.

• 现有 GNN 框架易集成
• 将图算子的计算和调度解耦
• 构建了一个结合图张量和图循环语义的特定于 GNN 的算子抽象
• 探索基于抽象的各种调度策略, 权衡并行性、局部性与效率

1 介绍

GNN 模型:

• 具有巨大的体系结构空间, 使用的图算子的可变性和复杂性迅速增加.
• 在具有独特特征的不同图结构数据集上进行操作, 并对不同的图数据集和图算子表现出不同的模式和瓶颈, 从而缺乏并行的自适应性.
• 不同于传统图算法, 没有边界结点带有的复杂控制流, 而是在图遍历时涉及特征维度和更复杂的计算.

现有 GNN 框架:
依赖于手写实现; 只能在有限的 GNN 模型和数据集范围内实现最佳性能; 原因在于对不同的图算子和输入图使用了固定的执行策略.

提出了 uGrapher, 一个支持图算子的统一高性能接口, 可以轻松集成到现有 GNN 框架中.

• 将图算子的计算和调度解耦, 以适应不同 GNN 算子和数据集
• 将图算子抽象为嵌套稀疏-稠密 for 循环的统一形式
  • 嵌套循环最内层: 捕获不同图算子的语义
  • 嵌套循环最外层: 提供统一全面的并行空间的探索机会
• 基于统一抽象, 探索 GPU 上不同图算子对应的不同循环变换的各种执行策略及其权衡关系
• 为上层 GNN 框架提供统一易用的 API
• 灵活可扩展, 提供高性能的自动 CUDA 代码生成.

本文贡献:

• 针对不同的图算子和数据集分析了现有 GNN 框架在内核级别的低效性
• 针对 GNN 中的所有图算子提出了一个统一抽象, 其为 GPU 上不同并行执行策略定义了一个全面的优化空间.
• 基于统一抽象, 能够自动为所有图算子提供高性能的 CUDA 代码生成, 而只需简单的算子信息, 带来显著的灵活性和可扩展性.
• 设计了一个统一 API uGrapher, 支持现有框架中的所有图算子并探索其在不同数据集上的最佳并行执行策略.

2 背景和动机

2.1 图神经网络

GNN 模型的输出是输入图中每个结点的 $d$ 维嵌入向量.

GNN 模型:
图 $G=(V,E)$ 上的操作可分为三个阶段:

• $u$, $v$: 结点索引
• $e$: 结点 $u$ 和 $v$ 之间的边的索引
• $h_v$: 结点 $v$ 的特征嵌入
• $m_e$: 边 $e$ 关联的消息

图算子定义:
图算子: 需要遍历输入图结构的运算符.
图算子包括 $message-creation$(消息创建)、$message-aggregation$(消息聚合), 以及 $fused-aggregation$(融合聚合) 三类.

图算子的复杂性:

图数据的可变性:

2.2 GPU 上执行效率分析

由于缺乏系统的优化方法, 现​​有 GNN 框架使用的底层 CUDA 内核存在效率低下和不灵活的问题.
总结: 现有的 GNN 框架依赖于具有固定执行策略的手写内核, 由于图相关操作的多样性和现实世界图结构的多样性, 执行 GNN 模型的效率低下.

3 统一图算子抽象

3.1 图算子的嵌套循环表示法

3.2 统一抽象设计

图算子的三个执行阶段:

1. 将数据从结点移动到边
2. 对所有边执行边计算
3. 执行从边到相关联结点的归约函数

统一抽象:

两个由用户定义的动态算子:

• $edge\_op$: 每条边的边计算
• $gather\_op$: 边到结点的归约操作

三个嵌入张量 $A,B,C$ 的类型:

• 源结点嵌入张量($Src\_V$)
• 目标结点嵌入张量($Dst\_V$)
• 边嵌入张量($Edge$)
• $NULL$

统一抽象公式:
$$C_{c\_idx,f}=\sum _{dst\in V}\sum _{src\in V}\sum _{f\in F}\psi (C_{c\_idx,f},\rho (A_{a\_idx,f},B_{b_{i}dx,f}))$$

• $\psi$: $edge\_op$
• $\rho$: $gather\_op$

4 优化空间分析

4.1 权衡空间

图算子在 GPU 上性能的权衡空间, 即三个维度的优化空间:

• 局部性: 描述了程序中空间和时间重用量.
  作用: 可增加缓存命中率
  方法: 对嵌套循环进行平铺或分块
• 并行性: 指可以并发执行的计算量.
  作用: 可提高资源利用率, 隐藏内存访问延迟
  方法: 启动更多线程、warp、线程块
• 工作效率: 表示为开销值的倒数.
  同一算子的不同执行策略可能引入额外的计算; 存在写冲突时需要原子指令, 引入加锁开销.

没有一种策略可以同时提高这三个指标.

4.2 并行策略探索

两个细粒度参数:

• V/E-分组(V/E grouping): 将多条边或结点组合成一个组 (由 1 个线程执行).
• 特征平铺(feature tiling): 利用特征维度的并行性来启动更多线程 (特征映射到的 warp 数量).

uGrapher 支持的并行策略:

不同并行策略对局部性、并行性和工作效率的影响:

4.3 最优执行策略分析

不同策略在不同情况下取得最佳结果.

5 uGrapher 的详细设计

uGrapher 的两个主要特征:

• 能为各种图算子提供完整的语义表示
• 通过自动探索灵活动态的并行策略来实现高效执行

5.1 统一接口设计

$uGrapher$ API:

• $graph\_tensor$: 图数据
• $op\_info$: 有关计算 $edge\_op$, $gather\_op$ 和输入张量的信息
• $parallel\_info$: 指定并行策略

5.2 uGrapher 实现

解耦调度和计算: 利用基于模板的编程.
首先为并行策略手动实现 CUDA 内核模板, 然后在模板中预留支持各种图算子的设备函数接口.

代码生成实现: 实现了一个自动化端到端的代码生成程序, 确保正确性并针对不同图算子生成的 CUDA 内核进行优化.

• 第一个过程: 当 $op\_info$ 的成员 $edge\_op$ 或 $gathe\_op$ 是 $NULL$ 时, 融合最里面的两个代码语句, 从而减少寄存器使用和读/写开销.
• 第二个过程: 生成最终的设备函数代码, 可以通过分析不同线程是否竞争相同的数据来选择使用原子操作.

通过全局函数与设备函数的自由组合实现了针对不同算子的灵活高效实现. 前者提供对不同并行策略的支持; 后者提供对图算子中不同算法的支持.

5.3 与现有框架集成

使用 pybind11 实现由代码生成器生成的 CUDA 内核的 Python 调用接口.
在不更改用户代码的情况下将 uGraphe 实现为被调用的底层接口.

5.4 自适应预测最优并行策略

利用梯度提升框架 LightGBM 训练预测模型以在并行化空间中选择最优策略.

6 方法

实验、baselines、benchmarks、batchsize、数据集

7 评估

性能: Figure 13, Figure 14(不同 GPU), Figure 15(不同数据集)

性能指标: Figure 16
执行策略: Table 9, Figure 18

---

笔者总结

本文的核心在于对 GNN 中的图算子进行了统一抽象, 并将图算子的计算和调度解耦, 由此提出了支持现有 GNN 框架和提供高性能的 CUDA 代码生成的统一 API, uGrapher.