回顾性逆向工程苹果神经引擎
原题Retrospectively Reverse-Engineering Apple's Neural Engine
推荐理由三年逆向还原 M1 苹果神经引擎:计算核、激活、调度与任务描述符一层层摊开,并点明 M5 把 ANE 并入 GPU、独立 NPU 时代收束。端侧 AI 到底跑在哪、旧假设如何过期,硬件与系统工程师少见的完整底图。
读前速览
作者对 M1 芯片上的苹果神经引擎进行逆向工程,详细解析其计算、激活、调度器和任务描述符设计,揭示 CNN 时代的假设以及向 GPU 集成 AI 工作负载的转变。
计算核心架构与数据流设计
ANE 包含 16 个并行计算核心,每个核心有 256 条 MAC 通道,排列在乘法器-加法器-累加器数据通路中。累加器在 2^15 处饱和,数据流针对密集 CNN 层典型的可预测复用模式进行了优化,不过当数据适当流式传输时,同一 MAC 阵列仍可支持 Transformer 工作负载。
通过查找表实现非线性激活
激活函数在 MAC 之后应用,无需内存往返。tanh 通过 33 项分段线性 LUT 和线性插值实现,ReLU 和恒等使用模式 0/1 且无表,线性缩放/偏置折叠到卷积权重中,共享同一插值硬件。
调度器与任务队列机制
驱动程序通过将任务描述符流写入内存并触发 TM_PUSH 来提交工作,TM_PUSH 选择八个任务队列之一。每个队列持有两个乒乓描述符槽和一个 32 项的 BAR 表,提供基 IOVA 地址,允许 ANE 在编译时通过固定偏移解析虚拟地址。
任务描述符格式与示例
每个固定大小的任务描述符由头部、内核 DMA 配置、公共字段(格式、维度、分组)和瓦片源/目标块组成。一个 1x1 卷积 TD 的十六进制转储展示了 EON、LogEvents、DebugEvents、RBase、WBase、KBase、tileH、ActiveNE 和地址寄存器等字段如何编码操作。
ANE 实用性的影响与未来集成
作者得出结论,ANE 固执的数据流将其限制在具有可预测复用的工作负载上,使其作为通用加速器不太有用。M5 芯片中将 ANE 核心折叠到 GPU 中,并以 LLM 性能为卖点,标志着从独立 NPU 向以 GPU 为中心的 AI 处理的转变。
完整摘要
文章详细介绍了对 M1 芯片中苹果神经引擎(ANE)的回顾性逆向工程。首先讲述了作者三年后的动机,然后逐步剖析 ANE 的计算核心:16 个核心,每个核心有 256 条 MAC 通道,数据通路,累加器在 2^15 处饱和,以及数据流如何针对 CNN 工作负载中可预测的复用模式进行调优。接着描述了非线性激活模块,展示了 33 项分段线性 LUT 如何实现 tanh,而 ReLU 和恒等模式共享同一硬件,以及线性缩放和偏置如何折叠到卷积权重中。调度器部分解释了任务管理器通过 TM_PUSH 提交命令、乒乓任务队列,以及为固定大小描述符提供基 IOVA 地址的 BAR 表。任务描述符格式通过一个 1x1 卷积 TD 的十六进制转储进行剖析,分解了头部字段、内核 DMA、公共字段、瓦片源/目标,以及通过 BAR 偏移解析虚拟地址的方式。最后,作者反思了 ANE 作为通用加速器的有限实用性,指出其固执的架构限制了工作负载多样性,以及 M5 将 ANE 核心集成到 GPU 中标志着独立 NPU 时代的终结,反映了从 CNN 时代的 NPU 向 GPU 处理基于 Transformer 模型的转变。
目录 · 15 节
https://github.com/eiln/ane/tree/main
M1 die shot: https://mastodon.social/@dougall/115149886886125067
M5(2025)的主打卖点是“LLM 性能”,而且他们还顺手把 ANE 核心整合进了 GPU 核心里——我早知道这事迟早会发生,但它正式让我觉得,独立 NPU 的末日已经开始了。所以,为了纪念 ANE 似乎就此退场,我们要做件更加没用的事:回头对 M1 上的 ANE 进行逆向工程,把我们当初没做完的事做完。已经三年了(操),而我现在应该比刚开始做这件事时懂得更多。The M5 (2025)'s headline feature was "LLM performance", and they also conveniently folded the ANE cores inside the GPU cores — I knew it was coming, but it officially feels like the beginning of the end for the standalone NPU. So, in honor of the ANE’s apparent demise, we will do something even more useless: go back and reverse-engineer the ANE on the M1, finish what we started. It's been three years (fuck), and I should know more than I did when I first worked on this.
如果说三年前的目标是通过在 ANE 上运行运算,让它变得实用;那么这一次,重点更多在于梳理其完整的内部架构——计算单元、数据通路、调度器、内存和执行模型——因为这些内部设计决策揭示了苹果在 A11 Bionic(2017)中率先将哪些关于机器学习工作负载的假设固化进芯片,以及这说明了什么:从 CNN 时代的 NPU 到如今运行 Transformer 工作负载的 GPU,期间发生了怎样的转变。If the goal three years ago was to make the ANE useful by running ops on it; this time, it's more about mapping the full internal architecture — compute, datapath, scheduler, memory, and execution model — because those internal design decisions reveal the assumptions about ML workloads that Apple was willing to commit to silicon first in the A11 Bionic (2017), and what that says about the shift from CNN-era NPUs to today's GPUs running transformer workloads.
1. 计算1. Compute
16 个计算核心可能是 ANE 中最不值得关注的部分。苹果最初瞄准的是密集型图像处理 CNN 工作负载,这类工作负载由具有可预测重用性的密集张量归约构成。M1 ANE 的计算核心是由大量乘加(MAC)单元组成的大型并行阵列,但仅凭这一点,几乎无法说明它是为哪些工作负载设计的,也无法说明它擅长处理什么。The 16 compute cores are probably the least interesting part of the ANE. Apple originally targeted dense image-processing CNN workloads, which consists of dense tensor reductions with predictable reuse. The M1 ANE compute core is a large parallel array of multiply-accumulate (MAC) units, but that alone says almost nothing about what workloads it was designed for and accels at.
卷积层对一个激活窗口和学习得到的卷积核权重做点积,注意力则对查询向量和键向量做点积。点积就是点积,而 MAC 做的正是这件事。让 ANE 专门适用于 2017 年的 CNN 模型的,并不是 MAC 本身,而是围绕 MAC 的数据流:MAC 的输入和输出何时、在何处进入、停留和移动。Transformer 打破的假设,尤其是在自回归解码中被打破的假设,是数据会按可预测的模式重复使用;ANE 正是利用了这一点,设计出足够高效的数据流,使其能在手机上运行。M5 的决策表明,ANE 的计算核心对 Transformer 依然有用,只是要置于不同的数据流之中。A convolutional layer does a dot product between an activation window and learned kernel weights, and attention does a dot product between a query and key vector. A dot product is a dot product, and a MAC does just that. What specialized ANE to the 2017 CNN models is not the MAC, but dataflow surrounding the MACs: when and where MAC inputs and outputs enter, stay, move. The assumption that transformers broke, especially with autoregressive decode, was predictable reuse patterns, which the ANE exploited to architect a dataflow efficient enough to run on phones. The M5 decision confirms that ANE's compute core remained still useful for transformers, but inside a different dataflow.
不过,以下是 16 个计算核心各自内部的数据通路:Still, here's the datapath inside each of the 16 compute cores:
┌────────────────────── core ─────────────────────┐
│ ┌───────── 256× MACs ─────────┐ ┌────────────┐ │
│ │ MAD ─► add ─► accumulator │─►│ activation │ │
│ │ ▲ │ │ └────────────┘ │
│ │ └──────────┘ │ │
│ └─────────────────────────────┘ │
└─────────────────────────────────────────────────┘
乘累加Multiply-Accumulate
ANE 有 16 个并行计算核心。每个计算核心有 128 条 FP16(或 256 条 INT8)并行乘加(MAC)运算通道。每条 MAC 运算通道执行以下递推关系:ANE has 16 parallel compute cores. Each compute core has 128 FP16 (or 256 INT8) parallel multiply-accumulate (MAC) lanes. Each MAC lane performs the recurrence:
\[ s\leftarrow s+a\times b \]将两个操作数 \(a\) 和 \(b\) 相乘,然后将乘积加到当前累加和(累加器)中。Multiply two operands \(a\) and \(b\), and then add the product to the running sum (accumulator).
在 T 个周期内重复 MAC 运算,即可计算包含 T 项的点积:Repeating the MAC operation over T cycles computes a T-term dot product:
\[ s_T=s_0 + \sum_{t=0}^{T-1} a_t \, b_t. \]因此,一个 MAC 通道会随时间执行标量归约。16 核 ANE 拥有 2048 条并行 MAC 通道,A MAC lane thus performs a scalar reduction over time. A 16-core ANE has 2048 parallel MAC lanes,
\[ 128\ \text{lanes/core}\times16\ \text{cores} = 2048\ \text{parallel MAC lanes} \]因此,每个周期都会在空间维度上并行执行 2048 次归约,时间是唯一的归约轴:So each cycle performs 2048 parallel reductions spatially, with time being the only reduction axis:
\[ S_T[q,p] = S_0[q,p] + \sum_{t=0}^{T-1} a_t[q,p]\,b_t[q]. \]单个 MAC 通道并不知道自己正在沿矩阵或张量的哪个维度进行归约。需要注意的是,点积、矩阵乘法和卷积之间的区别,取决于操作数如何映射到核心上并进行调度。ANE 核心(稍后讨论的内核存储器除外)不会将 4 通道 CNN 层编码到硬件中。An individual MAC lane does not know what dimension of the matrix or tensor it is reducing over. It's important to note that a dot product vs matrix multiplication vs convolution arises from how the operands are mapped and scheduled onto the core. The ANE core (with the exception of kernel memory, discussed later) does not encode a 4-channel CNN layer into the hardware.
在内部,MAC 数据通路由一个乘法器、一个加法器和一个 32 位累加器寄存器组成。每个周期,加法器都会将乘法器刚产生的输出与前一次的累加结果相加,并将结果作为新的累加值。Internally, the MAC datapath consists of a multiplier, adder, and a 32-bit accumulator register. Each cycle, the adder adds the fresh multiplier output with the previous sum, which then becomes the new running sum.
operand a ──┐ ┌────────────┐ p[31:0] ┌──────────────┐ s_next[31:0] ┌─────────────┐
├──►│ MULTIPLIER │────────────►│ 32-BIT ADDER │────────────────►│ ACCUMULATOR │
operand b ──┘ └────────────┘ └──────▲───────┘ └──────┬──────┘
│ │ s[31:0]
└────────────────────────────────┘
该反馈路径将部分和保存在 MAC 通道本地的存储器中,因此在 MAC 周期之间无需从远处的外部存储器中读取。This feedback path keeps the partial sum in memory local to the MAC lane, so it does not need fetched from an external memory far away, between MAC cycles.
关于分辨率,它会在读出时使用 FP16 进行定点归约。乘数为 16 位,在 32 位寄存器中以 Q16.16 格式累加,然后通过符号扩展等方式读出为 FP16。若要以整数(十六进制)FP16 表示形式进行操作并探测累加器的范围,可以构建一个 CoreML ANE 程序,使用全为 (1) 的向量计算点积:这样每个乘数都会产生一个有界的 v,但累加器中的运行总和会持续增长:Regarding resolution, it does fixed-point reduction with FP16 at readout. The multiplier is 16-bit, accumulated in a 32-bit register as Q16.16, then read out as FP16 via sign-extend and etc. Working in integer (hex) FP16 representation, to probe the accumulator range, build a CoreML ANE program that computes a dot product with a vector of all (1)s, so each multiplier results in a bounded v, but the running sum in the accumulator keeps growing:
\[ s=\sum_{i=0}^{255}v=256v. \]| (v) | CPU hex | CPU value | ANE hex | CoreML value |
|---|---|---|---|---|
127.9375 |
0x77ff |
32752 |
0x77ff |
32752 |
128 |
0x7800 |
32768 |
0x7c00 |
+∞ |
−128 |
0xf800 |
−32768 |
0xf800 |
−32768 |
−128.125 |
0xf801 |
−32800 |
0xfc00 |
−∞ |
由于 32768 本身就是一个有效的 FP16 字(0x7800),ANE 的 0x7c00 不可能是 FP16 输出溢出;钳位发生在累加器内部,阈值为 \(2^{15}\)。因此,累加器会饱和到 \(2^{15}\),这恰好是带 16 位小数位的有符号 32 位定点数的取值范围。Since 32768 is itself a valid FP16 word (0x7800), the ANE's 0x7c00 can't be FP16 output overflow, the clamp happens inside the accumulator, at \(2^{15}\). Thus the accumulator saturates at \(2^{15}\), exactly the range of a signed 32-bit fixed-point value with 16 fractional bits.
非线性激活函数Nonlinear Activation
对于融合层,ANE 的计算如下:For a fused layer, the ANE computes:
\[ y = f(\sum_k x_k w_k + b) \]重要的是,完成的 MAC 求和会直接送入 MAC 后的激活模块,从而避免中间的内存往返。这之所以可行,是因为激活操作是逐点进行的:标量归约一旦完成,其激活结果便只取决于该标量,因此可以立即应用激活。Importantly, completed MAC sums feed directly into the post-MAC activation block, avoiding an intermediate memory round-trip. This is possible because the activation is pointwise: once a scalar reduction is complete, its activation depends only on that scalar and can be applied immediately.
要确定 ANE 如何实现 tanh(),请编译一个仅包含单个 TANH 激活层的 CoreML 模型,并检查生成的已编译硬件寄存器文件(hwx)。系数区域包含从 0x4288 开始的 33 个连续 FP16 字:To determine how the ANE implements tanh(), compile a CoreML model containing a single TANH activation layer and inspect the resulting compiled hardware register file (hwx). The coefficient region contains 33 consecutive FP16 words beginning at 0x4288:
00004270: 3120 3001 0000 0000 0000 0000 0000 0000
00004280: 0000 0044 0000 003c 0000 f52f d633 bc35 # 0.000000 0.124329 0.244873 0.358398
00004290: 6537 7038 1539 a239 183a 793a c93a 0a3b # 0.462158 0.554688 0.635254 0.704102 0.761719 0.809082 0.848145 0.879883
000042a0: 3e3b 673b 883b a23b b63b c63b d33b dd3b # 0.905273 0.925293 0.941406 0.954102 0.963867 0.971680 0.978027 0.982910
000042b0: e53b eb3b ef3b f33b f63b f83b fa3b fb3b # 0.986816 0.989746 0.991699 0.993652 0.995117 0.996094 0.997070 0.997559
000042c0: fc3b fd3b fe3b fe3b ff3b 0000 0000 0000 # 0.998047 0.998535 0.999023 0.999023 0.999512
000042d0: 003c 0300 6000 0000 0000 0000 0000 0000
这 33 个 FP16 字与 \(\tanh(x)\) 的 33 个 IEEE 小端序(LE)FP16 量化样本相匹配:Those 33 FP16 words match 33 IEEE LE FP16 quantized samples of \(\tanh(x)\):
\[ T_i=\operatorname{round}_{16}\!\left(\tanh(i/8)\right), \qquad i=0,1,\ldots,32. \]
现在切换到 RELU 激活层:Now switch to RELU activation layer:
| activation program | NonlinearMode |
lookup coefficients |
|---|---|---|
| identity | 0 | none |
| ReLU | 1 | none |
| tanh | 2 | 33 FP16 words |
因此,模式 2 选择一个自定义的 33 项查找表。33 个点定义了 32 个区间。当 \(R=3\) 时,节点为Thus, mode 2 selects a custom 33-entry lookup table. 33 points defines 32 intervals. With \(R=3\), the knots are
\[ x_i=\frac{i}{8},\qquad i=0,\ldots,32, \]以 \(1/8\) 的间距覆盖 \([0,4]\)。输入映射到表中的方式为 \(u=2^R|x|\),因此 \(R\) 决定了节点间距。其分辨率比 33 号分箱更平滑;我怀疑相邻条目之间进行了线性插值。为验证这一点,构建一个只有单个尖峰的冲激 LUT:covering \([0,4]\) with spacing \(1/8\). The input maps into the table as \(u=2^R|x|\), so \(R\) sets the knot spacing. The resolution is smoother than its 33 bin; I suspect that adjacent entries are linearly interpolated. To test, build an impulse LUT with a single spike:
\[ T_8=1,\qquad T_k=0\ \text{for }k\ne8,\qquad R=3. \]
然后让输入扫过 \(T_8\) 周围的两个单元格。测得的输出呈三角形:幅值从 \(|x|=7/8\) 时的 \(0\) 线性上升到 \(|x|=1\) 时的 \(1\),随后线性下降到 \(|x|=9/8\) 时的 \(0\)。Then sweep the input across the two cells around \(T_8\). The measured output forms a triangle: magnitude rises linearly from \(0\) at \(|x|=7/8\) to \(1\) at \(|x|=1\), then falls linearly to \(0\) at \(|x|=9/8\).
因此,我们知道模式 2 实现了一个包含 33 个条目的分段线性 LUT。\(R\) 将输入缩放到 LUT 坐标系中,Thus, we know that mode 2 implements a 33-entry piecewise-linear LUT. \(R\) scales the input into LUT coordinates,
\[ u=2^R|x|, \]so the knot spacing is \(\Delta x=2^{-R}\). \(\lfloor u\rfloor\) and \(\lceil u\rceil\) select the adjacent entries, and \(\alpha=u-\lfloor u\rfloor\) gives the interpolation weight between them.
缩放与偏置Scaling and Bias
CoreML 还支持线性缩放和偏置变换 \(ax + b\)。于是我猜想,\(ax + b\) 可能可以共用模式 2 的线性插值硬件。为了确认这一点,构建一个带有固定缩放系数和偏移量的 ReLU CoreML 模型:CoreML also supports a linear scaling and bias \(ax + b\) transform. I then suspected \(ax + b\) could share the linear interpolation hardware of mode 2. To confirm, construct a CoreML model with a ReLU with a constant scale and offset:
\[ z=4x-2,\qquad y=\operatorname{ReLU}\left(\frac{z}{2}+1\right), \]如果编译器将常量缩放因子和偏移量折叠进卷积中:If the compiler folds the constant scale and offset into the convolution:
\[ W'=\frac12W=2,\qquad b'=\frac12b+1=0, \] \[ y=\operatorname{ReLU}(2x). \]解码 model.espresso.weights 可确认,ReLU 上恰好进行了这种折叠变换:Decoding model.espresso.weights confirms exactly this folded transformation on ReLU:
authored convolution: W = 4, b = -2
activation affine: s = 0.5, c = 1
compiled convolution: W' = 2, b' = 0
而寄存器文件的十六进制差异展示了偏置和激活函数如何在编译时融合到同一条 MAC 后处理路径中:And the register file hexdiff shows how bias and activation are fused into the same post-MAC path at compile time:
| Probe | Tasks | BiasMode | PostScaleMode | NonlinearMode |
|---|---|---|---|---|
| Plain convolution | 1 | 0 | 0 | 0 |
Explicit Core ML Bias |
1 | 1 | 0 | 0 |
| Bias + ReLU | 1 | 1 | 0 | 1 |
| Bias + tanh | 1 | 1 | 0 | 2 |
极其邪门的想法:用非线性插值来计算额外的一次内核执行,或者把 int8 权重量化为 int4 权重。Extremely cursed idea: use nonlinear interpolation to compute an additional kernel pass, or quantize int8 into int4 weights.
2. 调度器2. Scheduler
ANE 驱动的源代码无聊得令人失望。驱动从不向 ANE 下发 CONV、MATMUL 或 RELU 操作码来执行。所有神经网络运算都已编译成由任务描述符(TD)组成的命令流,驱动软件只需将任务加载到内存中,通过 (TM_ADDR, TM_SIZE) 设置指向不透明任务数据块的指针,然后敲响门铃 (TM_PUSH),提交已准备好的任务。The ane driver source code is disappointingly boring. The driver never gives the ANE a CONV, MATMUL, or RELU opcode to run. All the neural operations have all already been compiled into a command stream of task descriptors (TDs), and the driver software simply loads the task to memory, sets the pointer to the opaque task blob via (TM_ADDR, TM_SIZE), and submits the staged task by ringing the doorbell (TM_PUSH).
static void ane_tm_push_tq(struct ane_device *ane, struct ane_request *req)
{
int qid = req->qid;
tm_write32(ane, TM_ADDR, tq_read32(ane, TQ_ADDR1(qid)));
tm_write32(ane, TM_INFO, tq_read32(ane, TQ_SIZE1(qid)) | req->td_count);
tm_write32(ane, TM_PUSH, TQ_PRTY_TABLE[qid] | (qid & 7) << 8); // magic
}
https://github.com/eiln/ane/blob/main/ane/src/ane_tm.c#L87
随后,硬件会一直负责处理该提交,直到完成;完成后,它会向 ARM64 核心触发中断。The hardware then owns the submission until completion, and raises an interrupt to the ARM64 core when it's done.
static void ane_tm_handle_irq(struct ane_device *ane)
{
int line;
line = 0;
for (u32 n = 0; n < tm_read32(ane, TM_IRQ_EVTC(line)); n++) {
这个(平淡无奇的)命令提交前端类似于 GPU 的命令提交前端,比如 NVIDIA 的 pushbuffer/PBDMA。软件提交驻留在内存中的命令流,GPU 的命令处理器遍历命令流并分发命令,而无需知道这些命令会执行什么。This (boring) command submission frontend resembles that of a GPU's, think NVIDIA's pushbuffer/PBDMA. The software submits a command stream resident in memory, and the GPU's command processor walks over command stream and dispatches the commands, without knowing what that command executes.
TM_ADDR 和 TM_INFO 是全局暂存寄存器;只有写入 TM_PUSH 后才会发生任何事情(“魔法”),而 TM_PUSH 会以原子方式提交暂存的启动状态。具体而言,TM_INFO 存储所提供流中的描述符总数:TM_ADDR and TM_INFO are global staging registers, and TM_PUSH atomically commits that staged launch state, given that nothing happens until TM_PUSH is written ("magic"). TM_INFO in particular stores the total number of descriptors in the supplied stream:
TM_INFO[31:16] = descriptor_dwords - 1
TM_INFO[15:0] = descriptor_count
TM_INFO 寄存器自然映射到一个硬件计数器:TM_INFO register naturally maps onto a hardware counter:
if (fetch) begin
if (word_ctr == descriptor_dwords_minus_1) begin
word_ctr <= 0;
desc_ctr <= desc_ctr + 1;
end else begin
word_ctr <= word_ctr + 1;
end
end
为什么是“减 1”?将长度编码为 length - 1,是一种对 RTL 友好的做法,可以让零起始计数器在关键路径之外终止。但请注意,与 GPU 命令会解析环形缓冲区中由可变长度描述符组成的流不同,ANE 只接收总数,这表明描述符的大小是固定的。Why the "minus 1"? Encoding length - 1 is an RTL-friendly way to terminate a zero-based counter out of the critical path. But note how, compared to GPU commands which parse a variable-length stream of descriptors in a ringbuffer, ANE only receives the total count, indicating that descriptors are fixed-size.
任务队列Task Queue
再深入一层,任务管理器从任务队列(TQ)中选择的是什么?Going one layer deeper, what's in a task queue (TQ) that the task manager selects from?
+------------------+
CPU / driver ------>| Task Manager |
| |
| schedule / fetch |
| / dispatch |
+--------+---------+
|
+------------------+------------------+
| | |
v v v
+---------+ +---------+ +---------+
| TQ 0 | ... | TQ 3 | ... | TQ 7 |
| BAR[32] | | BAR[32] | | BAR[32] |
| NID | | NID | | NID |
| state | | state | | state |
+---------+ +---------+ +---------+
同一个寄存器块有 8 个副本(由 qid(0…7)索引),结构如下:There's 8 copies of the same register block (indexed by qid (0…7)), structured as:
TQ[qid] + 0x000 STATUS
0x010 PRIORITY
0x014 VACANT
0x01c INFO
0x020 BAR1[0..31] // task1
0x0a0 NID1
0x0a4 SIZE2
0x0a8 ADDR2
0x0ac BAR2[0..31] // task2
0x12c NID2
0x130 SIZE1
0x134 ADDR1
next qid: +0x148
每个 TQ 包含:Each TQ holds:
- (1) 每个 TQ 的调度状态(状态、优先级和空闲情况)(1) Per-TQ scheduling state (status, priority, and vacancy)
- (2) 两组命令流描述符,每个 TQ 一组(ADDR1/ADDR2、SIZE1/SIZE2、NID1/NID2 和 32 个 BAR)。这两个槽位显然采用乒乓式暂存方案:一个槽位执行时,软件可以修改另一个槽位。(2) Two sets of command stream descriptors, per-TQ (ADDR1/ADDR2, SIZE1/SIZE2, NID1/NID2, and 32 BARs). The two slots are certainly a ping-pong staging scheme to let one slot execute, while software modifies the other slot.
注意 TM_PUSH 如何通过附加一个 qid 来执行 TM_ADDR/TM_SIZE 中引用的任务:Notice how TM_PUSH executes a task referenced in TM_ADDR/TM_SIZE by attaching a qid:
tm_write32(ane, TM_PUSH, TQ_PRTY_TABLE[qid] | (qid & 7) << 8); // magic
自然的理解是,描述符流指定要运行的任务,而 qid 选择该描述符运行所处的启动上下文。常驻 TQ 上下文(BAR、NID)很像 GPU 硬件通道。下面是我的驱动程序填充单个 TQ 并启动它的代码:The natural interpretation is that the descriptor stream specifies what task to run, while the qid selects the launch context the descriptor runs under. The resident TQ context (BAR, NID) is much like a GPU hardware channel. Here's my driver populating a single TQ to launch it:
int ane_tm_enqueue(struct ane_device *ane, struct ane_request *req)
{
int qid = req->qid;
tq_write32(ane, TQ_STATUS(qid), 0x1);
for (int bdx = 0; bdx < ANE_TILE_COUNT; bdx++) {
tq_write32(ane, TQ_BAR1(qid, bdx), req->bar[bdx]);
}
tq_write32(ane, TQ_SIZE1(qid), ((req->td_size >> 2) - 1) << 0x10);
tq_write32(ane, TQ_ADDR1(qid), req->btsp_iova);
tq_write32(ane, TQ_NID1(qid), (req->nid & 0xff) << 8 | 1);
return 0;
}
https://github.com/eiln/ane/blob/main/ane/src/ane_tm.c#L70
这里唯一重要的是 32 项的 BAR 表(基址寄存器)。接下来我们会讨论任务描述符,但编译后的 ANE 命令流只通过相对偏移量引用虚拟地址,而 BAR 提供基准 IOVA(IOMMU 外设虚拟地址)重定位地址。访问 ANE 虚拟地址时,需要在编译时提供一个硬编码的 BAR 基址偏移量;这意味着它没有 GPU 风格的 load/store 指令,无法从虚拟地址动态发起加载/存储操作。The only thing important here is the 32-entry BAR table (base address register). We'll get into task descriptors next, but the compiled ANE command stream only references virtual addresses by relative offsets, and BAR provides the base IOVA (IOMMU peripheral virtual address) relocation address. An ANE virtual address access needs a hard-coded BAR base offset supplied at compile time, meaning it lacks GPU-style load/store instructions that dynamically issue load/stores from virtual address.
任务描述Task Descriptor
任务管理器会遍历并执行固定大小的任务描述符链:The task manager walks over and executes chain of fixed-size task descriptors:
for (int i = 0; i < td_count; i++)
execute_task(td_block, i);
每个 TD 中包含什么?以下是最简单的 1x1 卷积的 TD 十六进制转储:What's in each TD? Here's a hexdump of the TD for the simplest 1x1 convolution:
# M1 h13, 1x1 convolution: X[1,8,4,1] -> Y[1,3,4,1]
# TD header KernelDMASrc Common TileDMASrc L2 PE NE TileDMADst
00000000: 02000000 00000000 0000042a 00000000 # Header: EON=1 LogEvents=0x42a
00000010: 00fff86a 00000000 30009800 00000000 # Header: DebugEvents=0xfff86a SPL TSR TSE SrcLoc=1 DstLoc=1
00000020: 03025024 00000021 f401f800 00000040 # Header: RBase0=4 WBase=5 KBase0=1 ENE=3 KernelDMA: packet
00000030: 00000000 00000081 00000081 00000081 # KernelDMA.Config[0..2]: En=1 Hint=2
00000040: 00000080 00000080 00000080 00000080 # KernelDMA.Config[3..6]: En=0 Hint=2
00000050: 00000080 00000080 00000080 00000080 # KernelDMA.Config[7..10]: En=0 Hint=2
00000060: 00000080 00000080 00000080 00000080 # KernelDMA.Config[11..14]: En=0 Hint=2
00000070: 00000080 00000000 00000040 00000080 # KernelDMA.Config[15]: En=0 Hint=2; Base[0..2]=0,1,2
00000080: 00000000 00000000 00000000 00000000 # KernelDMA.Base[3..6]=0
00000090: 00000000 00000000 00000000 00000000 # KernelDMA.Base[7..10]=0
000000a0: 00000000 00000000 00000000 00000000 # KernelDMA.Base[11..14]=0
000000b0: 00000000 00000040 00000040 00000040 # KernelDMA.Base[15]=0; Size[0..2]=1
000000c0: 00000040 00000040 00000040 00000040 # KernelDMA.Size[3..6]=1
000000d0: 00000040 00000040 00000040 00000040 # KernelDMA.Size[7..10]=1
000000e0: 00000040 00000040 00000040 00000040 # KernelDMA.Size[11..14]=1
000000f0: 00000040 00000080 00000080 00000080 # KernelDMA.Size[15]=1
00000100: 00000080 00000000 00000000 00000000
00000110: 00000000 00000040 00000040 00000040
00000120: 00000040 3c000000 00040001 00000001 # Common: packet; Win=1 Hin=4
00000130: 00000022 00000008 00000003 00040001 # Common: InFmt=2 OutFmt=2 Cin=8 Cout=3 Wout=1 Hout=4
00000140: 00000001 5000a021 00002041 00010001 # Common.Conv: Kw=1 Kh=1 Sx=1 Sy=1 Groups=1
00000150: 00000004 00000000 00000000 04144405 # Common: tileH=4 ActiveNE=2 AccDB=1
00000160: 00100000 00000000 6c013800 00033881 # Common: NID=1 TileSrc: packet; enabled
00000170: 00008880 00000000 00000040 00000100 # TileSrc: base=0 row=1 plane=4
00000180: 00000800 00000800 00000000 00000000 # TileSrc: depth=32 group=32
00000190: 00000000 00000000 00000000 00000000
000001a0: 00000000 01002031 00000000 00000100 # TileSrc.Fmt: mode=1 trunc=3 mem=2 intlv=1
000001b0: 00000000 00000000 00000000 00000000
000001c0: 00000000 00000000 00000000 00000000 # TileSrc.PixelOffset[1..3]=0
000001d0: 00000000 00000000 00000000 44004800 # L2: packet
000001e0: 00000000 00500172 00000000 00000010 # L2.Source: base=0 channel=1
000001f0: 00000080 00000080 00000080 00000000 # L2.Source: row=8
00000200: 00000000 00000000 00000000 00000000
00000210: 0050017a 00000200 00000000 00000000 # L2.Result: base=0x20 channel=0 row=0
00000220: 00000000 00000000 0c008800 00000000 # PE: packet
00000230: 00000000 00000000 00000000 1000c800 # PE: zero NE: packet
00000240: 00000082 00101c00 00000000 00000000 # NE: KernelFmt=2 BinaryPoint=28
00000250: 00003c00 18017800 040000c1 00000000 # NE: PostScale=0x3c00 TileDst: packet; En=1 Base=0
00000260: 00000040 00000100 00000300 00000300 # TileDst: row=1 plane=4 depth=12 group=12
00000270: 01302031 # TileDst.Fmt: mode=1 trunc=3 mem=2 intlv=1 zpad
重要的是,TD 并不是可执行的指令流。ANE 没有 ISA。TD 是一系列“ControlDMA”(这个名称是我编的)突发写入数据包,用于向 ANE 的硬件配置寄存器写入数据,例如输入维度、输入/输出地址、激活函数。每个 ControlDMA 数据包由一个 32 位传输字和随后连续的 N 个 32 位寄存器值组成:Important is that a TD is not an executable instruction stream. ANE has no ISA. TD is a sequence of "ControlDMA" (I made this name up) burst-write packets writes to the ANE's hardware configuration registers, such as input dimension, input/output address, activation function. Each ControlDMA packet consists of a 32-bit transfer word followed by N consecutive 32-bit register values:
31 26 25 2 1 0
+-----------------------+-----------------------------+----+
| register count minus 1| first register base index | 00 |
+-----------------------+-----------------------------+----+
再次注意这个“减 1”的终止计数。ControlDMA 是一种灵活的单向 DMA 引擎,可将 N 个 32 位字从 IOMMU 虚拟 DRAM 复制到 ANE 的物理寄存器空间。例如,TD 中 KernelDMASrc 的数据包头为 0xf401f800:Notice the “minus 1” termination count again. ControlDMA is a flexible unidirectional DMA engine that copies N 32-bit words from IOMMU virtual DRAM into the ANE’s physical register space. For example, KernelDMASrc's packet header in TD is 0xf401f800:
count = (0xf401f800 >> 26) + 1 = 62 words
register base = 0xf401f800 & 0x03fffffc = 0x1f800
这不是一条 LOAD_WEIGHTS 指令,而是将 0xf4 或 62 个连续字复制到从 0x1f800 开始的 KernelDMA 寄存器偏移处。这些 KernelDMA 配置值可以告知 KernelDMA 从哪里加载权重。This is not a LOAD_WEIGHTS instruction. It's copying 0xf4 or 62 consecutive words into the KernelDMA register offset starting at 0x1f800. And those KernelDMA configuration values can tell KernelDMA where to load the weights from.
| Start byte | Section | Information |
|---|---|---|
0x000 |
Header | Dependencies, chaining, and BAR selectors |
0x028 |
KernelDMASrc |
0xf401f800; 16 coefficient-DMA lanes |
0x124 |
Common |
0x3c000000; tensor and convolution geometry |
0x168 |
TileDMASrc |
0x6c013800; activation-source DMA |
0x1dc |
L2 |
0x44004800; local source/result configuration |
0x228 |
Processing engine | 0x0c008800; PE configuration |
0x23c |
Neural engine | 0x1000c800; MAC and post-processing configuration |
0x254 |
TileDMADst |
0x18017800; result-destination DMA |
0x274 |
End | 628 bytes total |
由于每个部分都写入一个 MMIO 寄存器块,因此 TD 可以清晰地划分为 ANE 的数据通路部分:Since each section writes to one MMIO register block, TD divides cleanly into ANE's datapath sections:
| Starting address | Size | Block name | What |
|---|---|---|---|
0x26bc00000 |
0x4000 |
Common | Broadcast configuration selector; inferred |
0x26bc04000 |
0x4000 |
L2 | L2 backing/register aperture |
0x26bc08000 |
0x4000 |
PE | Processing-element configuration |
0x26bc0c000 |
0x4000 |
NE / MAC | Kernel format, MAC, bias, scaling, and nonlinear controls |
0x26bc10000 |
0x3000 |
Unknown | Unidentified register bank |
0x26bc13000 |
0x4000 |
Tile DMA source | Input-tile addresses, strides, formats, and DMA controls |
0x26bc17000 |
0x4000 |
Tile DMA destination | Output-tile addresses, strides, formats, and DMA controls |
0x26bc1b000 |
0x4000 |
Unknown / tunables | Unidentified configuration and tunable registers |
0x26bc1f000 |
0x4000 |
Kernel | Kernel backing / kernel DMA-source aperture |
0x26bc23000 |
0x1000 |
Unknown | Unidentified register bank |
0x26bc24000 |
0x1000 |
Task Manager | Task submission, execution state, events, and completion |
0x26bc25000 |
0x1000 |
Task Queues | Eight queues containing TD stacks, NIDs, priorities, and request pointers |
TD 本质上是 ANE 数据通路寄存器的一份序列化寄存器文件转储。每个“ANE 程序”实际上就是数据通路单次处理过程的配置。我们可以配置固定数据通路的运行方式(受其提供的调节选项限制),但无法改变数据通路能够执行哪些操作,也无法改变这些操作的执行顺序。A TD is effectively a serialized register-file dump of the ANE’s datapath registers. Each “ANE program” is simply the configuration for one pass through the datapath. We can configure how the fixed datapath operates (subject to the knobs it exposes), but not what operations the datapath is capable of performing, or how those operations are sequenced.
当用于执行 TD 的“magic”原子字写入任务管理器时,大致会发生以下一系列事情:When the "magic" atomic word is written to task manager to execute a TD, roughly, the sequence of what happens:
- ControlDMA 将 TD 复制到配置寄存器中。ControlDMA copies TD into the configuration registers.
- KernelDMA 将内核 W 复制到内核内存(KMem)中。KernelDMA copies kernel W into kernel memory (KMem).
- TileDMA 将输入 X 从 DRAM 复制到 L2。TileDMA copies input \(X\) from DRAM into L2.
- 每个 MAC 核心根据其权重对一行进行归约,生成 \(Y\) 的一行。Each MAC core reduces a row by its weights, producing one row of \(Y\).
- 对 \(X\) 的所有行重复步骤 2–3。Steps 2–3 repeat for all rows of \(X\).
- 应用后处理,并将完成的结果存储在 L2 中。Postprocessing is applied, and the completed results are stored in L2.
- TileDMADst copies \(Y\) from L2 back to DRAM.
ANE 是一种固定功能的数据流引擎,而不是执行任意指令的 GPU。TD 配置的是一个特定领域的数据通路。限制硬件接口通常意味着面积更小、数据传输更具确定性、延迟更低、功耗更少。ANE 的编译器可以明确调度张量的操作,但这也意味着编译器必须明确调度张量的操作。这是一种权衡,但理由充分:通常在编译时我们就知道模型的样子。动态执行并非 ANE 的瓶颈。ANE 的处理器接口相对通用,它只负责启动任务,而这些任务可以描述 Transformer。ANE is a fixed-function dataflow engine, not a GPU executing arbitrary instructions. The TD configures a domain-specific datapath. Constraining the hardware interface usually means smaller area, deterministic movement, lower latency, and less power drawn. ANE's compiler can explicitly schedule what the tensors do, but that also means the compiler must explicitly schedule what the tensors do. This is a tradeoff, but a justified one: we usually know what the model looks like at compile time. Dynamic execution is not what limits ANE. ANE's processor interface is relatively generic, and it simply launches tasks, and the tasks can describe transformers.
例如,在编译时固定张量大小,并不意味着它无法处理可变长度张量:比如,可以通过按大小循环遍历不断增长的 KV 缓存;与这里真正的核心问题——内存流式传输带宽——相比,调度开销微不足道。真正让 ANE 更适合 CNN 而非 Transformer 的,是内存搬运。For example making tensor sizes fixed at compile time does not mean it can't handle variable-length tensors: for example, a growing KV cache can be traversed by looping over the size, and dispatch overhead is negligible relative to the elephant in the room here, that is, memory-streaming bandwidth. What actually shaped ANE for CNNs over transformers is memory movement.
3. 内存3. Memory
屋顶线Roofline
找出当前最慢的环节总是有益的,这样我们就能优化真正重要的部分。It's always good to identify our current slowest link, so we can optimize what actually matters.
Apple 的统一内存使 ANE 能够访问系统 DRAM 池中的缓冲区,而 CPU 和 GPU 也可以访问该内存池。这并不意味着 ANE 可以通过零拷贝直接从该 DRAM 池中流式读取数据。ANE 必须先将内存中的数据复制到其本地的“ANE 内存”或 SRAM 中。因此,任何受带宽限制的任务都会受限于 ANE 本地内存的流式传输吞吐量。Apple’s unified memory lets the ANE access buffers from the system DRAM pool accessible by the CPU and GPU. It does not mean the ANE zero-copy streams directly out of that DRAM pool. ANE must first copy any memory into its local "ANE memory" or SRAM. Any bandwidth-limited task will thus be limited by ANE's local memory streaming throughput.
M1 ANE 报告的性能为 11 TOP/s,系统 DRAM 带宽为 68 GB/s。一次 MAC 执行两次运算,但需要读取两个 FP16 操作数,即 4 字节:M1 ANE reports \(11\text{ TOP/s}\) at \(68\text{ GB/s}\) at system DRAM bandwidth. A MAC performs two operations but consumes two FP16 operands, or 4 bytes:
\[ \frac{2\text{ OP}}{4\text{ bytes}} =0.5\text{ OP/byte}. \]如果每个 MAC 操作数都从 DRAM 流式传输,要维持 11 TOP/s,就需要以流式方式传输。If every MAC operand streamed from DRAM, sustaining \(11\text{ TOP/s}\) would require streaming
\[ \frac{11\text{ TOP/s}}{0.5\text{ OP/byte}} = 22\text{ TB/s}. \]这相当于据报道的系统 DRAM 带宽(\(68\text{ GB/s}\))的 300 多倍。因此,通过从某种本地 ANE 内存池中读取数据并重复利用,或许可以达到 ANE MAC 的峰值吞吐量。which is over 300x times the reported \(68\text{ GB/s}\) system DRAM capacity. Thus peak ANE MAC throughput could be reached by fetching from some local ANE memory reservoir, and reusing it.
换句话说,M1 ANE 在 \(68\text{ GB/s}\) 下达到的 \(11\text{ TOP/s}\) 数值确定了屋顶线的脊点:Restated, the M1 ANE’s \(11\text{ TOP/s}\) at \(68\text{ GB/s}\) number sets the roofline ridge point:
\[ \frac{11\text{ TOP/s}}{68\text{ GB/s}} = 162\text{ OP/byte}. \]每从 DRAM 读取的一个字节,平均必须支撑至少 162 次运算,DRAM 带宽才不会成为瓶颈。换言之,工作负载必须提供足够的片上重用,使算术强度达到至少每字节 DRAM 流量 162 OP。低于 162:1 时,加快计算也不会提高每秒解码的 token 数。Each byte fetched from DRAM must support, on average, at least 162 operations for DRAM bandwidth to stop being the limiter. Equivalently, the workload must provide enough on-chip reuse to achieve an arithmetic intensity of at least 162 OP/byte DRAM traffic. Below the 162:1 ratio, speeding up compute won't increase decoded token/s.
内存层次结构Memory Hierarchy
即使(平均)DRAM 带宽足够,ANE 出于诸多原因也不会直接读取 DRAM,包括 DRAM 的确定性时序、物理布线、共享流量等。如果 ANE 的流量与 CPU、GPU、显示器等设备在 AXI 上竞争,就无法为 MAC 提供确定性时序。此外,如果 16 个核心要使用某个输入数据块,我们不希望发起 16 次相同的 DRAM 传输。“ANE 本地存储器”可以让某个操作产生的中间激活值直接供下一个操作使用,而不必先传到 DRAM 再传回来。Even if (average) DRAM bandwidth were sufficient, ANE does not read DRAM directly for many reasons, including DRAM deterministic timing, physical routing, shared traffic, etc. If ANE's traffic competes on AXI the CPU, GPU, display, and etc, it cannot provide deterministic timing to the MACs. Also, if 16 cores consume some input tile, we do not want to initiate 16 identical DRAM transfers. "ANE local memory" would allow intermediate activation produced by one operation be consumed by the next instead of traveling to DRAM and back.
Apple 有多种方式来组织本地存储层次结构。乘加专利描述了阵列周围的数据缓冲路径。Apple had several ways to organize local memory hierarchy. The multiply-accumulate patent describes the the data buffer paths around the array.
Unified DRAM
│
▼
┌───────────────────────────────────────────────┐
│ shared ANE L2 memory, 2 MiB │
└────────┬───────────────┬───────────────────┬──┘
│ │ │
▼ ▼ ▼
┌────────────┐ ┌────────────┐ ... ┌────────────┐
│ core 0 │ │ core 1 │ │ core N │
│ ┌────────┐ │ │ ┌────────┐ │ │ ┌────────┐ │
│ │ L1 │ │ │ │ L1 │ │ │ │ L1 │ │
│ └────────┘ │ │ └────────┘ │ │ └────────┘ │
│ ┌────────┐ │ │ ┌────────┐ │ │ ┌────────┐ │
│ │ KMem │ │ │ │ KMem │ │ │ │ KMem │ │
│ │ 64 KiB │ │ │ │ 64 KiB │ │ │ │ 64 KiB │ │
│ └────────┘ │ │ └────────┘ │ │ └────────┘ │
└────────────┘ └────────────┘ └────────────┘
- KMem:每个内核配备 16 个 64 KiB 的“L1” SRAM。总计 1 MiB。KMem: 16x per-core 64 KiB "L1" SRAM for kernel. Total 1 MiB.
- L1:每个核心的 MAC 输入“L1”暂存区(共 16 个)。L1: 16x per-core MAC input "L1" staging area.
- L2:所有核心共享一个 2 MiB 的 L2 缓存。L2: 1x shared 2 MiB L2 across all cores.
我不会假装自己从没反编译过那些玩意儿。ANE ARM64 固件的任务调试例程会:(1)从 KMem 索引 0 到 15 各转储 0x10000 字节;(2)然后单独转储一份 0x200000 字节的 L2 数据:I'm not gonna pretend like I've never decompiled shit. The ANE ARM64 firmware's task-debug routine (1) dumps 0x10000 bytes from KMem indices 0 through 15 (2) then dumps one separate 0x200000-byte L2 dump:
_DAT_26bc30000 = 0; // core 0
uVar7 = 0;
do {
*(undefined4 *)((long)pvVar2 + uVar7) = *(undefined4 *)(&DAT_26bc34000 + uVar7);
bVar1 = uVar7 < 0xfffc;
uVar7 = uVar7 + 4;
} while (bVar1);
CDebugUtility::fileWrite(this,pvVar2,0x10000,"./td_%d/kmem_%d-%d-%d_%d.bin"); // KMem #0 64 KiB
// ... repeat
_DAT_26bc30000 = 0xf; // core 15
uVar7 = 0;
do {
*(undefined4 *)((long)pvVar2 + uVar7) = *(undefined4 *)(&DAT_26bc34000 + uVar7);
bVar1 = uVar7 < 0xfffc;
uVar7 = uVar7 + 4;
} while (bVar1);
CDebugUtility::fileWrite(this,pvVar2,0x10000,"./td_%d/kmem_%d-%d-%d_%d.bin"); // KMem #15 64 KiB
uVar7 = 0;
do {
*(undefined4 *)((long)pvVar2 + uVar7) = *(undefined4 *)(&DAT_26bd00000 + uVar7);
bVar1 = uVar7 < 0x1ffffc;
uVar7 = uVar7 + 4;
} while (bVar1);
CDebugUtility::fileWrite(this,pvVar2,0x200000,"./td_%d/l2_%d-%d-%d.bin"); // L2 2 MiB
DMA 引擎DMA Engines
┌──────────┐ ┌────────────────────────────────────── ANE ──────────────────────────────────────┐
│ │ │ │
│ │ │ ┌─────────────┐ ┌────────────────────┐ │
│ ├──────►│ │ Control DMA │─────►│ Hardware Registers │ │
│ │ │ └─────────────┘ └────────────────────┘ │
│ │ │ ┌─────────── MAC ────────────┐ │
│ │ │ ┌──────────────┐ │ ┌──────────────────────┐ │ │
│ ├──────►│ │ KernelDmaSrc │─────────────────────────────►│ │ Kernel Memory │ │ │
│ DRAM │ │ └──────────────┘ │ └──────────┬───────────┘ │ │
│ │ │ │ ▼ │ │
│ │ │ ┌──────────────┐ ┌────────────────┐ │ ┌──────────────────────┐ │ │
│ ├──────►│ │ TileDmaSrc │──────►│ L2 │◄───►│ │ MAC Array │ │ │
│ │ │ └──────────────┘ │ Tile Memory │ │ └──────────────────────┘ │ │
│ │ │ ┌──────────────┐ │ │ └────────────────────────────┘ │
│ │◄──────┤ │ TileDmaDst │◄──────│ │ │
│ │ │ └──────────────┘ └────────────────┘ │
│ │ │ │
└──────────┘ └─────────────────────────────────────────────────────────────────────────────────┘
有三个用于与 MAC 收发数据的 DMA 引擎:There are three DMA engines to/from the MACs:
- KernelDMASrc[0..15]:十六个逻辑系数通道将权重从 DRAM 复制到对应的各核心 KMem 存储体中。
KernelDMASrc[0..15]: sixteen logical coefficient lanes copy weights from DRAM into the matching per-core KMem banks. - TileDMASrc:将瓦片从 DRAM 复制到 L2。
TileDMASrc: copies tiles from DRAM to L2. - TileDMADst:将分块从 L2 复制到 DRAM。
TileDMADst: copies tiles fromL2to DRAM.
仅凭十六个 KernelDMASrc 寄存器通道,并不能证明存在十六个物理上独立的 DMA 前端。通道扩展实验确实表明,逻辑通道能够并发推进:启用一个或十六个通道时,32 个有效任务的每个启用通道包含 1 MiB 系数数据,其总耗时基本相同。不过,它们仍可能汇入共享的请求生成逻辑、交叉开关、缓存和 DRAM 仲裁。The sixteen KernelDMASrc register lanes do not by themselves prove sixteen physically independent DMA front ends. A lane-scaling experiment does show that the logical lanes make concurrent progress: 32 valid tasks with 1 MiB aggregate coefficients per enabled lane take essentially the same time at one and sixteen lanes. They may still merge into shared request-generation, crossbar, cache, and DRAM arbitration.
这三个 tile/kernel DMA 引擎用起来很简单。你给它提供任意 src、dst 和 size,它就会传输那块内存。不过,三个 DMA 引擎的存在揭示了一些有意思的假设:The three tile/kernel DMA engines are chill. You supply it any src, dst, and size, and it will transfer that block of memory. However the existence of three DMA engines reveals some interesting assumptions:
- Kernel 竟然有一条专用路径。将内核与 tile 分开,意味着承认内核是一种独立的操作数类别,且其生命周期不同。Kernel gets a dedicated path, at all. Splitting kernel vs tile is a commitment that kernel is a distinct operand class with a different lifecycle.
- KernelDMA 是单向的,只用于加载。他们认为权重不会被写回。KernelDMA is unidirectional load-only. They thought that weights would not be written back.
- TileDMA 是双向共享的。中间的 tile 输出可以直接反馈,无需往返 DRAM。TileDMA is bidirectional and shared. Intermediate tile outputs can be fed back without a round-trip to DRAM.
- 内核存储器是每个核心私有的(复制了 16 份)。他们认为权重会在每个核心上被重复使用很多次。Kernel memory is private per core (replicated 16x). They thought that weights would be reused many times per-core.
- 新内核必须从 DRAM 加载,而不是从 L2 加载。他们可能没想到要对超过 64 KiB × 16 的大型权重进行流式传输。New kernels must be loaded from DRAM, and not L2. They probably didn't think to stream large weights that exceeds 64 KiB x 16.
内核 L1Kernel L1
卷积并不是对 \(A\) 和 \(B\) 进行对称的乘加运算。卷积会将同一个卷积核 \(w[k]\) 滑过输入:A convolution is not a symmetric multiply-add of \(A\) and \(B\). A convolution slides the same kernel \(w[k]\) across the input:
\[ y[p]=\sum_k x[p+k]\,w[k]. \]Shoutout 2k2 textbook 
请注意,两个抽头核 \(w_0\) 和 \(w_1\) 可以只加载一次,然后在整个输出期间重复使用:Notice how the two tap kernel \(w_0\) and \(w_1\) can be loaded once and then reused across the whole duration of the output:
MAD 0 MAD 1
coefficient w0 w1
× ×
shift 0: a0 a1 → y0
shift 1: a1 a2 → y1
shift 2: a2 a3 → y2
因此,内核可以驻留在本地内存中并重复使用,同时移入新的输入;这正是苹果的数据通路专利所描述的内容(https://patents.google.com/patent/US20190340491A1/en)。So the kernel can reside in local memory, and reused, while new inputs shift in; which is what Apple's datapath patent describes (https://patents.google.com/patent/US20190340491A1/en).
预期的稳态是:The intended steady state is:
┌── activation ───────┐
DRAM → L2 ──────────┤ × → MAC
└── output ←──────────┘
DRAM → KMem ─────────────────── coefficient
如果要设计一款专门用于大量执行卷积的 ASIC——其中卷积核并不是频繁动态流入的(咳,KV)——那么利用操作数的不对称性来减少卷积核的内存搬运,显然是不二之选,因为我们已经证明,在 MAC 单元达到饱和之前,ANE 仍然深受内存吞吐量限制(比例为 162:1)。If designing an ASIC to churn convolutions — where the kernel isn't something that's frequently and dynamically streamed in (ahem KV) — it's a no brainer to exploit the operand asymmetry and reduce kernel memory movement, because we've shown that the ANE is still deeply in the memory throughput-limited (162:1 ratio) region before the MACs can be saturated.
但为什么不允许 L2 → KMem?这样看来,统一 tile 和 kernel 路径甚至更容易。But why not allow L2 → KMem? It seems easier, even, to unify tile and kernel paths.
(1)这样 kernel 流量就不会和 tile 争抢 L2 带宽了?我认为,L2 竞争并非 Apple 省略 L2 -> KMem 路径的原因:常驻 kmem 之所以存在,正是因为预期 kernel 不会频繁加载。如果稳态下 kmem 流量可以忽略,那么只需将其优先级设得低于 tile 的 L2 访问即可。(1) So kernel traffic doesn't compete with tile for L2 bandwidth? I argue that L2 contention was not why Apple omitted the L2 -> KMem path: resident kmem exists at all because kernels were expected to be loaded infrequently. If kmem traffic is negligible in steady state, it could simply be given lower priority than tile L2 accesses.
(2)因为内核 L2 分发会增加复杂性?苹果已经将 L2 端点分发到每个核心;我认为,沿用同一 tile 路径的内核路径也没那么糟。(2) Because kernel L2 distribution adds complexity? Apple already distributes L2 endpoints to each core; I argue that a kernel path riding the same tile path is not that bad.
ANE 的物理布局以(字面意义上的)中央 L2 SRAM 矩形为中心,矩形两侧各有(7+7)个核心,顶部一侧有 2 个核心,底部一侧则是共享控制逻辑。两侧的 7+7 个核心从青色垂直主干水平接入 L2;但顶部的两个核心需要将相同的宽水平接口旋转后垂直引出,这很可能形成了顶部入口处那个显眼的垂直梳状结构。ANE's physical layout is centered around (literally) the center L2 SRAM rectangle, with (7+7) cores along each side of the rectangle, 2 cores on the top side, and shared control logic on the bottom side. The 7+7 side cores take L2 ingress horizontally from the cyan vertical trunk; but the two top cores need the same horizontal wide interface rotated and escaped vertically, which likely produces that conspicuous vertical comb in the top ingress.
当然,我完全是在纸上谈兵,不过我认为,加入 kmem L2 多路复用器,确实不至于有那么难。如果内核取数回到 DRAM,ANE 的解码性能就不会下降得那么严重。这让我觉得,Apple 只是从未预料到驻留在 L2 中的张量会变成内核;在 2017 年,这种假设是合理的。而且 Apple 也喜欢开发彼此隔离、模块化的模块,完全将 1 MiB 的 kmem(只读,这样能在 SRAM 布线方面节省一点面积)和 2 MiB 的 tile L2 分开开发,可能会更容易。Granted I am fully armchair engineering here, but adding the kmem L2 mux, I argue, really could not have been that bad. ANE decode performance would not have been as tanked if kernel fetches go back to DRAM. It makes me think that Apple simply never expected L2-resident tensors to become kernels, which was a valid assumption in 2017. And Apple also likes developing isolated modular modules, probably would've been easier to completely isolate development of the 1 MiB kmem (read-only, which would save a little area in SRAM routing) and 2 MiB tile L2.
4. 结束了吗?4. Is it over?
DRAM 吞吐量DRAM Throughput
GPU 比 ANE 快吗?Is the GPU faster than the ANE?
Transformer 单 token 解码是权重重用率和计算/内存比最差的情况,因为生成一个 token 就需要流式读取相当于整个模型大小的权重。不过,如果 ANE 和 GPU 都受读取带宽限制,那么无论峰值计算能力如何,读取带宽更高的一方每秒解码的 token 就会更多。Transformer single-token decode is the worst case for weight reuse and compute per memory ratio, because we need to stream the whole model’s worth of weights to generate one token. However, if both the ANE and GPU are read-bandwidth bound, whichever one that has higher read bandwidth will decode more token/s, regardless of peak compute capacity.
既然 ANE 和 GPU 共用同一 DRAM,那么可以合理地认为,ANE 在 DRAM 访问方面不一定会比 GPU 吃亏。Now since ANE and GPU share the same DRAM, it's fair game: the ANE is not necessarily penalized in DRAM access compared to the GPU.
- 如果 ANE 在单 token 解码时比 GPU 慢,那是因为它的 DMA 控制器无法维持足够多的并行请求,使 DQ 达到饱和。If ANE is slower than GPU at single token decode, it's because its DMA controller cannot maintain enough parallel requests to saturate DQ.
要测量 ANE 与 GPU 的 DRAM 读取吞吐量,请生成一个受读取带宽限制的工作负载(缓冲区远大于缓存,填充真实的伪随机数据,并且只读取一次),测量执行时间,并针对不同的读取大小重复此过程:To measure ANE vs GPU DRAM read throughput, generate a read-bandwidth-bound workload (buffers are much larger than the caches, filled with real pseudo-random data, and consumed once), and measure the execution time, and repeat for different read sizes:
\[ s=\frac{\text{change in measured execution time}} {\text{change in payload size}} \]拟合斜率回答了这样一个问题:多进行一次 DRAM 读取需要增加多少执行时间?其倒数就是设备持续的 DRAM 读取带宽。The fitted slope answers: how much additional execution time does one extra DRAM read require? The reciprocal is the device's sustained DRAM read bandwidth.
- ANE KernelDMA:37.99 GB/s:每次执行时间对应的 CoreML 内核大小。ANE KernelDMA: 37.99 GB/s: CoreML kernel size per execution time.
- ANE TileDMA: 59.08 GB/s: CoreML tile src size per execution time.
- GPU:77.70 GB/s:通过一个每次读取每个 private uint4 一次并写入数据相关校验和的着色器,测得 Metal 缓冲区读取速度。GPU: 77.70 GB/s: Metal buffer reads per execution time, via a shader that reads every private uint4 once and writes a data-dependent checksum.
ANE 的内核(操作数 A)带宽最高可达 38 GB/s,tile(操作数 B)则可达 60 GB/s。我们能否将两者相加,达到 38 + 60 = 98 GB/s,从而逼近 M3 的 100 GB/s DRAM 带宽上限?ANE's kernel (operand A) maxes out at 38 GB/s, and tile (operand B) at 60 GB/s. Can we issue 38 + 60 = 98 GB/s to hit M3's 100 GB/s DRAM ceiling?
不。实验表明,kernel+tile 的组合运行时间等于各自单独运行时间之和。如果请求之间存在任何重叠,那么较短路径带来的额外耗时就会很少甚至没有,但执行时间(而非吞吐量)严格单调。No. Experiment shows that the kernel+tile combined runtime matched the sum of the isolated runtimes. If the requests overlapped at all, then the shorter path would contribute little or no additional time, but execution time (not throughput) is strictly monotonic.
\[T_{AB}=0.001+0.939T_A+0.981T_B\]
因此,ANE 的内核和 tile DMA 请求是串行发送的(一次一个),这意味着 ANE DRAM 的吞吐量简直被双重干爆了:Thus, ANE's kernel and tile DMA requests are sent serially (one at a time), meaning ANE DRAM throughput is double-fucked:
- 单独运行的内核和 tile DMA 的读取带宽都低于 GPU 的读取带宽(GB/s)。Both isolated kernel and tile DMA are lower than GPU's read GB/s.
- 内核和 tile DMA 的时间也是相加的。Kernel and tile DMA times are also additive.
除非……?Unless…?
当 ANE 解码受限于 DRAM 屋顶线时,像 10 -> 25 tok/s 这样高达 2.5 倍的显著提升,只能来自约 2.5 倍的内存流式带宽提升。With ANE decode pinned to the DRAM roofline, a drastic 2.5x improvement like 10 -> 25 tok/s can only come from ~2.5× higher memory streaming bandwidth.
让 ANE 再次达到 50 GB/s
但是……如果我们能额外增加 50 GB/s 的 kernelDMA 吞吐量呢?But … what if we could add 50 GB/s of additional kernelDMA throughput ?