Retrospectively Reverse-Engineering Apple's Neural Engine

I stopped moving connected the reverse-engineered Apple Neural Engine (ANE) driver 3 years ago, upon a sad mini realization that the ANE artifact is conscionable not that useful, and I could beryllium doing much useful things, and moved onto upstreaming other, much useful, blocks. The ANE's architecture was excessively opinionated to build a general-purpose accelerator level astir it, and a linux driver efficaciously opening ANE hardware API entree could not broaden the people of workloads it could do. Even macOS only regularly uses their ain ANE to make upsampled preview images successful Finder.

https://github.com/eiln/ane/tree/main

M1 dice shot: https://mastodon.social/@dougall/115149886886125067

The M5 (2025)'s header characteristic was "LLM performance", and they besides conveniently folded the ANE cores wrong the GPU cores — I knew it was coming, but it officially feels for illustration the opening of the extremity for the standalone NPU. So, successful grant of the ANE’s evident demise, we will do thing moreover much useless: spell backmost and reverse-engineer the ANE connected the M1, decorativeness what we started. It's been 3 years (fuck), and I should cognize much than I did erstwhile I first worked connected this.

If the extremity 3 years agone was to make the ANE useful by moving ops connected it; this time, it's much astir mapping the afloat soul architecture — compute, datapath, scheduler, memory, and execution exemplary — because those soul creation decisions uncover the assumptions astir ML workloads that Apple was consenting to perpetrate to silicon first successful the A11 Bionic (2017), and what that says astir the displacement from CNN-era NPUs to today's GPUs moving transformer workloads.

1. Compute

The 16 compute cores are astir apt the slightest absorbing portion of the ANE. Apple primitively targeted dense image-processing CNN workloads, which consists of dense tensor reductions pinch predictable reuse. The M1 ANE compute halfway is simply a ample parallel array of multiply-accumulate (MAC) units, but that unsocial says almost thing astir what workloads it was designed for and accels at.

A convolutional furniture does a dot merchandise betwixt an activation model and learned kernel weights, and attraction does a dot merchandise betwixt a query and cardinal vector. A dot merchandise is simply a dot product, and a MAC does conscionable that. What specialized ANE to the 2017 CNN models is not the MAC, but dataflow surrounding the MACs: erstwhile and wherever MAC inputs and outputs enter, stay, move. The presumption that transformers broke, particularly pinch autoregressive decode, was predictable reuse patterns, which the ANE exploited to designer a dataflow businesslike capable to tally connected phones. The M5 determination confirms that ANE's compute halfway remained still useful for transformers, but wrong a different dataflow.

Still, here's the datapath wrong each of the 16 compute cores:

---

Multiply-Accumulate

ANE has 16 parallel compute cores. Each compute halfway has 128 FP16 (or 256 INT8) parallel multiply-accumulate (MAC) lanes. Each MAC lane performs the recurrence:

\[ s\leftarrow s+a\times b \]

Multiply 2 operands \(a\) and \(b\), and past adhd the merchandise to the moving sum (accumulator).

Repeating the MAC cognition complete T cycles computes a T-term dot product:

\[ s_T=s_0 + \sum_{t=0}^{T-1} a_t \, b_t. \]

A MAC lane frankincense performs a scalar simplification complete time. A 16-core ANE has 2048 parallel MAC lanes,

\[ 128\ \text{lanes/core}\times16\ \text{cores} = 2048\ \text{parallel MAC lanes} \]

So each rhythm performs 2048 parallel reductions spatially, pinch clip being the only simplification axis:

\[ S_T[q,p] = S_0[q,p] + \sum_{t=0}^{T-1} a_t[q,p]\,b_t[q]. \]

An individual MAC lane does not cognize what magnitude of the matrix aliases tensor it is reducing over. It's important to statement that a dot merchandise vs matrix multiplication vs convolution arises from really the operands are mapped and scheduled onto the core. The ANE halfway (with the objection of kernel memory, discussed later) does not encode a 4-channel CNN furniture into the hardware.

---

Internally, the MAC datapath consists of a multiplier, adder, and a 32-bit accumulator register. Each cycle, the adder adds the caller multiplier output pinch the erstwhile sum, which past becomes the caller moving sum.

This feedback way keeps the partial sum successful representation section to the MAC lane, truthful it does not request fetched from an outer representation acold away, betwixt MAC cycles.

Regarding resolution, it does fixed-point simplification pinch FP16 astatine readout. The multiplier is 16-bit, accumulated successful a 32-bit registry arsenic Q16.16, past publication retired arsenic FP16 via sign-extend and etc. Working successful integer (hex) FP16 representation, to probe the accumulator range, build a CoreML ANE programme that computes a dot merchandise pinch a vector of each (1)s, truthful each multiplier results successful a bounded v, but the moving sum successful the accumulator keeps growing:

\[ s=\sum_{i=0}^{255}v=256v. \](v)CPU hexCPU valueANE hexCoreML value

127.9375	0x77ff	32752	0x77ff	32752
128	0x7800	32768	0x7c00	+∞
−128	0xf800	−32768	0xf800	−32768
−128.125	0xf801	−32800	0xfc00	−∞

Since 32768 is itself a valid FP16 connection (0x7800), the ANE's 0x7c00 can't beryllium FP16 output overflow, the clamp happens wrong the accumulator, astatine \(2^{15}\). Thus the accumulator saturates astatine \(2^{15}\), precisely the scope of a signed 32-bit fixed-point worth pinch 16 fractional bits.

---

Nonlinear Activation

For a fused layer, the ANE computes:

\[ y = f(\sum_k x_k w_k + b) \]

Importantly, completed MAC sums provender straight into the post-MAC activation block, avoiding an intermediate representation round-trip. This is imaginable because the activation is pointwise: erstwhile a scalar simplification is complete, its activation depends only connected that scalar and tin beryllium applied immediately.

To find really the ANE implements tanh(), compile a CoreML exemplary containing a azygous TANH activation layer and inspect the resulting compiled hardware registry record (hwx). The coefficient region contains 33 consecutive FP16 words opening astatine 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

Those 33 FP16 words lucifer 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. \]

Now move to RELU activation layer:

activation programNonlinearModelookup coefficients

identity	0	none
ReLU	1	none
tanh	2	33 FP16 words

Thus, mode 2 selects a civilization 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, \]

covering \([0,4]\) pinch spacing \(1/8\). The input maps into the array arsenic \(u=2^R|x|\), so \(R\) sets the knot spacing. The solution is smoother than its 33 bin; I fishy that adjacent entries are linearly interpolated. To test, build an impulse LUT pinch a azygous spike:

\[ T_8=1,\qquad T_k=0\ \text{for }k\ne8,\qquad R=3. \]

Then expanse the input crossed the 2 cells astir \(T_8\). The measured output forms a triangle: magnitude rises linearly from \(0\) astatine \(|x|=7/8\) to \(1\) astatine \(|x|=1\), past falls linearly to \(0\) astatine \(|x|=9/8\).

Thus, we cognize 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\) prime the adjacent entries, and \(\alpha=u-\lfloor u\rfloor\) gives the interpolation weight betwixt them.

---

Scaling and Bias

CoreML besides supports a linear scaling and bias \(ax + b\) transform. I past suspected \(ax + b\) could stock the linear interpolation hardware of mode 2. To confirm, conception a CoreML exemplary pinch a ReLU pinch a changeless standard and offset:

\[ z=4x-2,\qquad y=\operatorname{ReLU}\left(\frac{z}{2}+1\right), \]

If the compiler folds the changeless standard and offset into the convolution:

\[ W'=\frac12W=2,\qquad b'=\frac12b+1=0, \]\[ y=\operatorname{ReLU}(2x). \]

Decoding model.espresso.weights confirms precisely this folded translator connected ReLU:

And the registry record hexdiff shows really bias and activation are fused into the aforesaid post-MAC way astatine compile time:

ProbeTasksBiasModePostScaleModeNonlinearMode

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

Extremely cursed idea: usage nonlinear interpolation to compute an further kernel pass, aliases quantize int8 into int4 weights.

2. Scheduler

The ane driver root code is disappointingly boring. The driver ne'er gives the ANE a CONV, MATMUL, aliases RELU opcode to run. All the neural operations person each already been compiled into a bid watercourse of task descriptors (TDs), and the driver package 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).

https://github.com/eiln/ane/blob/main/ane/src/ane_tm.c#L87

The hardware past owns the submission until completion, and raises an interrupt to the ARM64 halfway erstwhile it's done.

This (boring) bid submission frontend resembles that of a GPU's, deliberation NVIDIA's pushbuffer/PBDMA. The package submits a bid watercourse resident successful memory, and the GPU's bid processor walks complete bid watercourse and dispatches the commands, without knowing what that bid executes.

TM_ADDR and TM_INFO are world staging registers, and TM_PUSH atomically commits that staged motorboat state, fixed that thing happens until TM_PUSH is written ("magic"). TM_INFO successful peculiar stores the full number of descriptors successful the supplied stream:

TM_INFO[31:16] = descriptor_dwords - 1 TM_INFO[15:0] = descriptor_count

TM_INFO registry people maps onto a hardware counter:

Why the "minus 1"? Encoding magnitude - 1 is an RTL-friendly measurement to terminate a zero-based antagonistic retired of the captious path. But statement how, compared to GPU commands which parse a variable-length watercourse of descriptors successful a ringbuffer, ANE only receives the full count, indicating that descriptors are fixed-size.

Task Queue

Going 1 furniture deeper, what's successful a task queue (TQ) that the task head 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 | | authorities | | authorities | | authorities | +---------+ +---------+ +---------+

There's 8 copies of the aforesaid registry artifact (indexed by qid (0…7)), system as:

Each TQ holds:

• (1) Per-TQ scheduling authorities (status, priority, and vacancy)
• (2) Two sets of bid watercourse descriptors, per-TQ (ADDR1/ADDR2, SIZE1/SIZE2, NID1/NID2, and 32 BARs). The 2 slots are surely a ping-pong staging strategy to fto 1 slot execute, while package modifies the different slot.

Notice really TM_PUSH executes a task referenced successful TM_ADDR/TM_SIZE by attaching a qid:

The earthy mentation is that the descriptor watercourse specifies what task to run, while the qid selects the motorboat discourse the descriptor runs under. The resident TQ discourse (BAR, NID) is overmuch for illustration a GPU hardware channel. Here's my driver populating a azygous TQ to motorboat it:

https://github.com/eiln/ane/blob/main/ane/src/ane_tm.c#L70

The only point important present is the 32-entry BAR array (base reside register). We'll get into task descriptors next, but the compiled ANE bid watercourse only references virtual addresses by comparative offsets, and BAR provides the guidelines IOVA (IOMMU peripheral virtual address) relocation address. An ANE virtual reside entree needs a hard-coded BAR guidelines offset supplied astatine compile time, meaning it lacks GPU-style load/store instructions that dynamically rumor load/stores from virtual address.

---

Task Descriptor

The task head walks complete and executes concatenation of fixed-size task descriptors:

What's successful 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

Important is that a TD is not an executable instruction stream. ANE has nary ISA. TD is simply a series of "ControlDMA" (I made this sanction up) burst-write packets writes to the ANE's hardware configuration registers, specified arsenic input dimension, input/output address, activation function. Each ControlDMA packet consists of a 32-bit transportation connection followed by N consecutive 32-bit registry values:

31 26 25 2 1 0 +-----------------------+-----------------------------+----+ | registry count minus 1| first registry guidelines scale | 00 | +-----------------------+-----------------------------+----+

Notice the “minus 1” termination count again. ControlDMA is simply a elastic unidirectional DMA motor that copies N 32-bit words from IOMMU virtual DRAM into the ANE’s beingness registry space. For example, KernelDMASrc's packet header successful TD is 0xf401f800:

This is not a LOAD_WEIGHTS instruction. It's copying 0xf4 aliases 62 consecutive words into the KernelDMA registry offset starting astatine 0x1f800. And those KernelDMA configuration values tin show KernelDMA wherever to load the weights from.

Start byteSectionInformation

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; section 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

Since each conception writes to 1 MMIO registry block, TD divides cleanly into ANE's datapath sections:

Starting addressSizeBlock nameWhat

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 registry 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 registry bank
0x26bc24000	0x1000	Task Manager	Task submission, execution state, events, and completion
0x26bc25000	0x1000	Task Queues	Eight queues containing TD stacks, NIDs, priorities, and petition pointers

A TD is efficaciously a serialized register-file dump of the ANE’s datapath registers. Each “ANE program” is simply the configuration for 1 walk done the datapath. We tin configure really the fixed datapath operates (subject to the knobs it exposes), but not what operations the datapath is tin of performing, aliases really those operations are sequenced.

When the "magic" atomic connection is written to task head to execute a TD, roughly, the series of what happens:

0. ControlDMA copies TD into the configuration registers.
1. KernelDMA copies kernel W into kernel representation (KMem).
2. TileDMA copies input \(X\) from DRAM into L2.
3. Each MAC halfway reduces a statement by its weights, producing 1 statement of \(Y\).
4. Steps 2–3 repetition for each rows of \(X\).
5. Postprocessing is applied, and the completed results are stored successful L2.
6. TileDMADst copies \(Y\) from L2 backmost to DRAM.

ANE is simply a fixed-function dataflow engine, not a GPU executing arbitrary instructions. The TD configures a domain-specific datapath. Constraining the hardware interface usually intends smaller area, deterministic movement, little latency, and little powerfulness drawn. ANE's compiler tin explicitly schedule what the tensors do, but that besides intends the compiler must explicitly schedule what the tensors do. This is simply a tradeoff, but a justified one: we usually cognize what the exemplary looks for illustration astatine compile time. Dynamic execution is not what limits ANE. ANE's processor interface is comparatively generic, and it simply launches tasks, and the tasks tin picture transformers.

For illustration making tensor sizes fixed astatine compile clip does not mean it can't grip variable-length tensors: for example, a increasing KV cache tin beryllium traversed by looping complete the size, and dispatch overhead is negligible comparative to the elephant successful the room here, that is, memory-streaming bandwidth. What really shaped ANE for CNNs complete transformers is representation movement.

3. Memory

Roofline

It's ever bully to place our existent slowest link, truthful we tin optimize what really matters.

Apple’s unified representation lets the ANE entree buffers from the strategy DRAM excavation accessible by the CPU and GPU. It does not mean the ANE zero-copy streams straight retired of that DRAM pool. ANE must first transcript immoderate representation into its section "ANE memory" aliases SRAM. Any bandwidth-limited task will frankincense beryllium constricted by ANE's section representation streaming throughput.

M1 ANE reports \(11\text{ TOP/s}\) astatine \(68\text{ GB/s}\) astatine strategy DRAM bandwidth. A MAC performs 2 operations but consumes 2 FP16 operands, aliases 4 bytes:

\[ \frac{2\text{ OP}}{4\text{ bytes}} =0.5\text{ OP/byte}. \]

If each 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}. \]

which is complete 300x times the reported \(68\text{ GB/s}\) strategy DRAM capacity. Thus highest ANE MAC throughput could beryllium reached by fetching from immoderate section ANE representation reservoir, and reusing it.

Restated, the M1 ANE’s \(11\text{ TOP/s}\) astatine \(68\text{ GB/s}\) number sets the roofline ridge point:

\[ \frac{11\text{ TOP/s}}{68\text{ GB/s}} = 162\text{ OP/byte}. \]

Each byte fetched from DRAM must support, connected average, astatine slightest 162 operations for DRAM bandwidth to extremity being the limiter. Equivalently, the workload must supply capable on-chip reuse to execute an arithmetic strength of astatine slightest 162 OP/byte DRAM traffic. Below the 162:1 ratio, speeding up compute won't summation decoded token/s.

---

Memory Hierarchy

Even if (average) DRAM bandwidth were sufficient, ANE does not publication DRAM straight for galore reasons, including DRAM deterministic timing, beingness routing, shared traffic, etc. If ANE's postulation competes connected AXI the CPU, GPU, display, and etc, it cannot supply deterministic timing to the MACs. Also, if 16 cores devour immoderate input tile, we do not want to initiate 16 identical DRAM transfers. "ANE section memory" would let intermediate activation produced by 1 cognition beryllium consumed by the adjacent alternatively of walking to DRAM and back.

Apple had respective ways to shape section representation hierarchy. The multiply-accumulate patent describes the the information buffer paths astir the array.

Unified DRAM │ ▼ ┌───────────────────────────────────────────────┐ │ shared ANE L2 memory, 2 MiB │ └────────┬───────────────┬───────────────────┬──┘ │ │ │ ▼ ▼ ▼ ┌────────────┐ ┌────────────┐ ... ┌────────────┐ │ halfway 0 │ │ halfway 1 │ │ halfway N │ │ ┌────────┐ │ │ ┌────────┐ │ │ ┌────────┐ │ │ │ L1 │ │ │ │ L1 │ │ │ │ L1 │ │ │ └────────┘ │ │ └────────┘ │ │ └────────┘ │ │ ┌────────┐ │ │ ┌────────┐ │ │ ┌────────┐ │ │ │ KMem │ │ │ │ KMem │ │ │ │ KMem │ │ │ │ 64 KiB │ │ │ │ 64 KiB │ │ │ │ 64 KiB │ │ │ └────────┘ │ │ └────────┘ │ │ └────────┘ │ └────────────┘ └────────────┘ └────────────┘

• KMem: 16x per-core 64 KiB "L1" SRAM for kernel. Total 1 MiB.
• L1: 16x per-core MAC input "L1" staging area.
• L2: 1x shared 2 MiB L2 crossed each cores.

I'm not gonna dress for illustration I've ne'er decompiled shit. The ANE ARM64 firmware's task-debug routine (1) dumps 0x10000 bytes from KMem indices 0 done 15 (2) past dumps 1 abstracted 0x200000-byte L2 dump:

---

DMA Engines

┌──────────┐ ┌────────────────────────────────────── ANE ──────────────────────────────────────┐ │ │ │ │ │ │ │ ┌─────────────┐ ┌────────────────────┐ │ │ ├──────►│ │ Control DMA │─────►│ Hardware Registers │ │ │ │ │ └─────────────┘ └────────────────────┘ │ │ │ │ ┌─────────── MAC ────────────┐ │ │ │ │ ┌──────────────┐ │ ┌──────────────────────┐ │ │ │ ├──────►│ │ KernelDmaSrc │─────────────────────────────►│ │ Kernel Memory │ │ │ │ DRAM │ │ └──────────────┘ │ └──────────┬───────────┘ │ │ │ │ │ │ ▼ │ │ │ │ │ ┌──────────────┐ ┌────────────────┐ │ ┌──────────────────────┐ │ │ │ ├──────►│ │ TileDmaSrc │──────►│ L2 │◄───►│ │ MAC Array │ │ │ │ │ │ └──────────────┘ │ Tile Memory │ │ └──────────────────────┘ │ │ │ │ │ ┌──────────────┐ │ │ └────────────────────────────┘ │ │ │◄──────┤ │ TileDmaDst │◄──────│ │ │ │ │ │ └──────────────┘ └────────────────┘ │ │ │ │ │ └──────────┘ └─────────────────────────────────────────────────────────────────────────────────┘

There are 3 DMA engines to/from the MACs:

• KernelDMASrc[0..15]: sixteen logical coefficient lanes transcript weights from DRAM into the matching per-core KMem banks.
• TileDMASrc: copies tiles from DRAM to L2.
• TileDMADst: copies tiles from L2 to DRAM.

The sixteen KernelDMASrc registry lanes do not by themselves beryllium sixteen physically independent DMA beforehand ends. A lane-scaling research does show that the logical lanes make concurrent progress: 32 valid tasks pinch 1 MiB aggregate coefficients per enabled lane return fundamentally the aforesaid clip astatine one and sixteen lanes. They whitethorn still merge into shared request-generation, crossbar, cache, and DRAM arbitration.

The 3 tile/kernel DMA engines are chill. You proviso it immoderate src, dst, and size, and it will transportation that artifact of memory. However the beingness of 3 DMA engines reveals immoderate absorbing assumptions:

• Kernel gets a dedicated path, astatine all. Splitting kernel vs tile is simply a committedness that kernel is simply a chopped operand people pinch a different lifecycle.
• KernelDMA is unidirectional load-only. They thought that weights would not beryllium written back.
• TileDMA is bidirectional and shared. Intermediate tile outputs tin beryllium fed backmost without a round-trip to DRAM.
• Kernel representation is backstage per halfway (replicated 16x). They thought that weights would beryllium reused galore times per-core.
• New kernels must beryllium loaded from DRAM, and not L2. They astir apt didn't deliberation to watercourse ample weights that exceeds 64 KiB x 16.

---

Kernel L1

A convolution is not a symmetric multiply-add of \(A\) and \(B\). A convolution slides the aforesaid kernel \(w[k]\) crossed the input:

\[ y[p]=\sum_k x[p+k]\,w[k]. \]

Shoutout 2k2 textbook

Notice really the 2 pat kernel \(w_0\) and \(w_1\) tin beryllium loaded erstwhile and past reused crossed the full long of the output:

So the kernel tin reside successful section memory, and reused, while caller inputs displacement in; which is what Apple's datapath patent describes (https://patents.google.com/patent/US20190340491A1/en).

The intended dependable authorities is:

┌── activation ───────┐ DRAM → L2 ──────────┤ × → MAC └── output ←──────────┘ DRAM → KMem ─────────────────── coefficient

If designing an ASIC to churn convolutions — wherever the kernel isn't thing that's often and dynamically streamed successful (ahem KV) — it's a nary brainer to utilization the operand asymmetry and trim kernel representation movement, because we've shown that the ANE is still profoundly successful the representation throughput-limited (162:1 ratio) region earlier the MACs tin beryllium saturated.

But why not let L2 → KMem? It seems easier, even, to unify tile and kernel paths.

(1) So kernel postulation doesn't compete pinch tile for L2 bandwidth? I reason that L2 contention was not why Apple omitted the L2 -> KMem path: resident kmem exists astatine each because kernels were expected to beryllium loaded infrequently. If kmem postulation is negligible successful dependable state, it could simply beryllium fixed little privilege than tile L2 accesses.

(2) Because kernel L2 distribution adds complexity? Apple already distributes L2 endpoints to each core; I reason that a kernel way riding the aforesaid tile way is not that bad.

ANE's beingness layout is centered astir (literally) the halfway L2 SRAM rectangle, pinch (7+7) cores on each broadside of the rectangle, 2 cores connected the apical side, and shared power logic connected the bottommost side. The 7+7 broadside cores return L2 ingress horizontally from the cyan vertical trunk; but the 2 apical cores request the aforesaid horizontal wide interface rotated and escaped vertically, which apt produces that conspicuous vertical comb successful the apical ingress.

Granted I americium afloat armchair engineering here, but adding the kmem L2 mux, I argue, really could not person been that bad. ANE decode capacity would not person been arsenic tanked if kernel fetches spell backmost to DRAM. It makes maine deliberation that Apple simply ne'er expected L2-resident tensors to go kernels, which was a valid presumption successful 2017. And Apple besides likes processing isolated modular modules, astir apt would've been easier to wholly isolate improvement of the 1 MiB kmem (read-only, which would prevention a small area successful SRAM routing) and 2 MiB tile L2.

4. Is it over?

DRAM Throughput

Is the GPU faster than the ANE?

Transformer single-token decode is the worst lawsuit for weight reuse and compute per representation ratio, because we request to watercourse the full model’s worthy of weights to make 1 token. However, if some the ANE and GPU are read-bandwidth bound, whichever 1 that has higher publication bandwidth will decode much token/s, sloppy of highest compute capacity.

Now since ANE and GPU stock the aforesaid DRAM, it's adjacent game: the ANE is not needfully penalized successful DRAM entree compared to the GPU.

• If ANE is slower than GPU astatine azygous token decode, it's because its DMA controller cannot support capable parallel requests to saturate DQ.

To measurement ANE vs GPU DRAM publication throughput, make a read-bandwidth-bound workload (buffers are overmuch larger than the caches, filled pinch existent pseudo-random data, and consumed once), and measurement the execution time, and repetition for different publication sizes:

\[ s=\frac{\text{change successful measured execution time}} {\text{change successful payload size}} \]

The fitted slope answers: really overmuch further execution clip does 1 other DRAM publication require? The reciprocal is the device's sustained DRAM publication bandwidth.

• 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: Metal buffer sounds per execution time, via a shader that sounds each backstage uint4 erstwhile and writes a data-dependent checksum.

ANE's kernel (operand A) maxes retired astatine 38 GB/s, and tile (operand B) astatine 60 GB/s. Can we rumor 38 + 60 = 98 GB/s to deed M3's 100 GB/s DRAM ceiling?

No. Experiment shows that the kernel+tile mixed runtime matched the sum of the isolated runtimes. If the requests overlapped astatine all, past the shorter way would lend small aliases nary further time, but execution clip (not throughput) is strictly monotonic.

\[T_{AB}=0.001+0.939T_A+0.981T_B\]

Thus, ANE's kernel and tile DMA requests are sent serially (one astatine a time), meaning ANE DRAM throughput is double-fucked:

• Both isolated kernel and tile DMA are little than GPU's publication GB/s.
• Kernel and tile DMA times are besides additive.

---

Unless…?

With ANE decode pinned to the DRAM roofline, a drastic 2.5x betterment for illustration 10 -> 25 tok/s tin only travel from ~2.5× higher representation streaming bandwidth.

Getting 50 GB/s Back Out of the ANE

But … what if we could adhd 50 GB/s of further kernelDMA throughput ?