Hector：一种用于在 GPU 架构上实现关系图神经网络的高效编程与编译框架
分布式、并行与集群计算
2024-04-10 v3 神经与进化计算
摘要
关系图神经网络（RGNNs）是带有专用结构以建模异构图中所不同类型节点与边的图神经网络。尽管 RGNNs 因通用性与准确性日益应用于许多实际场景，它们带来了性能与系统设计挑战：固有的内存密集型计算模式、编程接口与内核 API 之间的鸿沟，以及因与数据布局和异构性耦合而在优化内核中产生的繁重编程负担。为系统应对这些挑战，我们提出 Hector，一种新颖的两级中间表示及其代码生成框架，其（a）捕捉 RGNN 模型的关键属性，以及在内核间调度与物化中减少内存访问的机会，（b）生成具有灵活数据访问方案的代码以消除冗余数据拷贝，（c）将模型语义、数据布局与算子特定优化彼此解耦以降低编程负担。基于一个通用矩阵乘（GEMM）模板与节点/边遍历模板，Hector 在推理上实现最高 9.9 倍加速、在训练上实现最高 43.7 倍加速，优于最先进公开系统，所选模型为 RGCN、RGAT 和 HGT，运行于 Deep Graph Library（DGL）与 Open Graph Benchmark（OGB）提供的异构图。此外，Hector 在这些测试中未触发任何内存溢出（OOM）异常。我们还提出线性算子重排与紧凑物化，将系统进一步加速最高 3.8 倍。作为编程负担降低的 indicator，Hector 以 51 行代码表达这三个模型并生成总计 8K 行 CUDA 与 C++ 代码。
引用
@article{arxiv.2301.06284,
title = {Hector: An Efficient Programming and Compilation Framework for Implementing Relational Graph Neural Networks in GPU Architectures},
author = {Kun Wu and Mert Hidayetoğlu and Xiang Song and Sitao Huang and Da Zheng and Israt Nisa and Wen-mei Hwu},
journal= {arXiv preprint arXiv:2301.06284},
year = {2024}
}
备注
Accepted by ASPLOS