VMOC：基于MLIR的面向向量加速器上CNN部署的内存优化技术

《Journal of Systems Architecture》：VMOC: MLIR-based memory optimization for CNN deployment on vector accelerators

　　摘要：由于具有高能效和并行处理能力，向量加速器逐渐成为部署卷积神经网络的重要平台。在这样的架构中，高效利用内存对于成功部署至关重要。然而，现有的图分裂方法在向量加速器上存在诸多局限性，包括硬件适应性差、分裂成本建模不准确、编译效率低以及过度依赖人工干预。本文提出了一种基于MLIR的新型内存优化方法VMOC。VMOC能够自动识别内存热点子图，构建一个硬件感知的分裂成本模型来量化分裂带来的内存收益与执行开销，并采用自适应图分裂算法动态搜索既满足内存约束又具有正向收益的分裂方案。同时，还引入了基于图封装的调度优化方法来缓解图分裂带来的调度复杂性问题。这些优化策略被集成到MLIR框架中，实现了从ONNX模型到目标可执行代码的端到端自动编译，从而生成适用于向量加速器的高性能代码。实验结果表明，VMOC在GoogleNet、ResNet18、VGG16和YOLOv5s上的效果显著，其内存压缩比可达52.05%～96.18%，相比MAGIS的编译速度提升2.34× ～152.11倍，相比TVM的提升幅度为12.63× ～100.93倍。在FT-M7032和FT-XT4000向量加速器上，它的计算效率分别为62.30%～74.05%和50.25%～71.01%，始终优于VCNN。这些结果表明，VMOC有望成为在向量加速器上自动部署CNN的一种高效实用方法。

引言：随着卷积神经网络在边缘计算中的日益应用，向量加速器因其高能效和并行处理能力而成为重要的目标平台，例如英特尔的Habana Gaudi2、Greco AI专用芯片以及德州仪器的C7x DSP等。这类向量加速器通常采用多级内存架构，包括容量大但延迟高的DDR内存、容量中等的全局共享内存、以及容量小但速度快的阵列内存和标量内存。在CNN推理过程中，中间特征图通常存储在全局共享内存中，以充分利用片上带宽并减少对外部内存的访问。然而，随着网络深度和特征图尺寸的增加，某些中间张量的内存占用往往超过全局共享内存的容量，不得不回写到DDR内存中，从而导致频繁的片外数据传输开销。因此，高效管理内存资源仍是向量加速器上实现高性能CNN推理的关键挑战。

为了解决内存限制问题，现有方法通常采用图等价分裂变换，将大型运算符分解为多个子运算符，从而降低每次计算的内存占用。这类方法已被广泛研究用于资源受限的硬件平台上的部署。它们的核心思想是通过等价的图变换来降低峰值内存需求。然而，分裂本身并不一定能提升性能——分裂后的边界重叠会带来冗余计算，而中间结果的重新组织则会增加数据传输量。如果分裂粒度选择不当，这些额外的开销可能会超过减少DDR访问所带来的好处，进而导致推理时间延长。当将图分裂应用于向量加速器时，现有方法仍面临四个主要挑战：首先，适应性有限，对于并非内存瓶颈的张量进行分裂无法有效降低峰值内存占用；其次，缺乏成本建模，现有方法仅通过判断是否满足内存约束来确定分裂方式，却忽视了分裂带来的额外计算和数据传输开销可能超过将中间结果保留在高速片上内存所带来的内存访问优势；第三，编译效率低，分裂会显著增加图节点的数量，从而提高后续调度和搜索的成本；第四，依赖人工部署，现有的CNN部署流程仍然高度依赖人工优化，导致开发效率低下且可移植性差，严重制约了部署效率。

与传统主要着眼于确保计算可行性的分裂方法不同，本研究关注的是分裂后的实际执行效益。分裂不仅要满足中间张量的片上内存约束，还要确保额外的计算和数据传输开销低于从片上内存获取的好处。为了解决这些挑战，本文提出了VMOC，这是一种基于MLIR的、用于在向量加速器上部署CNN的内存优化方法。VMOC旨在提高向量加速器上的内存利用率，同时提升推理性能和编译效率。该工作的主要贡献如下：首先，提出了一种内存热点子图识别方法，可自动定位那些直接决定峰值内存占用的关键运算符；其次，设计了一个硬件感知的分裂成本模型，能够在全局共享内存容量约束下量化分裂带来的收益与开销，避免过度分裂导致的性能下降；第三，提出了一种自适应图分裂算法，能够针对不同的热点子图动态搜索出在满足全局共享内存容量约束的同时还能带来正向收益的最小可行分裂粒度，之后再进行子图重构；第四，引入了一种基于图封装的调度优化方法，通过将分裂后产生的局部子图封装为超节点，来抑制图结构膨胀并提升CNN的编译效率；最后，基于MLIR构建了一个端到端的CNN编译框架，该框架能够从ONNX模型自动生成高性能目标代码，提供完整的工具链支持以实现向向量加速器的高效部署，进一步丰富了MLIR生态系统。

向量加速器：向量加速器通常具有与x86 CPU和GPU等主流处理器不同的架构特点（如图1所示）。它们主要由多个通用数字信号处理器集群组成，这些处理器通过片上网络相互连接。其内存系统采用三级架构，包括容量大但延迟高的DDR内存、MB级别的多核共享全局共享内存，以及KB级别的高速阵列内存。

VMOC概述：本文提出了一种基于MLIR的VMOC方法，用于在向量加速器上部署CNN。如图2所示，整个框架由两个主要部分组成：VMOC优化流程和多层编译管线。第一部分是VMOC优化流程，包含四个核心阶段：首先自动识别导致峰值内存使用的记忆热点及其相关子图；随后进行自适应图分裂……

VMOC方法：本节详细介绍了VMOC的三个关键方面：热点识别、自适应图分割以及基于封装的调度优化。为便于后续的数学建模和公式推导，表1总结了常用的关键符号和表示方法。

VMOC代码生成器：为应对深度学习模型日益复杂的趋势以及向量加速器上较低的部署效率，基于MLIR实现了一个集成VMOC的自动代码生成器。该生成器构建了从ONNX模型到低级目标指令的完整编译链，从而实现向量加速器的自动优化代码生成。编译流程涉及三种关键方言：(a) ONNX方言用于表示模型计算的运算符级语义。

实验设置：本研究采用配备英特尔Xeon Gold 6248R CPU的工作站作为代码开发和编译的平台。最终目标代码将在国防科技大学自主研发的FT-M7032和FT-XT4000高性能向量处理器上运行。FT-M7032处理器拥有四核，配备64 GB DDR内存、8 MB全局共享内存、512 KB阵列内存以及32 KB标量内存。其理论带宽为……

结论：为解决在向量加速器上部署卷积神经网络时的内存效率和编译优化问题，本文提出了基于MLIR的VMOC内存优化方法。VMOC能够自动识别内存热点，构建硬件感知的分裂成本模型来量化分裂的收益与开销，针对内存热点自适应地应用图分裂，并通过基于子图封装的调度优化方法提升编译效率。

讨论：目前VMOC方法主要适用于典型的CNN模型，未来可进一步扩展到Transformer类模型。我们的初步方案预计包括以下四个阶段：(1) 热点识别与维度映射，通过识别热点，精准定位注意力机制、MLP和LayerNorm等模块中导致峰值内存占用的热点子图；(2) 成本模型重构……

作者贡献声明：陈晓荣负责原文撰写、软件实现、方法设计及数据整理工作；刘忠负责文章审阅与编辑以及整体指导工作。

利益冲突声明：作者声明不存在任何可能影响本文研究的已知财务利益或个人关系。

致谢：本研究得到了国防科技大学创新研究基金的支持，项目编号为23-ZZCX-JDZ-11。

陈晓荣|刘忠
中国长沙，国防科技大学计算机科学与技术学院，邮编410073