优化基于MLIR的Torq NPU编译器

优化基于MLIR的Torq NPU编译器：ENERZAi如何降低SLICE并行带来的DMA开销

本文介绍了ENERZAi如何在Synaptics Astra平台的Torq NPU上，针对基于IREE/MLIR的编译器，降低因SLICE并行化而产生的DMA开销。

本文最初发布于ENERZAi官方网站，经ENERZAi授权在此转载。

大家好，这里是ENERZAi。我们最近为Torq编译器的后端实现和模型优化做出了贡献，Torq编译器是一款开源的基于IREE/MLIR的编译器，服务于全球半导体公司Synaptics旗下Astra SL2610产品线中的Torq NPU。在本文中，我们将分享这一过程中遇到的挑战以及获得的经验。

具体来说，我们将深入探讨当一个运算被拆分到NPU的两个计算模块（SLICE）上并行执行时所产生的隐藏数据搬移（DMA）开销，分析我们在编译器的哪些层级中降低了这一开销，以及具体的实现方式。

再次感谢Synaptics为我们提供这次合作机会！

一个仍在成形中的NPU生态系统

过去几年，业界对专用领域架构的关注迅速增长。CPU和GPU已经演化出区分通用计算与大规模并行工作负载的能力。深度学习催生了类似的专业化需求：模型的大部分计算可以表示为张量运算的流程，这使得针对此类工作负载量身定制的架构变得愈发重要。基于模型的服务的快速增长正在加速这一趋势。

因此，许多公司正在开发并推出专门用于神经网络工作负载的处理单元：NPU，即神经网络处理单元。它们在边缘计算中尤为重要，因为模型需要在严格的功耗和内存预算下直接在设备端运行，同时仍要提供实用的智能能力。

然而，与CPU和GPU不同，NPU尚未拥有成熟的编译器栈、运行时以及建立在其之上的完善生态系统。各厂商的架构差异很大，与GPU相比商业化部署的时间仍然较短，业界对于硬件与软件边界应设在何处也尚未形成广泛共识。因此，各平台的API和基础设施存在显著差异。在某些情况下，NPU可能已经投入量产，而能够有效使用它的软件栈却仍在追赶进度。

看似显而易见的优化

多计算单元的NPU所暗示的第一个想法很简单：拆分工作并并行运行各个部分。Astra SL2619上的Torq NPU拥有两个SLICE（计算模块），因此乍看之下，将一个运算拆分为两半并分别在每个SLICE上运行似乎是显而易见的优化方案。

只需了解几个硬件术语，就能明白为何结果并不总是那么简单：

SLICE程序对已经放置在LRAM（一种小容量的片上内存）中的数据进行操作。

XRAM是在板载主DRAM中为NPU执行预留的区域。

数据通过DMA在XRAM和LRAM之间移动。输入数据需要先到达LRAM才能被SLICE使用，而结果在需要时会被复制回XRAM。

拆分运算可以创造更多的并行执行机会，但也可能产生更多的数据搬移。拆分决策是在编译器仍处于张量级中间表示（IR）阶段时做出的，此时具体的DMA任务和SLICE程序尚未生成。一旦这些部分被降级处理，每一部分可能会各自获得独立的加载、存储以及程序代码传输。如果两个部分需要相同的常量或输入数据，这些数据可能会被多次搬移。

还有一个微妙之处：拥有两个SLICE任务并不意味着两个SLICE会真正地重叠执行。当前的调度器依据IR中启动（start）和等待（wait）操作的顺序来分配任务。如果希望两者并发执行，两个程序都必须在对应的等待操作之前启动。

当像scf.forall这样的MLIR并行循环被展开为两个独立的部分时，就能生成第一种调度方式。然而，仅仅并排创建两个操作，仍可能产生第二种形式：如果每个操作后面紧跟着自己的等待操作，那么第二个SLICE要等到第一个完成后才会开始执行。

数据搬移又增加了一层复杂性。DMA可以与SLICE执行重叠，因此额外的流量不会自动转化为串行延迟。但这仍然给调度增加了更多需要吸纳的工作量，任何无法重叠的部分都会延长关键路径。

这成为我们在Torq项目中反复出现的一种模式：一项在张量IR层面看起来有益的优化，往往只有到达编译器更低层级时才会暴露出新的成本。并行化本身并不是难点，真正的难点在于让这种并行化所带来的成本足够早地显现出来，以便我们能够采取应对措施。

这一观察引出了一个编译器设计层面的问题：当一项决策在某一层级做出，而其代价却在另一层级显现时，每项优化应该放在哪一层？

Torq为何使用多个编译器层级

编译器通过原始模型与最终硬件代码之间一系列中间表示（IR）来描述模型。不同的优化需要从这些表示中获取不同的信息。

一些转换对许多目标平台都有用，而另一些则取决于目标平台的数据流、支持的运算模式或内存层次结构。这两类优化不能简单地割裂开来互不理会，因为它们的决策仍会相互影响。运算融合就是一个很好的例子：合并运算可以避免生成中间结果，但融合后的表达式可能无法适配NPU所提供的有限专用运算集合，而为了适配所付出的代价可能会超过融合所节省的成本。平铺（Tiling）也存在同样的张力。它可以在张量层级表达，但最终有效的平铺形状仍受硬件内存和分配策略的限制。

Torq通过在多个编译器层级之间划分关注点来处理这一问题。在MLIR中，这些层级以方言（dialect）的形式组织，每种方言代表特定抽象层级上的运算和概念。Torq尽可能长时间地将张量语义和高层转换保留在标准方言中。只有当涉及NPU程序和内存的具体信息变得必要时，Torq专用的方言才会出现。

Torq编译器以IREE插件的形式实现，并大量使用了MLIR。其主要层级包括：

linalg层：这是MLIR的内置方言，主要用于以数学语义表达张量运算。平铺和融合等张量级转换就在这一IR上进行。某些编译passes可能会参考目标平台信息，但此时硬件专用的程序描述符和内部数据路径尚未暴露。

torq_hl层：这是一种赋予运算Torq专用语义的方言，用于表示XRAM/LRAM传输以及SLICE程序的执行。在这一层级，哪些数据需要移入或移出NPU内存、以及何时启动NPU程序，都变得明确可见。

torq_hw层：这是一种映射到硬件所消费的描述符的方言。D/W/B/Q数据路径以及ALU/ACT配置等信息，会直接出现在定义实际SLICE任务的IR中。

在每个层级上观察同一个矩阵乘法运算，可以让这一渐进过程变得具体：

每个层级上的算术运算本身保持不变，但表示方式越来越具体。linalg承载张量级语义。torq_hl添加了将该运算作为NPU程序执行所需的信息。到达torq_hw层级时，D（输入数据）、W（权重）、B（偏置/缩放）、Q（输出）以及ALU/ACT配置，已经展现出接近硬件实际行为的样貌。

一旦周边的程序和数据搬移变得可见，底层代价就变得更加清晰：

在linalg中原本是单一的张量运算，现在已扩展为查找XRAM地址、分配LRAM、加载输入、启动并等待程序运行，最后将结果存回。算术运算本身没有改变，但为使该计算能在设备上执行所需做的准备工作的成本，现在已经可见。

这正是层级划分对我们工作至关重要的原因。在linalg层级做出的一次拆分，到了torq_hl层级可能会演变为额外的地址、DMA和程序启动。同时，一些问题最容易在这些细节出现之前就解决，而另一些问题则恰恰需要后续才会出现的信息。与其将每一项优化都塞进同一个层级，我们可以在相关信息变得可用的那个层级去应对对应的代价。

同样的结构也为编译器的持续演进留出了空间。NPU的运算方式、内存容量和数据路径可能会随着硬件代际的更新而变化。将高层语义与底层硬件细节分离开来，可以限制整个流水线随目标平台变化而需要调整的范围。

为切分所付出的代价买单

一旦我们能够看清额外的数据搬移出现在何处，我们的大部分工作就围绕一个主题展开：保留并行性，但尽可能去除它所带来的额外搬移。这些修复分布在不同的抽象层级，因为每一项修复都需要对程序有不同的视角。

在linalg层级，我们发现了一些连续的逐元素运算被逐个拆分的情况。复合激活函数和后处理运算经常被规范化为一系列linalg.generic运算的链条。如果每个运算都被放入独立的循环，中间值就可能跨越循环边界，之后又在XRAM中被实例化。

我们统一了附加在这类运算链条上的torq-fuse-group属性。编译器会将同一分组内的linalg运算合并进单一的scf.forall循环，而不同分组的运算则会被放入独立的循环中。通过在拆分开始之前就将整条链条分配到同一分组，编译器可以从一开始就为整条链构建一个循环，而不是事后再尝试合并循环。下面的示例为了便于阅读，将较长的链条压缩为三个运算。

在出现此问题的降级路径中，跨越不同scf.forall循环的中间值会在XRAM中被实例化，从而产生额外的数据往返。将整条链条保持在同一循环内，可以让这些中间值在LRAM中停留更长时间，减少相关的DMA进出流量。这里的重点与其说是"为融合而融合"，不如说是为了在整条链条中保持数据的片上驻留。

启用切分后erf运算（1x1024x64, bf16）的编译期时间线。每个面板中，上方一组是将逐元素运算链放入同一融合分组之前，下方一组是之后；底部面板放大了蓝色方框区域，即每次运行的前330微秒。①链条中各运算之间向XRAM的存储：每个运算都将两个SLICE半部分都写出。②紧接着将这些值重新加载回LRAM以供下一运算使用的加载操作。③将整条链条放入同一融合分组后，中间值保留在LRAM中，SLICE程序可以连续背靠背运行。存储次数从36次（2.6MB）降至7次（448KB），预估耗时从1834.9微秒降至888.0微秒。

另一个问题出现在编译器已经到达torq_hl层级之后。某些值仍以张量形式表示，但它们的内存位置此时已编码在其类型中。在原始IR中，编译器会先将完整的值转换到XRAM，然后才提取实际所需的较小SLICE部分。

我们将其重写为一种等价形式：在值仍处于LRAM中时就提取出SLICE部分，然后只转换这个较小的SLICE。这样可以让大张量保留在片上，同时在较小的SLICE上暴露出一次冗余的数据往返，而现有的编译器优化恰好可以将其消除。

max_pool运算（5×5, 1x15x20x384, i8）在将SLICE提取操作移到LRAM转XRAM操作之前（上方一组）和之后（下方一组）的对比。①完整的112.5KB结果被转换到XRAM。②每个SLICE随后提取出自己的56.2KB半部分并重新加载回LRAM。③改动之后，两个SLICE在值仍处于LRAM时就各自提取出自己的半部分，因此这次数据往返就此消失。预估耗时从369.3微秒降至312.9微秒。

这项改动与上述linalg分组改动都没有引入任何新的运算。在这两个案例中，有用的转换其实已经存在于编译器某处；关键在于如何安排IR，使其能够针对Torq的内存模型做出正确的处理。

在torq_hl层级，LRAM地址、加载/存储以及程序边界都已明确，因此可以直接应对某些冗余的数据搬移。如果同一份只读数据被分别加载给两个SLICE，我们希望只传输一次并共享使用。如果某个中间值被写入临时XRAM但之后永远不会再被读取，我们希望能移除那次存储操作。

有趣的是，处理这些情况所需的许多分析和转换其实早已存在于编译器中。问题在于torq_hl是一种自定义方言，因此编译器并不能自动了解其运算如何访问内存，或者是否可以安全地移动这些运算。我们没有重新实现这些分析，而是通过编译器接口暴露了这些属性。

SLICE程序参数的读/写方向通过AsyncAccessOpInterface描述，这是IREE中所用接口的一个简化版本。对于运算级别的内存效应和移动合法性，我们使用了MemoryEffectsOpInterface和ConditionallySpeculatable。一旦这些属性可用，现有的MLIR分析工具和实用程序也能够对torq_hl进行推理。

一个立竿见影的好处是消除死存储（dead store）。在没有读写信息的情况下，编译器必须假设每个SLICE程序参数都可能被读取和写入。一旦编译器知道后续程序只是覆盖写入某个值而非读取它，之前的溢出存储操作就可以被消除。

同样的信息在某个缓冲区因LRAM空间不足而被临时溢出到XRAM、随后又被重新加载时也很有帮助。如果中间的程序只读取该缓冲区，那么之前写入的XRAM副本仍然有效。编译器可以释放该LRAM分配并重用这份副本，而不必再分配一个新的XRAM位置并重复存储同一份数据。

一个conv1d层（卷积核127，步长64，8万个音频采样点转288通道）在使用AsyncAccessOpInterface之前（上方一组）和之后（下方一组）的对比。底部面板将两次运行中相同的循环迭代对齐（蓝色方框：改动前为1950-2720微秒，改动后为1842-2612微秒），并将各自起点归零对齐。①一次死存储：一个200KB的值被溢出到XRAM，但下一个程序只会覆盖写入它。②与之对应的200KB加载操作，也随之消失。③一个未改变的309KB缓冲区被再次存储到XRAM；改动之后，XRAM中已有的副本被直接重用。DMA输出次数从22次（4.5MB）降至13次（2.4MB），预估耗时从5091.1微秒降至4218.8微秒。

一个相关的问题是某个运算是否可以被移动。如果一个运算在每次循环迭代中都产生相同的值，那么最好在循环外执行一次并重用结果。此前，MLIR无法判断移动一个torq_hl运算是否合法。

我们只标记了那些实际上可以安全移动的运算；而那些含义依赖于其所处位置的运算，例如按索引收集数值的运算，则被排除在外。这使我们能够重用MLIR的getBackwardSlice和moveLoopInvariantCode，而不必编写单独的合法性判断逻辑。

综合来看，这些改动消除了部分因切分而引入的重复数据搬移。它们还改变了我们对这一问题的思考方式：仅仅把工作拆分好是不够的。如果并行化会产生额外开销，编译器就需要有办法在这种并行化真正转化为性能提升之前，先为这份开销买单。

我们学到了什么

回顾整个过程，具体的重写和接口本身的重要性，反而不及其背后所体现的模式。

首先，优化始于理解数据在芯片中如何流动。在Torq上，每个运算都有特定的可用输入路径、累加路径和后处理路径，这些约束决定了什么能够放入单一程序中。任务数量和DMA字节数是有用的诊断指标，但它们本身并不能作为性能的最终判定依据。即使融合减少了任务数量，我们仍需要追问融合后的运算选择了哪条内核路径，以及实际关键路径会发生怎样的变化。

其次，在受限的LRAM预算下，平铺、融合和切分三者紧密耦合。这三者都会影响哪些数据留在片上、留存多少、以及多久之后又需要再次移动。如果这些决策被独立做出，很容易导致各项优化相互掣肘。

答案并不是把所有东西都压缩进一个硬件专用的IR中。分层流水线并不能消除各项优化之间的耦合，但它为这种耦合提供了一个能够显现出来的场所。高层转换可以继续使用张量语义进行工作，而底层则暴露出这些决策最终产生的硬件代价。对我们而言，这种分层划分的意义，与其说是出于架构整洁性的考虑，不如说是为了在正确的时间获得正确的信息。

下一步计划

同样的经验也塑造了我们对编译器未来发展方向的思考。我们的目标是支持更广泛的模型和未来的硬件，而不是让每一个新的目标平台都变成又一堆硬编码的特例。

从EasyKernel到Nadya。目前，Torq的内核是用EasyKernel编写的，这是一种内部的C++领域专用语言（DSL）。内核可以用更接近普通循环的方式表达，然后再序列化为描述符，而不必直接编写底层的硬件描述符和指令流。这使得开发工作更加轻松，但每个张量运算的具体实现仍然存在于编译器内部。这使得支持新运算和未来硬件代际变得更加困难，也限制了对内核实现进行运行时性能分析和调优的机会。

我们正在探索的一个方向是将内核编写工作迁移到Nadya上，这是我们正在开发的一种内核编程语言。Nadya的设计目标是让开发者能够编写高性能程序，而不必要求每位作者都理解所有底层硬件优化细节。NPU执行某个运算的具体实现方式可以存在于外部内核代码中，而编译器则负责决定多个运算应如何协同调度。

让优化决策变得可搜索。整个编译流水线中包含许多由启发式方法驱动的决策：平铺尺寸、SLICE划分方式、是否进行融合等等。我们希望将这些选择参数化，而不是将其硬编码到各个独立的编译passes中，从而使其能够被调优循环系统地探索。这个调优循环可能存在于外部工具中，也可能随时间推移与编译器更紧密地集成；关键在于将搜索策略与转换逻辑分离开来。这使得同一套编译流水线能够评估不同的配置方案，并适应新硬件，而不必为每一个新目标重新编写启发式规则。

要让一款面向全新目标类型的编译器既强大又灵活，仍然存在许多待解决的问题。但这恰恰正是这个项目有趣之处。在ENERZAi，我们仍在一个一个地攻克这些难题。

一个开放的项目

这个项目最棒的一点在于，硬件规格和编译器都是公开的。编译器代码可在synaptics-torq/torq-compiler的Github仓库中获取，模型示例可在torq-examples中找到，官方文档位于synaptics-torq.github.io。任何人都可以查看我们正在处理的各类问题，感兴趣的贡献者也可以直接参与进来。

如果您想进一步了解Torq项目或ENERZAi的其他优化技术（如Nadya），欢迎随时与我们联系！

Jaeho Kim

ENERZAi