博客

使用 ExecuTorch 和 Arm SME2 加速设备端 ML 推理

交互式图像分割已成为全球各类热门应用中具有代表性的移动体验。简单来说,你在图像上点击(或画出粗略提示),应用即可通过生成像素掩码瞬间“抠出”物体。这使得诸如制作个性化贴纸、隔离主体以进行背景替换,或对图像部分区域进行选择性增强等常见功能成为可能。如果你使用过 Instagram 的抠图工具,就会发现物体掩码直接出现在设备上是多么容易。这些结果是由通过 ExecuTorch(PyTorch 的开源端侧推理运行时)和 Arm SME2 (可扩展矩阵扩展 2) 运行的紧凑型分割模型所驱动的。

本博客探讨了这些硬件和软件的进步如何使 SqueezeSAM(Instagram 抠图功能背后的端侧交互式分割模型)的图像分割速度提升高达 3.9 倍,以及这对移动应用开发者的广泛影响。正如这篇博客文章所述,SqueezeSAM 已部署在 Meta 的系列应用中。

本文受众:从事移动和边缘设备端侧 AI 部署的机器学习 (ML) 工程师和开发人员。如果您想了解 SME2 对推理性能的影响,以及如何通过算子级分析优化模型,本文将为您提供参考。

移动端侧 AI 的兴起

随着端侧 AI 的发展,一个核心问题是:当性能更强的模型能够在严苛的移动端功耗和延迟约束下运行得更快时,我们将能实现什么?在实践中,许多交互式移动 AI 功能和工作负载已经在 CPU 上运行,因为它总是可用且与应用无缝集成,同时在多种不同场景下提供高灵活性、低延迟和强大的性能。对于这些部署,性能通常取决于 CPU 执行矩阵密集型内核的效率,以及在计算不再成为瓶颈后还存在哪些限制因素。

SME2 是 Armv9 架构中引入的一套高级 CPU 指令集,旨在直接在设备上加速面向矩阵的计算工作负载。我们量化了 SME2 在 ExecuTorch 和 XNNPACK 部署中对端到端推理的加速效果,并通过算子级分析展示了改进的具体方向。全新支持 SME2 的 Arm CPU 已在用于旗舰智能手机和下一代 PC 的 Arm Lumex 计算子系统 (CSS) 中亮相,支持 SME2 的设备列表可查看此处

案例研究:利用 SME2 加速交互式图像分割

我们测量了 SME2 在使用 ExecuTorch 和 XNNPACK 作为后端(利用 Arm KleidiAI 优化内核以发挥 SME2 加速能力)进行 SqueezeSAM 端到端推理时对延迟的影响。

在开启 SME2 后,8 位整数 (INT8) 和 16 位浮点 (FP16) 推理均获得了显著的速度提升(图 1)。在默认功耗设置的单 CPU 核心上,INT8 延迟提升了 1.83 倍(从 556 毫秒降至 304 毫秒),而 FP16 提升了 3.9 倍(从 1,163 毫秒降至 298 毫秒)。如果没有 SME2,这些延迟对于响应式交互使用而言过高;有了 SME2,端到端推理在单核心上可达到约 300 毫秒的水平,使得端侧执行变得可行,同时为应用的其他部分留出了余量。

这些结果表明,SME2 实质性地加速了 CPU 上的量化 INT8 模型。同时,在本案例研究中,SME2 使 FP16 的延迟接近 INT8,这一点值得注意,因为它扩展了实际部署选项的范围,而不是简单地替代 INT8。这为开发人员提供了更大的灵活性,可以根据精度和工作流需求选择最合适的精度,特别是对于图像超分辨率、图像抠图、低光去噪和高动态范围 (HDR) 增强等对精度敏感的工作负载。如果没有这种级别的 FP16 加速,移动部署通常为了满足延迟目标而被迫倾向于使用 INT8,即使这意味着必须采用量化工作流并承担精度下降的风险。

除了基准测试数据,这些提速直接转化为空闲的 CPU 计算余量。这些余量可用于更丰富的体验,例如并行运行分割和增强功能(例如去噪或 HDR),同时保持相机预览和 UI 的响应速度;或者将抠图功能从单张图片扩展到视频实时主体跟踪;亦或用于降低功耗。

图 1. 在正常模式(默认移动功耗设置)下,在 1 个 CPU 核心上启用和禁用 SME2 时 SqueezeSAM 的端到端延迟。INT8 从 556 毫秒优化至 304 毫秒(1.83 倍)。FP16 从 1,163 毫秒优化至 298 毫秒(3.90 倍),在本案例研究中 FP16 的延迟已接近 INT8。

本文中的所有结果均为在搭载支持 SME2 的 Arm CPU 的 vivo X300 Android 旗舰智能手机上的受控测量结果。性能可能会因型号、硬件和特定设备设置而异。

技术栈:PyTorch、ExecuTorch、XNNPACK、Arm KleidiAI 和 SME2

框架连接方式

上图总结了本案例研究中使用的 CPU 执行栈。模型在 PyTorch 中定义,由 ExecuTorch 导出并运行,CPU 计算被委托给 XNNPACK 作为后端。XNNPACK 使用 Arm KleidiAI,这是 Arm 的轻量级优化 CPU 内核库,用于在 Arm CPU 上加速 ML 工作负载。这些内核可以在支持的设备上自动利用 SME2,同时为非 SME2 系统上的其他 CPU 特性提供优化实现。

当 ExecuTorch 在启用 XNNPACK 委托的情况下运行模型时,XNNPACK 会根据底层硬件的能力在运行时选择适当的内核实现。在支持 SME2 的设备上,这使得这些操作中的矩阵乘法计算无需更改模型架构或应用代码即可受益于 SME2 加速。一旦这些操作得到加速,推理管道的其他部分(如数据移动、布局转换和未委托的算子)通常会成为新的瓶颈。这就是为什么算子级分析对于理解端到端性能至关重要的原因。

案例研究模型

为了进行评估,我们使用了 SqueezeSAM,它采用了轻量级的、卷积密集的 UNet 架构,具有许多移动视觉模型的代表性。

模型结构自然映射为两类对端到端推理时间有重大影响的主要工作:

  1. 数学密集型操作:卷积(iGEMM,隐式通用矩阵乘法)和注意力/MLP 层(GEMM,通用矩阵乘法)
  2. 数据移动操作:转置、重塑和布局转换

平台说明:在许多基于 Armv9 架构的设备上,SME2 是作为 CPU 核心间的共享执行资源实现的,其扩展行为可能因片上系统 (SoC) 和 CPU 微架构而异。我们在评估中明确考虑了这一点,并在解释单核和多核结果时讨论了其影响。

结果:INT8 和 FP16(1 个 CPU 核心 vs 4 个 CPU 核心)

我们对同一模型在两种精度(INT8 和 FP16)下,开启和关闭 SME2 的情况进行了基准测试。我们重点关注单核执行(SME2 在此提供最大的相对增益),并报告了四核结果,以展示在 SME2 作为共享硬件资源时的绝对延迟和扩展行为。所有测量均报告仅模型推理的延迟。

该模型在开启和关闭 SME2 的 Android 智能手机上使用 ExecuTorch 执行,且在完全相同的软件和系统条件下运行。除非另有说明,否则结果反映的是没有热节流情况下的稳态性能。

所有结果均报告为“正常模式 | 无约束模式 (毫秒)”。正常模式对应于启用了系统电源策略的默认移动功耗设置,代表典型的终端用户行为。无约束模式对应于接通电源、保持唤醒的配置,其中 CPU 频率上限被有效移除;对于单核测量,无约束模式的结果锁定在最高性能(Ultra/Prime,本例中为 4.2 GHz)CPU 核心上。

在两种模式下,SME2 都表现出一致的相对提速趋势,表明其收益对系统电源策略具有稳健性,尽管绝对延迟有所不同。除非另有明确说明,本文的其余部分将重点关注正常模式的结果,因为它们更好地反映了在典型智能手机操作条件下的用户感知延迟。包含无约束模式结果是为了说明性能余量和硬件极限,应将其视为最佳行为,而非日常终端用户体验。

精度 核心数 SME2 关闭 (ms) SME2 开启 (ms) 加速比
INT8 1 556 | 334 304 | 172 1.83× | 1.95x
4 195 | 106 180 | 104  1.08× | 1.03x
FP16 1 1,163 | 735 298 | 173 3.90× | 4.26x
4 374 | 176 193 | 124 1.94× | 1.42x

表 1. 在 Android 手机上启用和禁用 SME2 时 SqueezeSAM 的端到端延迟结果,分别在 1 个 CPU 核心和 4 个 CPU 核心上测量(仅模型推理延迟)。数值报告为“正常模式 | 无约束模式”。

关于四核扩展的说明:四核上较小的加速比(例如 INT8 在正常模式下为 1.08 倍,而单核为 1.83 倍)与 SME2 作为共享资源的情况一致,同时也受到内存带宽和缓存行为等其他共享系统效应的影响。扩展特性可能因 SoC 和 CPU 实现而异。在生产部署中,如果满足延迟目标,通常优先选择 1 到 2 个核心以提高能效;当需要更低的绝对延迟且功率预算允许时,可以使用额外的核心。

为什么算子级分析很重要

端到端延迟告诉我们性能提高了多少,但不能告诉我们原因,或者下一步应该优化什么。为了了解 SME2 在哪里获得了收益以及下一个瓶颈是什么,我们使用算子级分析。

我们使用 ExecuTorch DevTools 中的分析工具 ETDump 收集每算子的计时信息,该工具记录推理过程中单个算子的执行时间。这使我们能够将端到端加速归因于模型的特定部分,如图 2 和表 2 所示。

为了使分析具有可操作性,我们将算子分为一小组类别,这些类别可以清晰地映射到常见的模型结构:

  • Convolution:Conv2d 层(通常使用 iGEMM 实现)
  • GEMM:Matmul 和线性层(注意力机制和 MLP 投影)
  • Elementwise:ReLU、GELU、Add、Mul 等逐点算子
  • Data Movement:转置、复制、转换、重塑和填充
  • Other:未委托的算子和框架开销

有了这种细分,我们可以解释 SME2 在何处帮助最大,以及矩阵计算加速后还剩下什么。

图 2. 在 Android 智能手机上启用和禁用 SME2 时 FP16 和 INT8 的算子类别细分(绝对时间)(1 个 Arm CPU 核心,默认移动功耗设置)。SME2 大幅减少了卷积和 GEMM 时间,数据移动在运行时中所占份额增大。

分类 INT8 SME2 关闭 (ms) INT8 SME2 开启 (ms) INT8 加速比 INT8 % (开启) FP16 SME2 关闭 (ms) FP16 SME2 开启 (ms) FP16 加速比 FP16 % (开启)
Convolution 309.7 69.8 4.4× 23.0% 881.2 98.1 9.0× 32.9%
GEMM 27.3 8.1 3.4× 2.7% 31.6 7.6 4.1× 2.6%
Elementwise  2.1 2.2 1.0× 0.7% 1.6 1.7 0.9× 0.6%
Data Movement 123.0 125.8 1.0× 41.4% 139.0 119.1 1.2× 39.9%
其他 93.7 98.2 1.0× 32.3% 109.6 71.7 1.5× 24.0%
E2E 555.8 304.1 1.83× 1,163.0 298.2 3.90×

表 2. INT8 和 FP16 的算子级细分,SME2 关闭 vs 开启(Android 手机,默认移动功耗设置下的 1 个 CPU 核心)。非 matmul 算子主要受运行时变动影响。

从端到端和算子级结果中得出的三个洞察

洞察 1:SME2 加速矩阵计算,将瓶颈转移至数据移动

SME2 实质性地降低了 INT8 和 FP16 的端到端延迟。在单个 Arm CPU 核心上,INT8 提高了 1.83 倍(556 毫秒至 304 毫秒),FP16 提高了 3.90 倍(1,163 毫秒至 298 毫秒)。即使在四核上,SME2 也显著降低了 FP16 的延迟(374 毫秒至 193 毫秒)。这些增益将单核执行带入约 300 毫秒的范围,使交互式端侧执行变得可行,同时为应用的其余部分保留了 CPU 余量。

算子级分析显示,SME2 大幅加速了矩阵密集型算子。在禁用 SME2 的情况下,卷积和 GEMM 主导了推理时间,分别占 INT8 运行时的 55.7% 和 FP16 运行时的 75.8%。在启用 SME2 后,这些算子在 GEMM 上加速了约 3-4 倍,在卷积/iGEMM 上加速了 4-9 倍,这是端到端提速的主要驱动因素。

矩阵计算加速后,数据移动和框架开销的相对成本增加,转移了进一步优化工作的重心。

洞察 2:转置驱动的数据移动占运行时的 ~40%

在 SME2 加速后,数据移动成为主要的运行时组件之一。在启用 SME2 的 INT8 运行中,数据移动占总运行时的 41.4%(FP16:39.9%)。ETDump 跟踪显示,约 85% 的数据移动时间来自转置算子,仅两种转置节点类型就消耗了该类别 80% 以上的时间。

这种开销是由模型和运行时不同部分之间的布局不匹配驱动的,而不是由算术强度驱动的。在实践中,当布局偏好不同的算子被串联执行时,就会出现这种情况,迫使进行重复的 NCHWNHWC 转换。在此模型中,当非线性激活位于 Conv2d 和 BatchNorm 之间导致无法融合时,归一化会作为便携式 NCHW 算子执行,而 XNNPACK 卷积内核则偏好 NHWC。这导致在 UNet 编码器-解码器块中出现重复的布局转换。

BatchNorm/GroupNorm (NCHW)  Transpose (NCHWNHWC)  Convolution (NHWC)  Transpose (NHWCNCHW)  BatchNorm/GroupNorm (NCHW)

由于此成本是由模型和运行时布局选择而非算术强度驱动的,因此分析对于发现它并将其转化为可操作的优化目标至关重要。

重要的是,这一分析洞察已被证明是切实可行的。作为第一步,我们在 ExecuTorch 中实施了针对性的图级优化,以减少归一化周围不必要的布局转换。在我们的实验中,除了 SME2 带来的增益外,这还为 INT8 带来了额外的约 70 毫秒 (23%) 延迟降低,为 FP16 带来了约 30 毫秒 (10%) 的延迟降低。

这些结果证实,转置密集型数据移动是一个有意义的优化机会,随着我们继续分析整个图的布局行为,可能会有进一步的改进。更广泛的发现及其影响将在后续文章中介绍。

洞察 3:在本案例研究中,FP16 的延迟接近 INT8

尽管 INT8 每个张量元素的内存带宽使用量是 FP16 的一半,但这并不会自动转化为成比例的端到端提速。在本案例研究中,启用 SME2 后,FP16 的延迟接近 INT8(单核上为 298 毫秒 vs 304 毫秒)。

算子细分揭示了原因。FP16 在卷积上看到了特别强劲的加速(9.0 倍,而 INT8 为 4.4 倍),这弥补了 INT8 在内存效率上的优势。同时,INT8 矩阵路径带有来自量化、缩放和更复杂内核调度逻辑的额外开销,降低了 INT8 的有效带宽优势。

最终结果是,SME2 扩大了可行精度选择的范围。INT8 仍然是一个有效的选择,而对于量化复杂性或精度权衡不可取的精度敏感型工作负载,FP16 变得更加实用。虽然在本案例研究中 FP16 的性能接近 INT8,但这种行为取决于工作负载,可能会随算子组合、张量形状和内存压力而变化。

动手实践:重现该工作流

要亲自尝试此工作流,我们提供了基于开源 SAM 模型的动手教程,其中演示了如何导出模型、使用 SME2 运行推理以及使用 ETDump 进行算子级分析。完整的设置说明和代码示例可在代码仓库学习路径中找到。

您将学到:

  • 如何将分割模型导出到启用了 XNNPACK 委托的 ExecuTorch
  • 如何构建模型并将其部署到支持 SME2 的 Android、iOS 和 macOS 设备上
  • 如何运行 ETDump 分析以收集每算子的计时信息
  • 如何识别和量化您自己模型中的数据移动和其他非数学瓶颈

结论:SME2 在实践中带来了什么改变

在此 SqueezeSAM 案例研究中,SME2 为 INT8 和 FP16 带来了显著的端侧 CPU 加速,实质性地改变了交互式移动工作负载的可行性。

这对开发者和产品团队意味着什么

  • 端侧机器学习在 CPU 上变得更加可行SME2 可实现高达 3.9 倍的端到端推理提速,在默认 Android 电源设置下的真实交互式移动模型上,将单核延迟从超过一秒降低到约 300 毫秒。使得基于 CPU 的端侧机器学习从边缘化变得对交互式工作负载切实可行,同时为应用的其余部分保留了余量
  • FP16 在某些情况下成为更可行的部署选项:通过大幅加速 FP16 并缩小与 INT8 的延迟差距,SME2 为开发人员提供了更大的灵活性,可以选择最能满足精度、工作流和延迟要求的精度,特别是对于精度敏感的工作负载。
  • 节省的计算余量促成了更丰富的体验:空闲的 CPU 预算可以重新投入到额外的端侧功能中,例如将分割与增强功能(如去噪或 HDR)结合使用,或将抠图从单张图像扩展到具有帧间主体跟踪的实时视频。
  • 分析揭示了下一个优化目标:一旦 SME2 加速了矩阵密集型算子(卷积/iGEMM 和 GEMM),瓶颈往往会转移到数据移动和未委托的算子上。使用 ETDump 进行算子级分析使这些成本变得可见且可操作。

两个具体建议,取决于您的起点:

  • 如果您目前没有部署端侧 ML,SME2 加速 CPU 可以使移动 CPU 部署成为数学密集型模型的可行第一步,通过分析提供验证性能和迭代的清晰路径。
  • 如果您已经在部署端侧模型,SME2 可以创造余量来扩展功能并改善用户体验,而分析则突出了影响力最大的下一次更改(对于 SqueezeSAM,转置驱动的布局转换约占总运行时的 40%)。

总之,SME2 加速和算子级分析为解锁直接的性能增益和识别端侧 AI 的高影响力优化提供了实用的工作流。

鸣谢

我们要感谢 Meta ExecuTorch 团队的 Bilgin Cagatay、Mergen Nachin、Digant Desai、Gregory Comer 和 Andrew Caples 对实际用例的指导以及对推理优化实现所做的贡献。我们还要感谢 Arm 的 Ray Hensberger、Ed Miller、Mary Bennion 和 Shantu Roy 在此项工作中给予的支持和指导。