逆向工程苹果神经引擎:架构演变与式微
三年前,我停止了逆向工程 Apple Neural Engine (ANE) 驱动 的工作。当时有个令人沮丧的小发现:ANE 这个硬件块其实没那么有用,我可以去做更有价值的事,于是转向了其他更有用的硬件块的上游化工作。ANE 的架构过于特化,很难围绕它搭建一个通用加速器平台;就算写一个 Linux 驱动直接开放 ANE 的硬件 API,也无法扩大它能处理的负载范围。事实上,macOS 里唯一常用 ANE 的场景,就只是在 Finder 里生成放大的预览图。
https://github.com/eiln/ane/tree/main
M1 芯片显微图:https://mastodon.social/@dougall/115149886886125067M5(2025)的主打卖点是“LLM 性能”,而他们顺水推舟地把 ANE 核心塞进了 GPU 核心里——我早就料到会这样,但独立 NPU 的末日,从这一刻起算是正式开始了。所以,为了纪念 ANE 明显的式微,我们要做一件更没用的事:回头去逆向工程 M1 上的 ANE,把当年没做完的事做完。一晃三年过去了(操),按理说现在的我应该比当年懂得更多。
三年前的目标是让 ANE 变得有用——在上面跑各种算子;这一次,目标更偏向于完整梳理它的内部架构——计算单元、数据通路、调度器、内存和执行模型。因为这些内部设计决策,揭示了苹果当初愿意最先固化到 A11 Bionic(2017)芯片里的那些对机器学习负载的假设,也反映了从 CNN 时代的 NPU 到如今跑 Transformer 负载的 GPU 这一转变背后的逻辑。
1. 计算单元
ANE 的 16 个计算核心可能是整个芯片里最无趣的部分。苹果最初的目标是密集的图像处理 CNN 负载,这类负载本质上是可预测复用的稠密张量归约运算。M1 的 ANE 计算核心是一个大规模并行的乘加(MAC)单元阵列,但仅凭这一点,几乎看不出它是为哪种负载设计、又擅长加速什么的。

卷积层执行的是激活窗口与学习得到的核权重之间的点积,而注意力机制执行的则是查询向量与键向量之间的点积。点积就是点积,MAC 恰好完成的就是这个任务。将 2017 年的 CNN 模型特化到 ANE 上的关键,并非 MAC 本身,而是围绕 MAC 的数据流:MAC 的输入和输出何时进入、停留于何处、如何移动。Transformer 打破的可预测复用模式(尤其是自回归解码过程)是 ANE 所依赖的,而 ANE 正是利用这一点设计出了足以在手机端运行的数据流效率。M5 的决策表明,ANE 的计算核心对于 Transformer 依然有用,只是处于一种不同的数据流之中。
尽管如此,每个计算核心内部的数据通路仍如下所示:
┌────────────────────── core ─────────────────────┐
│ ┌───────── 256× MACs ─────────┐ ┌────────────┐ │
│ │ MAD ─► add ─► accumulator │─►│ activation │ │
│ │ ▲ │ │ └────────────┘ │
│ │ └──────────┘ │ │
│ └─────────────────────────────┘ │
└─────────────────────────────────────────────────┘
Multiply-Accumulate
ANE 拥有 16 个并行计算核心。每个计算核心包含 128 个 FP16(或 256 个 INT8)并行乘累加(MAC)通道。每个 MAC 通道执行如下递推运算:
\[ s\leftarrow s+a\times b \]即先对两个操作数 \(a\) 和 \(b\) 执行乘法,再将乘积累加到滚动总和(累加器)中。
在 T 个周期内重复执行 MAC 操作,即可计算一个 T 项点积:
\[ s_T=s_0 + \sum_{t=0}^{T-1} a_t \, b_t. \]因此,单个 MAC 通道实际上执行的是沿时间维度的标量归约。16 核心的 ANE 拥有 2048 个并行 MAC 通道:
\[ 128\ \text{lanes/core}\times16\ \text{cores} = 2048\ \text{parallel MAC lanes} \]这意味着每个周期内会在空间维度上并行执行 2048 个归约操作,时间成为唯一的归约轴:
\[ S_T[q,p] = S_0[q,p] + \sum_{t=0}^{T-1} a_t[q,p]\,b_t[q]. \]单个 MAC 通道并不清楚自己正在归约的是矩阵或张量的哪个维度。值得注意的是,点积、矩阵乘法与卷积之间的区别,取决于操作数如何映射和调度到核心上。除了后文提到的核内存(kernel memory)外,ANE 核心并未在硬件层面将 4 通道的 CNN 层编码进去。
内部结构上,MAC 数据通路由一个乘法器、一个加法器和一个 32 位累加寄存器组成。每个周期,加法器将乘法器的最新输出与上一周期的累加值相加,结果即为新的累加和。
operand a ──┐ ┌────────────┐ p[31:0] ┌──────────────┐ s_next[31:0] ┌─────────────┐
├──►│ MULTIPLIER │────────────►│ 32-BIT ADDER │────────────────►│ ACCUMULATOR │
operand b ──┘ └────────────┘ └──────▲───────┘ └──────┬──────┘
│ │ s[31:0]
└────────────────────────────────┘
这条反馈路径将部分和保持在 MAC 通道内的本地存储中,无需在 MAC 周期之间从远程外部内存取数。
关于精度,ANE 采用定点累加,读出时转换为 FP16。乘法器位宽为 16 位,在 32 位寄存器中按 Q16.16 格式累加,最后通过符号扩展等操作转为 FP16 输出。为探测累加器的范围,可使用十六进制整数 FP16 表示法:构建一个 CoreML ANE 程序,计算全 1 向量与目标向量的点积,这样每次乘法结果有界,但累加器中的累加和会持续增大:
\[ 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 位有符号定点数的范围。
非线性激活
对于融合层,ANE 的计算公式如下:
\[ y = f(\sum_k x_k w_k + b) \]关键在于,MAC 求和完成后会直接送入后置激活单元,无需中间写回内存。这是因为激活函数是逐点运算:标量归约一旦完成,其激活值只取决于这个标量,可以立即计算。
想知道 ANE 是如何实现 tanh() 的,可以编译一个只含单个 TANH 激活层 的 CoreML 模型,然后查看编译后的硬件寄存器文件(hwx)。系数区域从 0x4288 开始,有 33 个连续的 FP16 字:
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 小端 FP16 量化采样点:
\[ T_i=\operatorname{round}_{16}\!\left(\tanh(i/8)\right), \qquad i=0,1,\ldots,32. \]
接下来把激活层换成 RELU:
| 激活程序 | NonlinearMode | 查找表系数 |
|---|---|---|
| 恒等 | 0 | 无 |
| ReLU | 1 | 无 |
| tanh | 2 | 33 个 FP16 字 |
可见模式 2 对应一个 33 项的定制查找表。33 个点划分出 32 个区间。取 \(R=3\) 时,节点为
$$ x_i=\frac{i}{8},\qquad i=0,\ldots,32, $$ 覆盖 $[0,4]$ 区间,步长为 $1/8$。输入映射到查找表时,$u=2^R|x|$,因此 $R$ 决定了节点间距。其分辨率比 33 个 bin 更平滑,我推测相邻条目之间进行了线性插值。为了验证,构建一个仅含单峰脉冲的 LUT: $$ 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$。
由此可知,模式 2 实现了一个具有 33 个条目的分段线性 LUT。$R$ 将输入缩放至 LUT 坐标:
$$ u=2^R|x|, $$因此节点间距为 $\Delta x=2^{-R}$。$\lfloor u\rfloor$ 和 $\lceil u\rceil$ 选定相邻条目,而 $\alpha=u-\lfloor u\rfloor$ 给出它们之间的插值权重。
缩放与偏置
CoreML 还支持 $ax + b$ 的线性缩放与偏置变换。我随即怀疑 $ax + b$ 可能共享模式 2 的线性插值硬件。为了确认,构建一个带有常数缩放和偏移的 ReLU 的 CoreML 模型:
$$ z=4x-2,\qquad y=\operatorname{ReLU}\left(\frac{z}{2}+1\right), $$如果编译器将常数缩放和偏置折叠进卷积:
$$ W'=\frac12W=2,\qquad b'=\frac12b+1=0, $$ $$ y=\operatorname{ReLU}(2x). $$解码 model.espresso.weights 证实了 ReLU 上确实存在这一折叠变换:
authored convolution: W = 4, b = -2
activation affine: s = 0.5, c = 1
compiled convolution: W' = 2, b' = 0

寄存器文件的 hexdiff 展示了编译时偏置和激活如何融合到同一条 MAC 后处理路径中:
| Probe | Tasks | BiasMode | PostScaleMode | NonlinearMode |
|---|---|---|---|---|
| 普通卷积 | 1 | 0 | 0 | 0 |
显式 Core ML Bias | 1 | 1 | 0 | 0 |
| Bias + ReLU | 1 | 1 | 0 | 1 |
| 偏置 + tanh | 1 | 1 | 0 | 2 |
一个极其离谱的想法:使用非线性插值来执行额外的内核遍历,或者将 int8 量化为 int4 权重。
2. 调度器
ane 驱动程序源代码乏善可陈。驱动程序从不向 ANE 下发执行 CONV、MATMUL 或 RELU 的指令码。所有神经运算早已编译成由任务描述符(TD)组成的命令流,驱动软件只需将任务加载到内存,通过 (TM_ADDR, TM_SIZE) 设置指向不透明任务块(task blob)的指针,然后敲铃(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 核心发出中断。
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 的命令处理器遍历该命令流并分派指令,而无需知晓指令的具体执行内容。
TM_ADDR 和 TM_INFO 是全局暂存寄存器,TM_PUSH 原子性地提交该暂存启动状态;在写入 TM_PUSH 之前,“魔法”尚未触发,什么都不会发生。特别地,TM_INFO 存储了所提供命令流中描述符的总数:
TM_INFO[31:16] = descriptor_dwords - 1
TM_INFO[15:0] = descriptor_count
TM_INFO 寄存器自然映射到一个硬件计数器:
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
为什么是"minus 1"?把 length - 1 编码进来,是一种对 RTL 友好的做法,可以让从零开始计数的计数器在关键路径之外终止。不过值得注意的是,GPU 命令需要在 ringbuffer 里解析变长的描述符流,而 ANE 只接收一个总数,这说明描述符是固定大小的。
任务队列
再深入一层,任务管理器从中调度的任务队列(TQ)里到底有什么?
+------------------+
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) 索引),结构如下:
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 包含:
- (1) 每个 TQ 各自的调度状态(status、priority 和 vacancy)
- (2) 两组命令流描述符(ADDR1/ADDR2、SIZE1/SIZE2、NID1/NID2 以及 32 个 BAR)。这两个槽位显然是一个乒乓暂存方案:一个槽位在执行时,软件可以修改另一个槽位。
注意 TM_PUSH 是通过附加一个 qid 来执行 TM_ADDR/TM_SIZE 所引用的任务的:
tm_write32(ane, TM_PUSH, TQ_PRTY_TABLE[qid] | (qid & 7) << 8); // magic
一种自然的理解是:描述符流指定要执行的任务,而 qid 则选择该描述符运行所处的启动上下文。驻留的 TQ 上下文(BAR、NID)非常类似于 GPU 的硬件通道。
以下是我驱动的代码,用于填充单个 TQ 以启动它:
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 表(基地址寄存器,Base Address Register)。接下来我们将讨论任务描述符,但编译后的 ANE 指令流仅通过相对偏移引用虚拟地址,而 BAR 提供了基础 IOVA(IOMMU 外设虚拟地址)重定位地址。ANE 的虚拟地址访问需要编译时提供的硬编码 BAR 基址偏移,这意味着它缺乏类似 GPU 的动态从虚拟地址发起 load/store 指令。
任务描述符
任务管理器会遍历并执行一串固定大小的任务描述符:
for (int i = 0; i < td_count; i++)
execute_task(td_block, i);
每个 TD 里包含什么?以下是最简单的 1x1 卷积的 TD 的十六进制转储:
# M1 h13, 1x1 卷积:X[1,8,4,1] -> Y[1,3,4,1]
# TD 头部 KernelDMASrc Common TileDMASrc L2 PE NE TileDMADst
00000000: 02000000 00000000 0000042a 00000000 # 头部: EON=1 LogEvents=0x42a
00000010: 00fff86a 00000000 30009800 00000000 # 头部: DebugEvents=0xfff86a SPL TSR TSE SrcLoc=1 DstLoc=1
00000020: 03025024 00000021 f401f800 00000040 # 头部: RBase0=4 WBase=5 KBase0=1 ENE=3 KernelDMA: 数据包
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: 数据包; 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: 数据包; 已启用
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: 数据包
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: 数据包
00000230: 00000000 00000000 00000000 1000c800 # PE: zero NE: 数据包
00000240: 00000082 00101c00 00000000 00000000 # NE: KernelFmt=2 BinaryPoint=28
00000250: 00003c00 18017800 040000c1 00000000 # NE: PostScale=0x3c00 TileDst: 数据包; 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 位寄存器值组成:
31 26 25 2 1 0
+-----------------------+-----------------------------+----+
| register count minus 1| first register base index | 00 |
+-----------------------+-----------------------------+----+
注意又是“减 1”的终止计数写法。ControlDMA 是一个灵活的单向 DMA 引擎,把 N 个 32 位字从 IOMMU 虚拟内存拷贝到 ANE 的物理寄存器空间。
比如,TD 中 KernelDMASrc 的包头是 0xf401f800:
count = (0xf401f800 >> 26) + 1 = 62 words
register base = 0xf401f800 & 0x03fffffc = 0x1f800
这并不是一条 LOAD_WEIGHTS 指令,而是把 0xf4 即 62 个连续的字拷贝到从 0x1f800 开始的 KernelDMA 寄存器偏移处。这些 KernelDMA 配置值会告诉 KernelDMA 从哪里加载权重。
| 起始字节 | 区块 | 信息 |
|---|---|---|
0x000 | Header | 依赖、链式连接和 BAR 选择器 |
0x028 | KernelDMASrc | 0xf401f800;16 条系数 DMA 通道 |
0x124 | Common | 0x3c000000;张量与卷积几何参数 |
0x168 | TileDMASrc | 0x6c013800;激活值源 DMA |
0x1dc | L2 | 0x44004800;本地源/结果配置 |
0x228 | Processing engine | 0x0c008800;PE 配置 |
0x23c | Neural engine | 0x1000c800;MAC 与后处理配置 |
0x254 | TileDMADst | 0x18017800;结果目标 DMA |
0x274 | End | 总计 628 字节 |
由于每个代码段都写入一个 MMIO 寄存器块,TD(任务描述符)可以干净利落地映射到 ANE 的数据通路(datapath)的各个部分:
| 起始地址 | 大小 | 模块名称 | 功能 |
|---|---|---|---|
0x26bc00000 | 0x4000 | Common | 广播配置选择器;通过推断得出 |
0x26bc04000 | 0x4000 | L2 | L2 备份/寄存器窗口 |
0x26bc08000 | 0x4000 | PE | 处理元素(Processing Element)配置 |
0x26bc0c000 | 0x4000 | NE / MAC | 卷积核格式、乘累加(MAC)、偏置、缩放及非线性控制 |
0x26bc10000 | 0x3000 | 未知 | 未识别的寄存器组 |
0x26bc13000 | 0x4000 | Tile DMA 源 | 输入图块(Tile)地址、步长、格式及 DMA 控制 |
0x26bc17000 | 0x4000 | Tile DMA 目标 | 输出图块地址、步长、格式及 DMA 控制 |
0x26bc1b000 | 0x4000 | 未知 / 可调参数 | 未识别的配置和可调寄存器 |
0x26bc1f000 | 0x4000 | Kernel | 卷积核备份 / 卷积核 DMA 源窗口 |
0x26bc23000 | 0x1000 | 未知 | 未识别的寄存器组 |
0x26bc24000 | 0x1000 | Task Manager | 任务提交、执行状态、事件及完成信号 |
0x26bc25000 | 0x1000 | Task Queues | 八个队列,包含 TD 栈、节点 ID(NID)、优先级及请求指针 |
TD 本质上就是 ANE 数据通路寄存器的序列化转储。每个“ANE 程序”仅仅是数据通路的一次遍历配置。我们可以配置固定数据通路的运行方式(取决于其暴露的控制接口),但无法改变数据通路本身支持的操作类型,也无法改变这些操作的执行顺序。
当向任务管理器写入“魔法”原子字以执行 TD 时,大致的执行序列如下:
- ControlDMA 将 TD 复制到配置寄存器中。
- KernelDMA 将权重矩阵 W 复制至内核内存(KMem)。
- TileDMA 将输入矩阵 \(X\) 从 DRAM 复制到 L2 缓存。
- 每个 MAC 核心将其对应行与权重进行归约运算,生成 \(Y\) 的一行。
- 重复步骤 2–3,直至覆盖 \(X\) 的所有行。
- 应用后处理,并将完成的计算结果存入 L2。
- TileDMADst 将 \(Y\) 从 L2 写回 DRAM。
ANE 是一个固定功能的数据流引擎,而非执行任意指令的 GPU。TD 构建的是面向特定领域的专用数据通路。限制硬件接口的代价通常是更小的芯片面积、确定性的数据移动、更低的延迟以及更低的功耗。ANE 的编译器可以显式调度张量操作,但这同时也意味着编译器必须显式地管理张量的执行顺序。这是一种权衡,但合情合理:编译时通常已知模型结构。限制 ANE 性能的并非动态执行能力。ANE 的处理器接口相对通用,它仅负责启动任务,而这些任务足以描述 Transformer 结构。
例如,编译时固定张量尺寸并不意味着无法处理变长张量:比如,增长的 KV cache 可以通过循环遍历其尺寸来处理,而调度开销相较于主要瓶颈——即内存流带宽——几乎可以忽略不计。真正让 ANE 在处理 CNN 时优于 Transformer 的因素是内存移动效率。
3. 内存
Roofline
识别当前最慢的环节总是有益的,这样我们才能优化真正重要的部分。
Apple 的统一内存架构允许 ANE 访问 CPU 和 GPU 也能访问的系统 DRAM 内存池中的缓冲区。但这并不意味着 ANE 直接从该 DRAM 池进行零拷贝流式传输。ANE 必须先将内存数据复制到其本地的“ANE 内存”或 SRAM。因此,任何受带宽限制的任务都将受限于 ANE 本地内存的流式吞吐量。
M1 ANE 在 \(68\text{ GB/s}\) 系统 DRAM 带宽下报告的峰值性能为 \(11\text{ TOP/s}\)。一次 MAC 运算执行两个操作,但消耗两个 FP16 操作数(即 4 字节):
\[ \frac{2\text{ OP}}{4\text{ bytes}} =0.5\text{ OP/byte}. \]如果每个 MAC 操作数都从 DRAM 流式传输,要维持 \(11\text{ TOP/s}\) 的吞吐量,需要的流式传输带宽为:
\[ \frac{11\text{ TOP/s}}{0.5\text{ OP/byte}} = 22\text{ TB/s}. \]这超过了报告的 \(68\text{ GB/s}\) 系统 DRAM 容量的 300 多倍。因此,ANE 的峰值 MAC 吞吐量只能通过在本地 ANE 内存池中获取并复用数据来实现。
换个说法,M1 ANE 在 \(68\text{ GB/s}\) 带宽下达到 \(11\text{ TOP/s}\),由此得到 roofline 拐点:
\[ \frac{11\text{ TOP/s}}{68\text{ GB/s}} = 162\text{ OP/byte}. \]从 DRAM 读入的每个字节平均要支撑至少 162 次运算,DRAM 带宽才不再是瓶颈。换句话说,工作负载必须提供足够的片上数据复用,使计算强度达到至少 162 OP/byte 的 DRAM 流量。低于 162:1 这个比例,再怎么加速计算也无法提高每秒解码的 token 数。
内存层级
即便(平均)DRAM 带宽够用,ANE 也不会直接读 DRAM,原因有很多:DRAM 的时序是确定性的、物理布线限制、与其他模块共享流量等等。如果 ANE 的流量要和 CPU、GPU、显示等在 AXI 总线上竞争,就无法给 MAC 提供确定性时序。另外,如果 16 个核都要用同一块输入数据,我们也不想发起 16 次一模一样的 DRAM 读取。"ANE 本地内存"可以让上一个操作产生的中间激活值直接被下一个操作消费,而不必绕到 DRAM 走一圈。
Apple 本地内存层级有多种组织方式。乘加专利描述了阵列周围的数据缓冲路径。
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 用于存放 kernel,总共 1 MiB。
- L1:每个核 16 个 MAC 输入的 "L1" 暂存区。
- L2:所有核共享的 2 MiB L2。
别装没拆过固件。ANE ARM64 固件里的任务调试例程会先倾出 16 个 KMem 索引对应的各 0x10000 字节数据,再单独倾出一块 0x200000 字节的 L2 数据:
_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 Engines
┌──────────┐ ┌────────────────────────────────────── ANE ──────────────────────────────────────┐
│ │ │ │
│ │ │ ┌─────────────┐ ┌────────────────────┐ │
│ ├──────►│ │ Control DMA │─────►│ Hardware Registers │ │
│ │ │ └─────────────┘ └────────────────────┘ │
│ │ │ ┌─────────── MAC ────────────┐ │
│ │ │ ┌──────────────┐ │ ┌──────────────────────┐ │ │
│ ├──────►│ │ KernelDmaSrc │─────────────────────────────►│ │ Kernel Memory │ │ │
│ DRAM │ │ └──────────────┘ │ └──────────┬───────────┘ │ │
│ │ │ │ ▼ │ │
│ │ │ ┌──────────────┐ ┌────────────────┐ │ ┌──────────────────────┐ │ │
│ ├──────►│ │ TileDmaSrc │──────►│ L2 │◄───►│ │ MAC Array │ │ │
│ │ │ └──────────────┘ │ Tile Memory │ │ └──────────────────────┘ │ │
│ │ │ ┌──────────────┐ │ │ └────────────────────────────┘ │
│ │◄──────┤ │ TileDmaDst │◄──────│ │ │
│ │ │ └──────────────┘ └────────────────┘ │
│ │ │ │
└──────────┘ └─────────────────────────────────────────────────────────────────────────────────┘
MAC 与外部之间共有三条 DMA 通路:
KernelDMASrc[0..15]:16 条逻辑系数通道,负责将权重从 DRAM 复制到对应核心的 KMem 存储 bank。TileDMASrc:将 Tile 从 DRAM 复制到 L2。TileDMADst:将 Tile 从L2复制回 DRAM。
十六个 KernelDMASrc 寄存器 lane 本身并不能证明存在十六个物理独立的 DMA 前端。不过 lane 扩展实验表明,这些逻辑 lane 确实在并发推进:启用 1 个 lane 和 16 个 lane 时,各 lane 携带 1 MiB 系数、共 32 个有效任务的耗时基本相同。它们仍可能汇聚到共享的请求生成、crossbar、缓存和 DRAM 仲裁上。
三个 tile/kernel DMA 引擎则很佛系:你给它任意 src、dst 和 size,它就把那块内存搬过去。但三个 DMA 引擎的这种设计,透露出一些有趣的假设:
- Kernel 居然有一条专属通路。把 kernel 和 tile 分开,等于承诺 kernel 是一种生命周期不同的独立操作数类型。
- KernelDMA 是单向、只读加载。设计者认为权重不会被写回。
- TileDMA 是双向共享的。中间 tile 输出可以直接回灌,不必绕道 DRAM。
- Kernel 内存是每个 core 私有的(复制 16 份)。设计者认为权重在每个 core 上会被反复重用。
- 新 kernel 必须从 DRAM 加载,而不是 L2。他们大概没考虑过流式加载超过 64 KiB × 16 的大权重。
Kernel L1
卷积并不是 \(A\) 和 \(B\) 的对称乘加。卷积是让同一个 kernel \(w[k]\) 在输入上滑动:
\[ y[p]=\sum_k x[p+k]\,w[k]. \]致敬 2k2 教科书

注意两个 tap 的 kernel \(w_0\) 和 \(w_1\) 只需加载一次,就能在整个输出期间反复使用:
MAD 0 MAD 1
coefficient w0 w1
× ×
shift 0: a0 a1 → y0
shift 1: a1 a2 → y1
shift 2: a2 a3 → y2
也就是说,kernel 可以驻留在本地内存中反复重用,新的输入不断滑入——这正是 Apple 数据通路专利所描述的方案(https://patents.google.com/patent/US20190340491A1/en)。
预期的稳态工作方式是:
┌── activation ───────┐
DRAM → L2 ──────────┤ × → MAC
└── output ←──────────┘
DRAM → KMem ─────────────────── coefficient
如果目标是设计一款专门处理卷积运算的 ASIC——在这种场景中,卷积核通常不会被频繁动态流式加载(没错,就是在说 KV)——那么充分利用操作数不对称性、减少卷积核内存访问量就显得理所当然。因为前文已经证实,在 MAC 单元达到饱和之前,ANE 仍处于深度受限于内存吞吐量的区间(比例高达 162:1)。
那为什么不开放 L2 → KMem 的路径呢?这样做似乎甚至更简单,可以让 tile 和 kernel 的数据通路实现统一。
(1) 是为了避免 kernel 流量与 tile 争抢 L2 带宽?我认为,L2 争用并非 Apple 省略 L2 -> KMem 路径的原因:常驻 kmem 的存在本身就预设了卷积核加载频率极低这一前提。如果 kmem 的流量在稳态下可以忽略不计,那么只需简单地将 tile 的 L2 访问设为更高优先级即可。
(2) 是因为 kernel 的 L2 分布增加了复杂度?Apple 已经为每个核心配置了独立的 L2 端点;我认为,让 kernel 路径复用相同的 tile 路径并不会带来太大麻烦。

ANE 的物理布局以(字面意义上的)中心 L2 SRAM 矩形为核心,矩形上下两侧各分布有 7+7 个核心,上侧另有 2 个核心,底部则是共享控制逻辑。 位于 7+7 侧的核心通过青色垂直主干横向获取 L2 输入;而顶部 2 个核心则需要将同样的横向宽接口旋转后垂直引出,这很可能造成了顶部输入端那根醒目的垂直梳状结构。
当然,我这是在完全基于想象进行工程推演,但我认为添加 kmem 的 L2 多路复用器其实没那么难。如果 kernel 取数回退到 DRAM,ANE 的解码性能也不会那么惨烈。这让我猜测,Apple 最初可能从未预见到 L2 常驻张量会变成卷积核,这在 2017 年是一个合理的假设。此外,Apple 倾向于开发孤立的模块化组件,将 1 MiB kmem(只读,能稍微节省 SRAM 布线面积)和 2 MiB tile L2 完全独立开发,大概率也更省事。
4. 这就结束了吗?
DRAM 吞吐量
GPU 比 ANE 更快吗?
Transformer 模型的单 token 解码是权重复用率和计算访存比最糟糕的场景,因为生成每一个 token 都需要将模型的全部权重加载到内存中。然而,如果 ANE 和 GPU 都受限于读取带宽,那么谁的读取带宽更高,谁就能解码更多的 token/s,这与峰值算力无关。
由于 ANE 和 GPU 共享同一块 DRAM,这种情况是公平的:ANE 在 DRAM 访问上未必比 GPU 处于劣势。
- 如果 ANE 的单 token 解码速度慢于 GPU,那是因为其 DMA 控制器无法维持足够的并发请求来打满 DQ。
要测量 ANE 与 GPU 的 DRAM 读取吞吐量,需构建一个受限于读取带宽的工作负载(缓冲区远大于缓存,填充真实的伪随机数据,且仅消费一次),测量执行时间,并对不同的读取尺寸重复上述操作:
\[ s=\frac{\text{change in measured execution time}} {\text{change in payload size}} \]拟合出的斜率回答了一个问题:每增加一次 DRAM 读取需要额外耗费多少执行时间?其倒数即为设备可持续的 DRAM 读取带宽。

- ANE KernelDMA:37.99 GB/s:即 CoreML 内核尺寸除以执行时间。
- ANE TileDMA:59.08 GB/s:即 CoreML tile src 尺寸除以执行时间。
- GPU:77.70 GB/s:通过 shader 读取每个私有 uint4 一次并写入一个数据依赖的校验和,计算 Metal 缓冲区读取量与执行时间之比。
ANE 的内核(操作数 A)带宽上限为 38 GB/s,tile(操作数 B)为 60 GB/s。我们能否以 38 + 60 = 98 GB/s 的速率发出请求,从而逼近 M3 的 100 GB/s DRAM 极限?

不行。实验表明,内核与 tile 组合运行的总时长,等于它们独立运行时长之和。如果请求有任何重叠,较短路径几乎不会增加额外时间;但执行时间(而非吞吐量)是严格单调递增的。
\[T_{AB}=0.001+0.939T_A+0.981T_B\]
由此可见,ANE 的 kernel 和 tile DMA 请求是串行发送的(一次只发一个),也就是说 ANE 的 DRAM 吞吐被双重拖累了:
- 单独跑 kernel 或 tile DMA 的带宽都低于 GPU 的读取 GB/s。
- 而且两者的耗时还是叠加的。
除非……?
ANE 的解码已经顶到 DRAM 瓶颈线了,想要 10 -> 25 tok/s 这种 2.5 倍的激进提升,只能靠把内存流带宽提高约 2.5 倍。
但是……如果能额外加上 50 GB/s 的 kernel DMA 吞吐呢?