基于消息的向量处理单元上 AI 负载的内存感知调度

《Future Generation Computer Systems》：Memory-aware scheduling of AI workloads on a Messaging-based vector processing unit

　　

摘要

由于人工智能（AI）应用的快速增长，人工智能加速器芯片需求旺盛。这些芯片采用大量小型算术单元以实现并行计算。除了计算需求之外，AI工作负载还对片上嵌入式内存施加了巨大压力，因为必须持续供应大量数据和指令。然而，实际设计受限于功耗和面积管理，内存容量受到约束，从而限制了能在计算单元附近存储的数据量。这在可用的计算并行性与执行期间维持该并行性的能力之间造成了差距。在本工作中，我们提出了一种内存感知的调度框架，根据内存可用性调节数据何时以及如何进入计算流水线。我们展示了如何将嵌入式内存结构化为分层的缓冲区和缓存层次结构，以在避免执行停顿的同时平衡性能。该框架不是在有数据可用时立即传输数据，而是调度数据移动，使操作数只有在可以被消费和存储而无冲突时才到达。这种协调交付确保了权重、激活值和部分和以一种支持连续执行的方式被供应。我们使用周期精确建模在MAVeC架构上评估了所提出的框架，该建模经过RTL实现校准，并在28 nm工艺、1 GHz频率下采用三种可扩展配置（16 × 16、32 × 32和64 × 64）的层次化片上内存系统进行测试。在AlexNet、VGG-16、VGG-19、ResNet-18和ResNet-50上，所提出的方法在内存调节执行下实现了稳定的高利用率和可预测的可扩展性。对于ResNet-50，64 × 64配置实现了39.73 ms的延迟、187.75 GOp/s的吞吐量、349 GOp/s/W的能效和25.17次每秒的推理率。内存流量分析表明，仅有6%的访问发生在片外，而94%在片上服务。在AlexNet、VGG-16、ResNet-18和ResNet-50上，89%至94%的所有内存访问保持在片上，证明了在层次化内存调节下持续的数据驻留。

引言

人工智能工作负载要求大量数据持续通过计算流水线移动。现代加速器通过处理单元、互连和数据编排的设计与架构创新来满足这一需求[1]、[2]。真正的决定因素在于执行过程中在计算架构上保持指令和操作数的不间断流动[3]。这一约束推动了广泛的加速器设计[4]、[5]、[6]、[7]、[8]，每种设计在数据交付方式、组件连接方式、内存组织方式以及工作负载在硬件资源上的映射方式上各不相同。随着工作负载在规模和复杂性上增长，负担日益转移到内存系统上，数据供应延迟可能导致即使高度并行的计算架构也无法得到充分利用[9]。

先前的研究通过广泛的加速器架构和映射框架来解决这些挑战，旨在改善数据局部性并减少移动开销。脉动架构如TPU[5]和TriM[10]采用结构化数据传播模式，在处理单元间复用数据。TPU在规则的二维处理单元网格上执行计算，权重被加载一次并本地保留。输入激活值沿一个方向流经阵列，而部分和沿另一个方向传播。这种组织使每个权重和激活值在替代之前可以参与多次乘累加操作，从而减少对片外内存的重复访问。相比之下，TriM将处理单元组织为三角形布局，并根据预定的执行顺序沿对角线调度输入激活值的移动。当激活值遍历相邻元素时，每个元素在离开阵列之前为多次计算做出贡献，实现了复用并保持可预测的通信模式。其他基于ASIC的空间加速器追求不同的策略。CARLA[11]将计算分布在集群化的可重构阵列上以增强可扩展性，Eyeriss[6]采用行驻留数据流，平衡权重、激活值和部分和的复用，以最小化跨内存层次结构的数据传输，而MAERI[7]利用可重构互连，将数据移动模式适应工作负载的结构。基于FPGA的平台如NEURAghe[12]通过可编程流水线和可重构数据路径提供灵活性，允许通过协调的软硬件调度针对各个神经网络层定制执行。与这些硬件进展相辅相成的是，映射框架如Timeloop[13]和MAESTRO[14]优化循环变换、平铺策略和内存分配，以在架构约束下最大化复用。

尽管取得了这些进展，许多加速器设计和映射框架仍主要依赖编译时决策进行数据放置和执行调度[15]、[16]、[17]。在这些方法中，数据复用、平铺策略和内存分配在执行前根据估计的复用模式和架构约束确定。近期研究引入了更具适应性的执行范式，包括运行时感知编排、动态内存管理、通信调度和灵活调度机制[18]、[19]、[20]、[21]、[22]。这些方法通过在执行期间动态管理资源或通信模式来改善执行适应性。然而，它们主要优化调度策略、内存分配或通信行为，而未明确调节数据在层次化片上内存系统中的时间驻留性。

因此，在不断变化的内存占用和通信需求下高效协调层次化片上内存资源仍然是一个基本挑战，特别是对于维持不间断的流水线执行而言。仅靠编译时策略无法完全捕捉数据在执行期间的演变时间行为。不同类型的数据表现出不同的生命周期和复用特征。权重可能长时间驻留，输入激活值必须及时到达以维持计算，而部分和持续到归约完成为止。随着这些数据流的重叠，片上内存层次结构的占用率持续变化，要求执行决策适应瞬时的内存状态而非静态的编译时假设[23]、[24]。

这种演变的占用率将内存容量从固定的设计参数转变为运行时约束。启动或维持计算的能力不仅取决于处理单元的可用性，还取决于是否存在足够的内存来接纳传入数据，同时保留正在进行的运行所需的数据[25]。当内存需求超过容量时，系统可能会经历背压，导致停顿或额外的数据传输，破坏流水线连续性。为解决这一局限性，我们引入了一种内存感知的调度框架，将执行与片上内存资源的可用性协调起来。所提出的方法没有将内存视为被动存储，而是赋予其主动角色来管理执行进展。只有在存在足够容量来容纳传入数据同时保留系统中已驻留数据的情况下，计算才被接纳。该框架的核心概念是片上驻留窗口，它表征不同类型数据占用内存的持续时间，并指导与权重加载、激活值接纳和部分和累加相关的决策。

我们利用基于消息的自适应向量计算（MAVeC）架构[26]来实现和评估该框架。MAVeC具有层次化片上内存组织，与消息驱动的执行模型紧密耦合，在该模型中数据和指令作为统一的数据包通过计算架构传播[27]。这种组织使处理单元和内存资源之间能够紧密协调，允许执行按照内存可用性推进。利用该平台，我们分析了权重、激活值和部分和的时间行为，并量化了内存感知调度对利用率、数据移动和整体系统性能的影响。本工作的主要贡献总结如下：
• 一种根据层次化片上内存可用性调节执行的内存感知调度框架
• 不同工作负载下内存层次结构时空行为的表征
• 在AlexNet、VGG-16、VGG-19、ResNet-18和ResNet-50等CNN模型上的利用率、数据移动、吞吐量、延迟和能效评估
• 与TriM[10]和Eyeriss[6]等其他空间加速器的比较

本文其余部分组织如下。第2节描述MAVeC架构及其层次化内存组织。第3节介绍所提出的内存感知调度框架及其与执行流水线的集成。第4节分析内存约束下的执行行为。第5节提供性能评估和与先前工作的比较。第6节总结全文。

消息驱动向量处理单元

本节阐述基于消息的向量处理单元（MAVeC）的硬件架构和执行流水线。我们首先描述消息驱动的执行模型和底层架构组织，随后是数据传播。然后我们呈现执行流水线以说明数据驻留如何在各内存层级间演变。这些机制揭示了MAVeC的一个关键特性：执行推进受层次化内存容量和数据驻留约束。

内存感知调度框架

第2节确立了MAVeC中流水线推进受片上内存容量、层次化数据传播和内存事务共同约束，而非仅受算术吞吐量限制。相应地，所提出的框架根据内存可用性以及权重、消息和中间数据在整个片上内存层次结构中的传播来调节计算接纳。为解决这些约束，我们引入了一种内存感知的任务调度框架，该框架……

内存感知执行行为

执行行为由L0和L1驻留窗口通过内存感知调度来管理。随着执行在连续折叠和块交互中推进，驻留权重和飞行中的数据的占用率在层次化内存约束下演变。尽管算术资源可能仍然可用，但权重在L0中的持久存在和飞行中的部分和在L1中的持久存在决定了何时可以接纳新计算。为表征这种行为，我们分析了四个代表性的……

实验方法

我们的评估基于经过验证的MAVeC硬件基线和周期精确执行模型，以表征性能和片上内存行为。MAVeC微架构已实现为可综合RTL，并通过面向ASIC的设计流程在针对台积电28 nm工艺的目标下进行运行。硬件功能、层次化路由和多级内存组织已在我们的先前MAVeC研究中建立[26]、[28]、[29]，因此完整的RTL验证和……

结论

本工作证明了空间加速器中的执行连续性可以通过对层次化片上内存驻留的显式调节进行系统性管理。已通过将确定性的折叠接纳控制集成到执行模型中，片上内存容量约束从被动的存储限制提升为主动的调度参数。由此产生的驻留窗口界定了重叠流水线操作期间的内存占用，并实现了……的可预测演变。

CRediT作者贡献声明

Md Rownak Hossain Chowdhury：可视化、方法论、形式化分析、概念化。Kazi Asifuzzaman：写作——审阅与编辑。Mostafizur Rahman：监督。

利益冲突声明

作者声明他们不存在已知的竞争性财务利益或个人关系，这些利益或关系可能会影响本论文中报告的工作。

Md Rownak Hossain Chowdhury | Kazi Asifuzzaman | Mostafizur Rahman
密苏里大学堪萨斯城分校电气与计算机工程系，美国堪萨斯城密苏里州