Relational graph neural networks (RGNNs) are graph neural networks (GNNs) with dedicated structures for modeling the different types of nodes and/or edges in heterogeneous graphs. While RGNNs have been increasingly adopted in many real-world applications due to their versatility and accuracy, they pose performance and system design challenges due to their inherent computation patterns, gap between the programming interface and kernel APIs, and heavy programming efforts in optimizing kernels caused by their coupling with data layout and heterogeneity. To systematically address these challenges, we propose Pigeon, a novel two-level intermediate representation (IR) and its code generator framework, that (a) represents the key properties of the RGNN models to bridge the gap between the programming interface and kernel APIs, (b) decouples model semantics, data layout, and operators-specific optimization from each other to reduce programming efforts, (c) expresses and leverages optimization opportunities in inter-operator transforms, data layout, and operator-specific schedules. By building on one general matrix multiply (GEMM) template and a node/edge traversal template, Pigeon achieves up to 7.8x speed-up in inference and 5.6x speed-up in training compared with the state-of-the-art public systems in select models, i.e., RGCN, RGAT, HGT, when running heterogeneous graphs provided by Deep Graph Library (DGL) and Open Graph Benchmark (OGB). Pigeon also triggers fewer out-of-memory (OOM) errors. In addition, we propose linear operator fusion and compact materialization to further accelerate the system by up to 2.2x.

翻译：摘要：关系图神经网络（RGNN）是具有专用结构的图神经网络（GNN），用于对异构图中的不同类型节点和/或边进行建模。尽管RGNN因其多功能性和准确性在众多实际应用中得到日益普及，但其固有的计算模式、编程接口与内核API之间的鸿沟，以及由于与数据布局和异构性耦合而导致的优化内核所需的大量编程工作，给性能和系统设计带来了挑战。为系统性地解决这些问题，我们提出Pigeon，一种新颖的两级中间表示（IR）及其代码生成器框架，该框架能够：（a）刻画RGNN模型的关键属性，弥合编程接口与内核API之间的差距；（b）将模型语义、数据布局和算子特定优化相互解耦，以减少编程工作量；（c）表达并利用算子间变换、数据布局以及算子特定调度中的优化机会。通过构建于通用矩阵乘法（GEMM）模板和节点/边遍历模板之上，Pigeon在处理Deep Graph Library（DGL）和Open Graph Benchmark（OGB）提供的异构图时，在选定模型（如RGCN、RGAT、HGT）上相比现有最先进公开系统实现了推理加速最高达7.8倍、训练加速最高达5.6倍。同时，Pigeon引发的内存溢出（OOM）错误更少。此外，我们提出线性算子融合与紧凑物化策略，进一步将系统加速最高达2.2倍。