中国科学院计算技术研究所团队多篇文章获ASPLOS 2024录用

2024年4月27日—5月1日，第29届“ACM面向程序语言和操作系统的体系结构支持国际会议”（ASPLOS: ACM International Conference on Architectural Support for Programming Languages and Operating Systems）在美国圣地亚戈隆重召开。ASPLOS是涵盖体系结构、编程语言、操作系统等多个领域的国际顶尖学术会议，与MICRO、ISCA和HPCA等四个CCF A类会议并称计算机体系结构领域的四大国际学术会议。此次会议共收到922篇投稿，录用170篇，录用率18.4%。处理器芯片全国重点实验室共计6篇文章入选。

1、《CIM-MLC: A Multi-level Compilation Stack for Computing-In-Memory Accelerators》

共同第一作者为实验室博士生屈松云和赵世新。存算一体架构通过在存储器中执行原位计算，具有高并行性、低计算开销等优势。近年来，研究人员提出了许多基于不同存储设备的存算一体芯片，展示了其优越性。然而，随着以神经网络、图计算和大模型等为代表的应用层逐渐复杂多样，同时，学术界和工业界推出的差异化存算一体芯片架构逐渐多样化，如何提供统一的中间层编译适应来自两端的变化，对于减少重复劳动并帮助硬件领域创新者快速展示应用潜力变得十分关键，而当前开源的AI处理编译框架缺乏对不同存算一体架构上具体计算粒度的优化考虑，导致无法充分发挥存算一体加速器的潜力。故针对存算一体架构存储器件与电路设计多样化，微架构层次多样化，编程接口多样化的特点，本工作提出了面向存算一体芯片的多层次开源编译框架—CIM-MLC。CIM-MLC首先建立了统一的硬件抽象，以实现与多样化存算一体架构的兼容。架构抽象分为三个层级，可适配已有存算架构研究中广泛采用的多层次结构，为用户提供了丰富的定制化空间。为支持不同的架构调度粒度，提出了三种计算模式抽象。用户可通过指定计算模式来确认最小调度粒度，以适应不同的编程接口。在不同的计算模式下，设计了存算一体加速器的基本操作指令，建立了元操作集。使用元操作集来表示定制化控制流。

在硬件抽象的基础上，CIM-MLC提出了一种多层次调度优化方法，从粗粒度到细粒度逐步执行神经网络计算的映射优化。编译器读取用户定义的架构信息和计算模式信息后，加载待映射的神经网络模型，通过图层级、矩阵计算、向量计算层级的调度优化，充分利用存算一体架构资源。

2、《Optimizing Deep Learning Inference via Global Analysis and Tensor Expressions》

第一作者为实验室博士生夏春伟。该工作提出了一个针对DNN推理的优化编译框架，将传统编译分析技术应用在DNN的推理优化当中，能够在计算图、张量表达式和TensorIR多个层级分析张量、元素和指令之间的依赖关系，并进行语义不变的张量表达式变换，从而可以大幅度的优化DNN推理应用的性能。优化DNN在人工智能处理器上的推理性能是编译优化领域一个非常重要的问题。随着深度学习模型越来越复杂，提升其推理性能也愈加困难。算子融合技术是重要的优化手段，现有的DNN编译器难以从全局挖掘数据复用机会，难以有效地跨算子边界进行指令优化。其生成的代码还存在着许多可以优化的空间。

此次提出的深度学习编译器Souffle结合传统编译分析技术和DNN的应用特征，在计算图、张量表达式和TensorIR多个层级分析张量、元素和指令之间的依赖关系，并进行语义不变的张量表达式变换。Souffle大幅减少了生成的计算核数量和访存开销，与厂商深度优化的编译器TensorRT相比，能够获得平均3.7的加速比。

3、《Optimizing Dynamic-Shape Neural Networks on Accelerators via On-the-Fly Micro-Kernel Polymerization》

第一作者为实验室博士生俞峰。该工作提出了一个针对动态形状场景的两阶段人工智能算子编译优化方法MikPoly。由于智能应用任务的复杂性不断增加以及现实世界中广泛存在的形状多样的数据，使得动态形状神经网络模型的需求与日俱增。例如，语言模型的序列长度会因不同的输入语句而异。张量计算中引入的动态形状特性为高性能算子库和人工智能编译器进行性能优化带来了新的挑战。高效处理这些具有动态形状的张量计算，对于充分发挥动态形状神经网络模型的潜力至关重要。为此，本文提出了一种基于微内核聚合的动态形状张量程序优化方法，以提升动态形状神经网络模型在智能加速器上的执行性能。MikPoly包含两个阶段，即离线阶段和在线阶段，在轻量、准确的代价模型的指导下为运行时给定的输入形状生成一个优化的张量程序。MikPoly将张量程序模板划分为内部离线循环和外部在线循环两部分。离线阶段，MikPoly依据微内核模板生成一组高度优化的固定形状的微内核，并为每个微内核建立相应的性能模型。在线阶段，MikPoly使用聚合模式并结合参数化的微内核将在线循环构成张量程序，并通过离线阶段获得的固定形状的微内核进行实例化。MikPoly利用轻量、准确的代价模型来评估特定聚合策略下完整张量程序的性能，为给定形状的张量算子选择最高效的程序实现提供依据。在GPU、NPU等平台的实验结果验证了MikPoly的有效性，相较于高性能算子库，MikPoly实现了平均 1.49 倍的性能提升。

4、《PrimePar: Efficient Spatial-temporal Tensor Partitioning for Large Transformer Model Training》

第一作者为实验室博士生王浩然。现今，大模型被广泛应用于各个领域，成为了学术和工业界关注的热点。然而大模型的训练十分困难，需要成千上万张高性能GPU并行训练。基于这种现状，优化大模型的并行训练策略，以使得有限的硬件资源能够充分地被利用成为了重要的任务。据我们观察，大语言模型的训练瓶颈是高通信、存储开销。现有的大模型并行训练策略仅考虑空间维度的张量划分，造成了昂贵且难以与计算折叠的聚合通信以及张量在设备间的重复存储。PrimePar形式化地探索了空间-时间维度的张量划分空间，考虑将子算子分配到不同设备同时计算或同一设备的不同时间点进行计算，从而形成更广泛的张量划分空间，提供了额外的并行优化空间以更充分地利用硬件资源。PrimePar提出了新的张量划分原语，将需要聚合通信的子算子分配到同一个设备的不同时间点计算，从而避免了昂贵的聚合通信，代之以廉价的且易于与计算折叠的点到点通信。同时，该张量划分原语消除了张量在设备间的重复存储，降低了系统整体的存储开销。为了实现对Transformer模型结构的最优张量划分搜索，PrimePar提出了分段动态规划的搜索算法。该算法将Transformer模型算子图分为若干段，先在每一段中分别使用动态规划得到分段最优解，再融合成完整模型的最优解，且动态规划和融合最优解都具有较小的计算复杂度。相对于目前主流的大模型训练框架Megatron-LM，PrimePar的张量划分策略实现了1.11~1.68倍的训练吞吐和68~93%的峰值存储占用。PrimePar加速训练和节省存储的效果对于更大的模型和并行度更为显著。

5、《ACES: Accelerating Sparse Matrix Multiplication with Adaptive Execution Flow and Concurrency-Aware Cache Optimizations》

共同第一作者为博士生龙博宇和伊利诺伊理工学院博士生鲁潇阳。稀疏矩阵乘法（Sparse matrix-matrix multiplication ，SpMM）是许多科学和机器学习应用中的关键计算部分。SpMM涉及大量不规则的内存访问，对传统的基于缓存的计算机架构提出了巨大的挑战。最近提出了专用 SpMM 加速器来增强 SpMM 性能。然而，当前的SpMM加速器在适应不同的稀疏模式、充分利用硬件的并行性和优化缓存性能方面仍然面临挑战。为了解决这些问题，本文中提出了 ACES，一种新型 SpMM 加速器。首先，ACES 具有自适执行流程，可以动态调整以适应不同的稀疏模式，并平衡了并行计算效率和数据重用。其次，ACES 在全局缓存中引入了局部性（locality）和并发性（concurrency）的协同优化。ACES 采用并发感知缓存管理，该策略考虑数据的局部性和并发性以做出最佳替换决策。此外，非阻塞缓冲区与全局缓存的集成进一步加强了并发性并减少了因数据读取导致的计算阻塞。该架构确保对自适应执行流程、高级缓存优化和细粒度并行处理的有效支持。我们的性能评估表明，ACES 的性能显着优于现有解决方案，例如Sparch（HPCA‘20）和Spada（ASPLOS’23），并分别提供了8.9倍和 2.1 倍的加速，标志着SpMM 加速的重大进步。

6、《ORIANNA: An Accelerator Generation Framework for Optimization-based Robotic Applications》

共同第一作者是甘一鸣助理研究员和天津大学硕士郝宇辉。随着自动化技术的进步，机器人已在工业制造、医疗保健、农业劳动及家居服务等多元领域实现广泛应用。各类机器人功能的实现，离不开基于优化的机器人算法的支持，如定位、规划和控制等，它们极大地增强了机器人的效能。然而，这些算法需要多次构建并求解大规模线性方程组，由此产生的较高计算延迟和能耗制约了机器人技术的实际应用效率。当前，架构领域的研究者们提出了两种应对方案。一种方案是研发可编程加速器。尽管这类加速器可以适用多种机器人算法，但其性能提升相对有限，根本原因在于其未能充分利用机器人算法中固有的稀疏性；另一种方案是针对特定算法设计专用硬件加速器。虽然这类加速器能实现卓越的性能提升，但是由于不同机器人算法间存在显著差异，这类加速器无法通用。如果在一个机器人应用中集成多个算法的专用加速器，会占用较大的芯片面积。为了解决上述问题，我们提出了ORIANNA：面向多个算法的机器人应用加速器生成方案。在ORIANNA框架中，通过利用因子图作为通用抽象，我们设计了一个灵活的因子图软件库，允许用户定制自己的基于优化的机器人算法。我们提供通用因子的实现，并允许用户扩展自己的定制因子。此外，我们还设计了一个编译器，用于将高级用户程序翻译成底层矩阵运算指令。用户只需采用构建因子图的方式来表达问题，ORIANNA软件库与编译器负责将基于优化的机器人算法映射到硬件加速器中。最后，ORIANNA能够为不同机器人应用生成支持多种算法的硬件加速器，并在用户指定的约束条件下优化加速器。我们在Xilinx ZC706 FPGA上对ORIANNA加速器进行了评估。与Intel i7-11700 CPU和ARM Cortex-A57 CPU相比，ORIANNA加速器分别实现了7倍和54倍的性能提升，并且降低了93%和71%的能耗。此外，ORIANNA加速器实现了与现有专用加速器相当的性能，同时减少了64%的硬件资源消耗。