SCALE: 一种面向消息传递图神经网络的结构中心型加速器

作者与出版信息

本研究由 Lingxiang Yin, Sanjay Gandham, Mingjie Lin 和 Hao Zheng 共同完成，作者来自美国佛罗里达中央大学电气与计算机工程系。该研究发表于2024年IEEE/ACM国际微架构研讨会（MICRO）。

研究背景

随着图结构数据的广泛应用，图神经网络（Graph Neural Networks, GNNs）已成为理解复杂图结构数据的有力工具，在社交网络、推荐系统和生物信息学等领域展现出良好性能。消息传递范式被广泛用于开发复杂的GNN模型，允许简洁地表示边和顶点级操作。然而，由于边和顶点操作的异构计算内核，导致计算复杂性显著增加。特别是这种不一致的计算特性在利用中间数据重用、确保边和顶点工作负载平衡以及维持系统可扩展性方面带来了新的挑战。

研究内容与方法

SCALE加速器设计

研究者提出了一种名为SCALE的结构中心型加速器，可以支持各种消息传递GNN模型，提高并行性、数据重用和可扩展性。SCALE的核心理念是在GNN原语中发现潜在的相似性，如共享的数据流结构，而不是严格遵循异构模型结构。这作为调和各种GNN计算内核不一致性的关键。

SCALE加速器包含三个关键设计： 1. 新型脉动阵列式架构：该架构可以支持各种数据流，如全归约（all-reduce），以改善不同GNN操作的并行性、数据重用和吞吐量。它通过同时融合特征聚合（以归约操作形式）和顶点更新操作来消除图不规则性，从而显著提高性能。 2. 度和顶点感知调度：这种调度可以解决顶点和边级操作中遇到的工作负载不平衡问题，能够动态分配等价的边和顶点数量给不同的处理单元。 3. 融合图和神经操作的连贯数据流：该数据流可以统一图和神经算子的数据移动，无需额外的通信和存储开销。

系统架构

SCALE的整体架构包含以下主要组件：多banks全局缓冲区、任务控制器、数据加载器、任务调度器和灵活的处理元素（PE）阵列。多banks全局缓冲区用于存储图信息和权重矩阵。任务控制器负责调度边和顶点任务，并将它们分配给任务调度单元数组。每个PE可以同时支持消息聚合和顶点更新阶段。

PE阵列是一种类脉动阵列架构，同时启用三种类型的并行性：边（特征聚合）并行、顶点（顶点更新）并行和操作符（聚合和更新之间的依赖关系）并行。PE阵列中每行以环形拓扑相互连接，包含一组链接交换机可禁用分段短链接之间的信号传播。

聚合与更新阶段设计

聚合阶段：利用边或特征并行性同时执行多个归约操作。研究者采用个体顶点特征而非整个特征向量操作，减少了跨聚合PE的带宽要求。同时，使用带有双缓冲的移位寄存器阵列，存储多个顶点和边的特征值，以重叠特征分发与特征聚合的延迟。

更新阶段：一旦顶点特征的累积完成，就会进行顶点更新。其基本计算是矩阵-向量乘法。由于权重矩阵在整个计算过程中保持不变，因此采用了权重固定（weight-stationary）数据流。具体实现是将权重矩阵均匀分布到PE阵列中，每个PE保存一个块，而聚合特征则在PE间传递以与每个权重分区进行计算。

工作负载平衡调度算法

研究者提出了度和顶点感知任务调度算法，这是针对装箱问题（bin-packing problem）的贪婪启发式算法。该方法首先使用度感知调度创建边平衡任务，然后使用修改的顶点感知调度将这些任务组合成均匀分布的任务组。这种两步方法确保了任务组之间边和顶点工作负载的平衡。

实验设置与评估

实验设置

研究者使用经验证的周期精确C++模拟器实现了SCALE，并与配置为高带宽内存（HBM）的Ramulator集成。评估中使用了三个引文图（Cora、Citeseer和Pubmed）、一个知识图（NELL）和一个大型帖子图（Reddit），以及几种流行的GNN模型：GCN、Gated-GCN（G-GCN）、GraphSAGE-Pool（GS-PL）和GIN。

为了进行功率、面积和时序特性评估，研究者使用Verilog实现了SCALE，并使用TSMC 32nm标准库的Synopsys Design Compiler进行合成。时钟频率设置为1GHz。

性能分析

在配置为32×16 PE阵列、4MB全局缓冲区的SCALE上进行测试，每个PE有6KB本地缓冲区和2个MAC单元。与其他最先进的GCN和GNN加速器相比，SCALE在GCN模型上平均比AWB-GCN和GCNAX分别快1.62×和2.01×。对于G-GCN、GS-PL和GIN等其他模型，SCALE平均比FlowGNN和ReGNN分别快1.57×和1.80×。

性能提升主要来自改进PE利用率的平衡工作负载和简化通信模式的连贯数据流。与基线架构相比，暴露的通信延迟减少了高达87.56%，工作负载平衡贡献了聚合和更新阶段高达50.35%的延迟减少。

可扩展性分析

与基于512 MACs配置的AWB-GCN相比，当配置为4K MACs时，SCALE展示了平均12.07×的加速比，而AWB-GCN、GCNAX、ReGNN和FlowGNN分别为7.61×、6.49×、7.3×和6.68×。SCALE的更好可扩展性主要归功于统一的数据流架构、平衡的工作负载和简化的互连。

能耗和面积分析

与基线相比，SCALE平均减少了36.8%的DRAM能耗和53.2%的全局缓冲区能耗。这种能耗降低来自寄存器级别的更高中间数据重用，从而减少了对全局缓冲区和DRAM的读/写操作。在面积分布方面，存储组件（全局缓冲区和本地缓冲区）占总面积的81.4%，而MAC和任务控制分别占12.2%和6.4%。

研究结论与价值

SCALE加速器可以支持广泛的GNN模型和图不规则性，显著提高性能和能源效率。模拟结果表明，SCALE与最先进的GNN加速器相比，平均实现了1.82×的加速和38.9%的能源降低。

研究亮点

1. 统一数据流架构：通过发现GNN原语中潜在的相似性，SCALE成功地统一了不同GNN操作的数据流，从而简化了通信复杂性并提高了数据重用。
2. 三层并行性：SCALE的PE阵列同时支持边并行、顶点并行和操作符并行，显著提高了计算效率。
3. 度和顶点感知调度：提出的调度算法能够同时平衡边和顶点工作负载，解决了图处理中的工作负载不平衡问题。
4. 可扩展性：与现有解决方案相比，SCALE展示了更好的可扩展性，尤其是在处理具有高度不规则性的大型图时。
5. 能效优化：通过寄存器级别的中间数据重用，SCALE显著降低了能耗，特别是在DRAM和全局缓冲区访问方面。

SCALE提供了一种创新的方法来加速消息传递图神经网络，通过结构中心的设计理念和统一的数据流，克服了传统GNN加速器面临的关键挑战。