一种针对异构人工智能加速器上高性能SIMV操作的、能适应架构变化的优化策略

《Journal of Systems Architecture》：An architecture-adaptive optimization strategy for high-performance SYMV on a heterogeneous AI accelerator

华为Ascend A2异构计算中对称矩阵向量乘法优化研究，提出架构自适应映射策略，通过三阶段软件流水线实现跨引擎执行重叠，结合对称感知遍历和转置无关对角块策略降低数据访问开销，在Ascend A2上实现1.3-1.6倍加速，并与A100、MI210平台库进行对比。

　　

姜浩|卢璐|梁志宏

华南理工大学计算机科学与工程学院，广州，510006，中国

摘要

新兴的AI加速器为高性能计算（HPC）工作负载提供了强大的计算密度，但解耦的执行引擎和软件管理的内存系统使得性能的可移植性变得复杂。本文研究了在华为Ascend A2上的对称矩阵-向量乘法（SYMV）内核，该架构具有分离的立方体（AIC）和向量（AIV）引擎。我们提出了一种架构适应性映射方法：(i) 将非对角线密集计算任务分配给AIC，而对角线计算和最终化处理留在AIV上；(ii) 通过三阶段软件流水线协调跨引擎执行，以实现DMA、计算和同步的重叠；(iii) 利用三角形存储结构下的对称性降低片外矩阵读取流量，并在AIV上采用无需转置的对角线计算策略。在Ascend A2上，所提出的内核相比供应商提供的matmul_gemv基线实现了1.3$\times$–1.6$\times$的加速比。我们还提供了与cuBLAS（A100）和rocBLAS（MI210）的跨平台对比结果。

引言

近年来，专为深度学习设计的AI加速器已成为传统高性能计算（HPC）工作负载的强大平台[1]、[2]、[3]、[4]。这些加速器的专用架构通常具有针对矩阵运算进行的大规模并行处理单元，为加速科学模拟和数据分析提供了巨大潜力[5]、[6]、[7]。然而，这些加速器的架构与通用CPU或GPU[8]、[9]有显著差异，给现有高性能算法的移植和优化带来了独特挑战[10]、[11]、[12]。与依赖硬件线程调度器和灵活内存层次结构的GPU不同[13]，像Ascend A2这样的AI加速器采用异构设计，其性能在很大程度上依赖于软件管理的显式流水线[14]。这种架构差异使得在GPU上证明有效的优化策略无法直接移植[15]，因此需要针对这些特定硬件约束制定新的设计方法。

在基础HPC内核中，对称矩阵-向量乘法（SYMV）特别值得关注，因为它在科学和工程计算中应用广泛，包括迭代求解器、分子动力学和特征值问题[16]、[17]、[18]。作为一种典型的内存带宽受限操作，SYMV的性能几乎完全取决于内存子系统的效率和数据访问模式[19]。因此，在像Ascend A2这样的优化为密集计算设计的架构上实现高性能SYMV的核心挑战在于，如何在严格的软件管理流水线模型下协调其功能不同的异构单元，以克服内存瓶颈。

在本文中，我们研究了如何在解耦的异构加速器上实现高性能SYMV，以华为Ascend A2作为具体实例。与单体GPU不同，Ascend A2暴露了分离的计算引擎——立方体（AIC）和向量（AIV）单元，这些单元之间的硬件支持有限，难以实现跨引擎数据可见性。因此，直接采用传统的SYMV映射方法可能会导致协调和内存服务开销出现在关键路径上。

我们提出了一种架构适应性SYMV设计：(i) 根据规律性和控制需求分解工作负载；(ii) 通过显式软件流水线协调跨引擎执行；(iii) 利用三角形存储结构下的对称性减少片外矩阵读取流量。此外，我们在AIV上引入了一种无需转置的对角线计算算法，以避免显式的片上数据重新排列。

我们的贡献如下：

• •
  针对Ascend A2的架构适应性工作负载分解：我们将规则的非对角线计算任务映射到AIC，将对角线计算和最终化处理映射到AIV，使SYMV结构与平台的解耦执行模型相匹配。
• •
  流水线跨引擎协调和对称性感知遍历：我们设计了三阶段软件流水线以实现显式协调和重叠，并采用三角形存储下的对称性感知数据流来减少片外矩阵读取流量。
• •
  在Ascend和跨平台环境下的全面评估：在Ascend A2上，所提出的内核相比供应商提供的matmul_gemv基线实现了1.3$\times$–1.6$\times$的加速比。我们还与cuBLAS（A100）和rocBLAS（MI210）进行了对比，并提供了分析结果以揭示规模依赖性和剩余瓶颈。

本文的其余部分安排如下：第2节回顾相关工作；第3节分析架构约束并阐述设计动机；第4节介绍内核设计；第5节评估性能并进行跨平台比较；第6节讨论观察到的现象、限制以及更广泛的适用性；第7节总结本文。

相关研究

基本线性代数子程序（BLAS）的性能优化，尤其是在多核和众核架构上，一直是数十年的研究课题[20]、[21]、[22]、[23]、[24]。ATLAS[25]和GotoBLAS[26]等库的基础工作确立了针对特定架构进行调优以实现高性能的重要性。在GPU平台上，针对SYMV等二级内核的优化策略已经成熟[18]、[27]、[28]、[29]。研究人员通常会利用...

背景与动机

在Ascend A2上优化SYMV内核需要精确理解算法特性和底层Da Vinci架构的具体约束。本节分析了这些因素，以确定需要我们提出的优化范式的根本性能瓶颈。

高性能SYMV内核设计

为了解决第3节中确定的架构约束，我们为Ascend A2设计了一种架构适应性SYMV映射方法。关键思想是将SYMV分解为具有不同控制和算术特性的组件，并将它们映射到A2的分离计算引擎（AIC/AIV）上，同时使用软件流水线来实现数据移动、计算和同步的重叠。

实验评估

本节对在Ascend A2上提出的异构SYMV内核进行了实证评估。我们将其与（i）Ascend上的实际供应商基线和（ii）NVIDIA A100及AMD MI210上的供应商优化SYMV库进行了基准测试。我们的评估重点包括：Ascend上的整体加速比、跨平台竞争力、硬件效率分析以及识别限制因素的实验。

讨论

本节讨论了评估中观察到的性能特征和限制，重点关注规模依赖性行为、瓶颈情况以及Ascend A2的解耦执行模型对带宽驱动的二级内核的影响。

结论

本文研究了如何在解耦的异构加速器上实现高性能SYMV，以华为Ascend A2作为具体案例。我们解决了由分离的计算引擎（AIC/AIV）带来的实际挑战，以及跨引擎协调和数据可见性的需求。

我们提出了一种架构适应性SYMV设计，该方法根据规律性分解工作负载，通过三阶段软件流水线协调执行，并减少了片外矩阵读取流量。

CRediT作者贡献声明

姜浩：撰写——原始草稿、验证、软件实现、方法论设计、概念化。卢璐：撰写——审阅与编辑、监督。梁志宏：监督、项目管理、资金筹集。

利益冲突声明

作者声明他们没有已知的竞争性财务利益或个人关系可能影响本文报告的工作。

致谢

本工作得到了中国广东省自然科学基金（2024A1515010204）和中国南方电网公司技术研究项目（ZBKJXM20232483）的支持。

姜浩，男，2001年出生，中国。他目前在中国华南理工大学计算机科学与工程学院攻读博士学位。他的研究兴趣包括异构计算并行加速和性能优化以及AI评估。