图与硅片之间：深入 Apple Neural Engine 编译器

当我继续深入编译器的深渊时，有一个想法终于开始在我心中成形，编译器不仅仅是把语言A转化成机器B的翻译器。这个描述在刚开始学习时很有用，但它忽略了编译器有趣的大部分内容。

编译器也是一种管理信息的系统，同时逐步改变程序的形式。在每个阶段，它都必须决定哪些仍然重要，哪些可以转换，哪些最终可以丢弃。

这也是我在写LLVM时感兴趣的重要原因，也直接影响了KGEN和MLIR。LLVM促使我思考为什么存在中间表示。KGEN进一步复杂化了这一问题，展示了为什么单一表示可能并不总是足够。

Mojo可以在LIT中保留源级信息，同时保持信息仍然有用的信息;在参数信息仍然重要时进入KGEN，并在编译器完成推理前通过HLCF保持结构化控制流。

大部分讨论都围绕着编译器想要无论是保存还是改造。我的下一个项目，斯宾塞·布林格尔森的苹果神经引擎：架构、编程与性能引入了稍微不那么可谈判的东西，那就是机器本身。

苹果神经引擎是一个异常不透明的目标，鉴于这是苹果，这应该会让人感到意外没人.苹果通过Core ML公开公开了ANE，而编译器、程序格式、驱动、固件以及其底层的大部分机制仍未被记录。

Bryngelson的论文是对该栈的独立逆向工程，基于对苹果私有软件的直接测量和分析。在本文中，我主要关注编译器之后，神经网络图被转化为ANE实际可执行的过程。

考虑一个神经网络图，由卷积、偏置加法和ReLU激活组成。在图层面，这三种操作是完全合理的。先进行卷积，将偏置加入结果，然后应用激活。

当我刚开始阅读编译器章节时，我或多或少以为这三个操作中某个逐步低级别的版本最终会成为加速器的工作。

但这种假设很快就被打破了。

ANE编译器可以将偏置和激活吸收到卷积本身。填充也可以折叠进卷积中，量化路径的部分可以消失在同一后台操作中。卷积的偏置成为硬件增益偏移机制的一部分，而激活则可以占据与操作相关的激活槽之一。

编译器中进入的是什么conv → bias → relu因此，可以将此阶段作为一个融合的ANE操作，而非三个独立调度。

这与我从KGEN学到的内容直接联系起来，只是方向相反。在KGEN中，我总是想起为什么编译器会保留一个概念。这里我开始思考，什么时候两个概念之间的区分不再有用。

当我推理神经网络结构时，卷积、偏向和ReLU作为独立的概念仍然有用。一旦编译器到达已经有表达偏置和激活的硬件，保持这些边界可能就无法带来任何好处。

这两种表示方式都没有比另一种更真实。

现在，我们的四个大致阶段如下：融合、合法化、调度和任务描述符分区，然后是内存和DMA优化。用不那么编译器化的语言来说，这个顺序相当直观。

编译器首先要理解请求了什么计算，以及其中几个部分是否可以合并成一个机器操作。然后它要确定最终的操作是否在该硬件上合法，决定如何排序和划分剩余工作，最后决定数据应存放在哪里以及在所有这些过程中如何移动。

融合与边界的实用性

融合是一个很好的起点，因为它让表征与硬件之间的关系异常具体。

进入编译器该部分的网络表示仍可将卷积、偏置和激活作为独立操作。融合发生在ANE编译器自身图机制中。这一区分很重要，因为它表明原始边界并非错误。

它们在某个阶段有用，在另一个阶段则不再必要。

同样的逻辑也适用于其他操作。转置有时会消失在邻接工作中，或使用专用机器，而不是成为独立的派遣。缩放和批次归一化可以被吸收进同一条增益偏移路径。

一个去量化步骤可以折叠进卷积的权重路径，而不是作为独立操作存在。编译器不断询问中间结果或操作边界是否仍代表目标实际需要看到的内容。

这是有极限的。卷积结果与恒定偏置之间的加法自然符合增益偏移机制，因为操作的一侧是常数的。两个独立计算的卷积输出之间的加法则不然。

两个操作数都是活值，因此加法仍然是真正的后端操作。Concat可以创建另一个融合边界，注意引入更严格的段边界，后续容量约束甚至可能迫使编译器拆分之前融合过的内容。

我喜欢这种细微差别，因为它避免了融合变成另一种模糊的“编译器优化让事情更快”的解释。编译器并不是因为操作越少就合并操作。当硬件能够表达它们的组合行为时，编译器才会将它们合并;当硬件无法表达它们的行为时，编译器才会将它们分开。

当有效的计算不是合法机器程序时

中间表示可以描述一个完全有效的计算。数学可以合理，图可以是良好构造的，编译器前端也能准确理解你的需求。但这些都不能保证这台机器能以那种形式执行计算。

ANE有一个包络线。张量秩是有限的，单个维度有边界，卷积核有目标特定的约束，系数必须符合特定的内存限制。如果操作超出这些边界，编译器必须对其进行铺砌、拆分、转换，或将图拆分为额外的段。

这听起来很直接，但我认为它暴露了一个区别，当人们说某件硬件“支持”某种操作时，很容易被忽视。可表示、有效、可降低和可执行不一定是同一个主张。

现在，编译器暴露了验证器，用于检查操作是否满足预期的形状、类型、维度、操作数计数和目标特定能力。然而，通过其中一个验证器并不一定意味着后端代码生成一定会成功。

在M1上，论文发现诸如top-k、sort、动态切片和三维卷积等操作可以通过部分验证表面，但在后端降低过程中仍然会失败。

这意味着“ANE支持这个”。但是，支持在哪里？

中间语言能表示吗？前端接受吗？验证者批准吗？后端能把它转成 ANE 操作吗？最终程序真的能在这一代硅片上执行吗？

这些是不同的界限，我们越往下看，这个词的用处就越小支持变成了无需说明我们指的是哪一种。

当矩阵乘法决定变成卷积时

我最喜欢的例子之一是矩阵乘法，因为它让我一直带着的另一个编译本能崩溃。

如果矩阵乘法的右手权重符合ANE的片上工作集预算，编译器可以将该矩阵乘法重写为常驻卷积。如果权重过大，则仍为拼贴矩阵乘法。在M1上，工作集阈值约为2 MB。

第一次读到这条时，我脑海里还有一部分想反驳我要求的是矩阵乘法，而不是卷积。

编译器负责保持计算过程，而不是我最初给该操作起的名字。“矩阵乘法”可能是理解数学的最佳表达方式，而“驻留卷积”可能是在这种加速器上更合适实现，因为权重在芯片上非常合适。

这时降低的感觉不再像翻译，更像是在可能实现相同意图的实现中选择。源代码级操作并非神圣。重要的是转换是否保留程序应计算的内容，同时更好地利用底层机器。

这台机器不仅开始影响怎么会这样执行一个操作，但该操作的哪个表示有价值。

最终，字节开始做出编译器决策

一旦图被融合并合法化，编译器仍需确定最终工作应如何实际运行。第22章的调度阶段线性化操作图，并将调度划分为符合引擎芯片内工作集预算的任务描述符划分。

M1让这一点非常具体，因为调度器正对着大约2MB的静态内存上限工作。编译器没有按任意数量的操作来划分工作，而是小心地增加了一层，估算其生命周期内的峰值存活内存压力，并在下一项工作超出允许预算时关闭当前分区。

这可能是章节中让我最清楚地理解的部分。

在编译器更高处，我可以整天讨论表示、图结构、融合规则和操作。最终，程序划分到哪里可以变成更不哲学化的做法，即实时数据是否仍然适用？

此时，信息不再只是编译器推理的抽象。它占据字节，而这些字节需要存在。

这个认识又稍微改变了我对编译器的定义。它确实在管理信息，但随着我们越来越接近机器，它也越来越像一个资源管理系统。编译器必须考虑可用的快速存储量，哪些值同时处于活跃状态，以及所选执行顺序是否能将工作集保持在物理限制内。

两个运算之间的箭头不是自由的

DMA，即直接内存访问，本质上是允许数据在内存和加速器之间移动的机制之一，而无需CPU亲自管理每个字节。

在高层次上，我可能会画出两个操作：A → B主要思考依赖性：B需要A产出的东西。

靠近硬件时，箭头开始获取负载。A把结果放哪里了？需要重新写入动态内存吗？B需要重新读取吗？结果能留在芯片上吗？这两个操作是否安排在一起？

它们的瓦片几何形状是否对齐？还有人需要这个值才能重新使用吗？

第22章描述了几种分配状态，区分了常驻芯片内的数据、从动态内存流式数据、存储在环形缓冲区、原地重写或直接从一个操作链到下一个操作的数据。

在合适的条件下，拥有一个消费者的生产者可以将结果留在芯片上，使后续操作直接消耗数据，而不必不必要地往返动态内存。

如果你完全生活在图表层面，这个细节很容易被忽略。箭头在那里看起来是空的。在这里，它变成了内存流量。

编译器还会重新排列目的缓冲步幅，以避免银行冲突、折叠填充和权重流。这些决策不如发明红外线或证明优化正确那么光鲜，但它们揭示了加速器编程的一个根本性：一旦计算成本足够低，移动计算所需的数据可能成为问题的重要部分。

同样，抽象也随问题变化而变化。神经网络图关注操作之间的依赖关系。内存优化器必须关注当满足这些依赖关系时，生成的字节物理上可用的位置。

代表性是有代价的

在我处理这些问题的过程中，Chris Lattner为我在KGEN文章中试图表达的内容增添了有用的补充。我对MLIR能够让不同表示法保留不同阶段重要抽象的能力产生了浓厚兴趣。

他的观点是，每增加一个国际关系也会增加复杂性，而这种复杂性必须有其合理性。目标不是尽可能多地保留表示，而是使用最少的表示，同时仍能表达每个阶段所需的核心抽象。

这和我在ANE编译器中看到的情况出乎意料地吻合。

网络图在区分时将卷积、偏置和激活分开。当硬件能够共同表达它们的行为时，融合表示会放弃它们。合法化保留了足够的形状和能力信息，以判断最终工作是否属于机器的包管内。

调度器越来越关注活值范围和内存压力，而分配者最终关注的是驻留率、步幅、瓦片和移动。

编译器保留抽象并不是因为有抽象是好事。它保留抽象是因为在抽象存在时，还有一个更容易回答的问题。

一旦问题得到解答，继续将这种代表性推迟可能变得多余。

这部分逐渐改变了我对降低的看法。我过去主要把它想象成从高层次向低层次的移动。我仍然认为这在方向上是正确的，但这忽略了程序为什么不断变化形状。

当表征对物理学作出回应时

回顾LLVM、KGEN，现在还有ANE编译器，我觉得我一直在不小心跟着同样的问题往下走。

对于LLVM，我想弄明白为什么中间表示如此重要。一个有用的IR保留了足够的结构，使编译器能够继续推理程序，而不是立刻接受某台机器的局限。

KGEN 和 MLIR 扩展了这一理念。编译的不同阶段可能需要不同的有用抽象，编译器可以根据所需信息的变化在这些阶段间切换。有趣的问题在于决定哪些仍值得存续，以及何时可以安全放手。

ANE引入了另一个约束，因为这些表示最终会遇到具有有限内存、特定执行单元、支持和不支持表单、目标特定限制以及实际数据移动成本的物理机器。

卷积、偏置和ReLU可能崩溃，因为该加速器拥有能够表达它们组合行为的机制。矩阵乘法可以成为卷积，因为该实现在特定内存约束下恰好更适合目标。

图可以被划分到与程序员组织方式无关的位置，因为下一个活跃张量会推动工作集超出加速器能承受的范围。中间值可以保持常驻或流式，取决于周围执行是否允许。

计算仍然必须在对我们最终重要的意义上含义相同：在相同的有效输入下，它应该会产生我们所要求的结果。但几乎所有关于该计算表现方式的表现都可能在下降过程中发生变化。

它们看起来越来越不像机械翻译一种语言到另一种语言的机器，更像是多个同一计算描述相互协商的地方。程序员有一个描述。模型图有另一个。

编译器创造更多。最终，加速器引入了这些之前的表示无法忽视的约束。

从网络描述到可加载的ANE程序结束。中间发生的是融合、重写、调度、分区、分配和数据移动，但在所有这些机制之下，是一个相当直接的进程。

编译器不断修改程序，直到我们描述计算的方式与机器物理执行方式兼容。