Retrospectieve Reverse-Engineering van Apple's Neural Engine

Drie jaar geleden ben ik gestopt met het werken aan de reverse-engineered Apple Neural Engine (ANE) driver. Dit kwam voort uit het droevige besef dat het ANE-blok simpelweg niet zo nuttig is en dat ik mijn tijd beter kon besteden aan het upstreamen van andere, nuttigere blokken. De architectuur van de ANE was te specifiek om er een algemeen acceleratorplatform omheen te bouwen, en een Linux-driver die effectief toegang geeft tot de ANE-hardware API zou de klasse van workloads die het kan uitvoeren niet verbreden. Zelfs macOS gebruikt de eigen ANE regelmatig alleen voor het genereren van upsampled preview-afbeeldingen in Finder.

De hoofdfeature van de M5 (2025) was "LLM-prestaties", waarbij ze bovendien de ANE-cores handig in de GPU-cores hebben geïntegreerd. Ik wist dat dit eraan kwam, maar officieel voelt dit als het begin van het einde voor de standalone NPU. Ter ere van het schijnbare einde van de ANE zullen we iets nog nutteloosers doen: teruggaan en de ANE op de M1 reverse-engineeren om te voltooien waar we aan begonnen zijn. Er zijn drie jaar voorbijgegaan, en ik zou nu meer moeten weten dan toen ik hier voor het eerst aan werkte.

Waar het doel drie jaar geleden was om de ANE nuttig te maken door er operaties op uit te voeren, gaat het dit keer meer om het in kaart brengen van de volledige interne architectuur — compute, datapath, scheduler, geheugen en executiemodel. Deze interne ontwerpbeslissingen onthullen namelijk de aannames over ML-workloads waar Apple in 2017 (met de A11 Bionic) voor koos in silicon, en wat dat zegt over de verschuiving van CNN-tijdperk NPUs naar de huidige GPU's die transformer-workloads draaien.

1. Compute

De 16 compute-cores zijn waarschijnlijk het minst interessante deel van de ANE. Apple richtte zich oorspronkelijk op dense image-processing CNN-workloads, die bestaan uit dense tensor-reducties met voorspelbare hergebruikpatronen. De compute-core van de M1 ANE is een groot parallel array van multiply-accumulate (MAC) units, maar dat alleen zegt weinig over de workloads waarvoor het is ontworpen.

Een convolutionele laag voert een dot-product uit tussen een activatievenster en geleerde kernelweights, en attention voert een dot-product uit tussen een query- en key-vector. Een dot-product is in beide gevallen een dot-product, en een MAC doet precies dat. Wat de ANE specifiek maakte voor de CNN-modellen van 2017, is niet de MAC zelf, maar de dataflow rondom de MACs: wanneer en waar MAC-inputs en -outputs binnenkomen, blijven en bewegen. De aanname van voorspelbare hergebruikpatronen — die door transformers, vooral bij autoregressieve decode, werd doorbroken — stelde de ANE in staat een dataflow te architectureren die efficiënt genoeg was om op telefoons te draaien. De beslissing bij de M5 bevestigt dat de compute-core van de ANE nog steeds nuttig was voor transformers, maar binnen een andere dataflow.

Dit is de datapath binnen elk van de 16 compute-cores:

┌────────────────────── core ─────────────────────┐
│ ┌───────── 256× MACs ─────────┐  ┌────────────┐ │
│ │ MAD ─► add ─► accumulator   │─►│ activation │ │
│ │        ▲          │         │  └────────────┘ │
│ │        └──────────┘         │                 │
│ └─────────────────────────────┘                 │
└─────────────────────────────────────────────────┘

Multiply-Accumulate

De ANE heeft 16 parallelle compute-cores. Elke compute-core heeft 128 FP16 (of 256 INT8) parallelle multiply-accumulate (MAC) lanes. Elke MAC-lane voert de volgende recurrentie uit:

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

Hierbij worden twee operanden \(a\) en \(b\) vermenigvuldigd, waarna het product wordt toegevoegd aan de lopende som (accumulator). Door de MAC-operatie over \(T\) cycli te herhalen, wordt een dot-product met \(T\) termen berekend:

\[ sT = s0 + \sum{t=0}^{T-1} at \, b_t \]

Een MAC-lane voert dus een scalaire reductie uit over tijd. Een 16-core ANE heeft 2048 parallelle MAC-lanes:

\[ 128\ \text{lanes/core} \times 16\ \text{cores} = 2048\ \text{parallelle MAC-lanes} \]

Elke cyclus worden er dus 2048 parallelle reducties ruimtelijk uitgevoerd, waarbij tijd de enige reductie-as is:

\[ ST[q,p] = S0[q,p] + \sum{t=0}^{T-1} at[q,p]\,b_t[q] \]

Een individuele MAC-lane weet niet over welke dimensie van de matrix of tensor hij reduceert. Het is belangrijk om op te merken dat het onderscheid tussen een dot-product, matrixvermenigvuldiging of convolutie voortvloeit uit de manier waarop de operanden op de core worden gemapt en gescheduld. De ANE-core (met uitzondering van het kernelgeheugen) codeert geen 4-kanaals CNN-laag direct in de hardware.

Intern bestaat de MAC-datapath uit een multiplier, een adder en een 32-bit accumulator-register. Elke cyclus voegt de adder de nieuwe multiplier-output toe aan de vorige som, die vervolgens de nieuwe lopende som wordt.

operand a ──┐   ┌────────────┐   p[31:0]   ┌──────────────┐   s_next[31:0]  ┌─────────────┐
├──►│ MULTIPLIER │────────────►│ 32-BIT ADDER │────────────────►│ ACCUMULATOR │
operand b ──┘   └────────────┘             └──────▲───────┘                 └──────┬──────┘
                                                 │ s[31:0]
                                                 └────────────────────────────────┘

Deze feedback-path houdt de partiële som in het geheugen lokaal bij de MAC-lane, zodat deze tussen MAC-cycli niet uit een extern, ver weg gelegen geheugen hoeft te worden opgehaald.

Wat betreft de resolutie: er wordt gebruikgemaakt van fixed-point reductie met FP16 bij het uitlezen. De multiplier is 16-bit, geaccumuleerd in een 32-bit register als Q16.16, en vervolgens uitgelezen als FP16 via sign-extend en dergelijke. Om het bereik van de accumulator te testen, kan een CoreML ANE-programma worden gebouwd dat een dot-product berekent met een vector van alleen maar (1)en. Hierdoor resulteert elke multiplier in een begrensde waarde \(v\), maar de lopende som in de accumulator blijft groeien:

\[ s = \sum_{i=0}^{255} v = 256v \]

(v)CPU hexCPU waardeANE hexCoreML waarde
127.93750x77ff327520x77ff32752
1280x7800327680x7c00$+\infty$
−1280xf800−327680xf800−32768
−128.1250xf801−328000xfc00$-\infty$

Omdat 32768 zelf een geldig FP16-woord is (0x7800), kan de 0x7c00 van de ANE geen FP16-output-overflow zijn; de clamp gebeurt binnen de accumulator, bij \(2^{15}\). De accumulator verzadigt dus bij \(2^{15}\), precies het bereik van een signed 32-bit fixed-point waarde met 16 fractionele bits.

Niet-lineaire Activering

Voor een gefuseerde laag berekent de ANE:

\[ y = f\left(\sumk xk w_k + b\right) \]

Belangrijk is dat voltooide MAC-sommen direct in het post-MAC activatieblok terechtkomen, waardoor een tussenliggende geheugenronde wordt vermeden. Dit is mogelijk omdat de activatie pointwise is: zodra een scalaire reductie is voltooid, hangt de activatie alleen af van die scalar en kan deze direct worden toegepast.

Om te bepalen hoe de ANE tanh() implementeert, kan een CoreML-model met een enkele TANH-activatielaag worden gecompileerd en het resulterende hardware-registerbestand (hwx) worden geïnspecteerd. De coëfficiëntregio bevat 33 opeenvolgende FP16-woorden beginnend bij 0x4288.

Deze 33 FP16-woorden komen overeen met 33 IEEE LE FP16 gekwantiseerde samples van \(\tanh(x)\):

\[ Ti = \operatorname{round}{16}\left(\tanh(i/8)\right), \quad i=0,1,\ldots,32 \]

Kijken we nu naar de RELU-activatielaag:

ActiveringsprogrammaNiet-lineairModusLookup coëfficiënten
identity0none-
ReLU1none-
tanh233 FP16 woorden-

Modus 2 selecteert dus een aangepaste lookup-tabel (LUT) met 33 vermeldingen. 33 punten definiëren 32 intervallen. Met \(R=3\) zijn de knopen:

\[ x_i = \frac{i}{8}, \quad i=0,\ldots,32 \]

Dit bestrijkt \([0,4]\) met een tussenruimte van \(1/8\). De input wordt in de tabel gemapt als \(u=2^R|x|\), waarbij \(R\) de afstand tussen de knopen bepaalt. De resolutie is vloeiender dan de 33 bins suggereren; ik vermoed dat aangrenzende vermeldingen lineair worden geïnterpoleerd. Om dit te testen, kan een impuls-LUT met één enkele spike worden gebouwd:

\[ T8=1, \quad Tk=0 \text{ voor } k \ne 8, \quad R=3 \]

Wanneer de input over de twee cellen rond \(T_8\) wordt gescand, vormt de gemeten output een driehoek: de magnitude stijgt lineair van \(0\) bij \(|x|=7/8\) naar \(1\) bij \(|x|=1\), en daalt vervolgens lineair naar \(0\) bij \(|x|=9/8\).

Hieruit weten we dat modus 2 een piecewise-lineaire LUT met 33 vermeldingen implementeert. \(R\) schaalt de input naar LUT-coördinaten:

\[ u = 2^R|x| \]

De afstand tussen de knopen is dus \(\Delta x = 2^{-R}\). \(\lfloor u\rfloor\) en \(\lceil u\rceil\) selecteren de aangrenzende vermeldingen, en \(\alpha = u - \lfloor u\rfloor\) geeft het interpolatiegewicht tussen hen.

Schaling en Bias

CoreML ondersteunt ook een lineaire schaling en bias-transformatie \(ax + b\). Ik vermoedde dat \(ax + b\) gebruik zou kunnen maken van de lineaire interpolatie-hardware van modus 2. Om dit te bevestigen, kan een CoreML-model worden gemaakt met een ReLU met een constante schaal en offset:

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

Als de compiler de constante schaal en offset in de convolutie voegt:

\[ W' = \frac{1}{2}W = 2, \quad b' = \frac{1}{2}b + 1 = 0 \implies y = \operatorname{ReLU}(2x) \]

Het decoderen van model.espresso.weights bevestigt precies deze gevouwen transformatie op ReLU:

  • Auteur convolutie: \(W = 4, b = -2\)
  • Activatie affine: \(s = 0.5, c = 1\)
  • Gecompileerde convolutie: \(W' = 2, b' = 0\)

Het registerbestand laat zien hoe bias en activatie tijdens het compileren in hetzelfde post-MAC pad worden gefuseerd:

ProbeTasksBiasModePostScaleModeNonlinearMode
Plain convolution1000
Explicit Core ML Bias1100
Bias + ReLU1101
Bias + tanh1102

Een zeer onorthodox idee: gebruik niet-lineaire interpolatie om een extra kernel-pass te berekenen, of kwantiseer int8 naar int4 weights.

2. Scheduler

De broncode van de ANE-driver is teleurstellend saai. De driver geeft de ANE nooit een CONV, MATMUL of RELU opcode om uit te voeren. Alle neurale operaties zijn al gecompileerd in een command-stream van task descriptors (TD's). De driver-software laadt simpelweg de taak in het geheugen, stelt de pointer naar de ondoorzichtige task-blob in via (TMADDR, TMSIZE) en verzendt de taak door de "doorbell" te luiden (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
}

De hardware beheert de verzending vervolgens tot voltooiing en stuurt een interrupt naar de ARM64-core wanneer het klaar is.

Deze (saaie) command-submission frontend lijkt op die van een GPU, denk aan NVIDIA's pushbuffer/PBDMA. De software verzendt een command-stream die in het geheugen staat, en de command processor van de GPU loopt door deze stream en dispatcht de commando's zonder te weten wat dat commando exact uitvoert.

TMADDR en TMINFO zijn globale staging-registers, en TMPUSH commit die staging-status atomair. TMINFO slaat specifiek het totaal aantal descriptors in de geleverde stream op:

  • TMINFO[31:16] = descriptor\dwords - 1
  • TMINFO[15:0] = descriptor\count

Het TM_INFO register mapt natuurlijk naar een 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

Waarom "minus 1"? Het coderen van lengte - 1 is een RTL-vriendelijke manier om een nul-gebaseerde counter te beëindigen buiten het kritieke pad. Merk op dat de ANE, in tegenstelling tot GPU-commando's die een variabele lengte van descriptors in een ringbuffer parsen, alleen het totaal aantal ontvangt, wat aangeeft dat descriptors een vaste grootte hebben.

Task Queue

Wat zit er in een task queue (TQ) waar de task manager uit selecteert?

CPU / driver ------>|   Task Manager   |
                    |
                    | schedule / fetch |
                    | / dispatch       |
                    +--------+---------+
                             |
                             v
            +---------------------------------------+
            |                |                      |
            v                v                      v
      +---------+      +---------+            +---------+
      | TQ 0    |  ... | TQ 3    |     ...    | TQ 7    |
      | BAR[32] |      | BAR[32] |            | BAR[32] |
      | NID     |      | NID     |            | NID     |
      | state   |      | state   |            | state   |
      +---------+      +---------+            +---------+

Er zijn 8 kopieën van hetzelfde registerblok (geïndexeerd door qid 0…7), gestructureerd als:

  • 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
  • volgende qid: +0x148

Elke TQ bevat:

  1. Per-TQ scheduling state (status, priority, en vacancy).
  2. Twee sets van command stream descriptors per TQ. De twee slots zijn waarschijnlijk een ping-pong staging schema, zodat één slot kan uitvoeren terwijl de software het andere slot wijzigt.

De natuurlijke interpretatie is dat de descriptor stream specificeert welke taak moet worden uitgevoerd, terwijl de qid de launch-context selecteert waaronder de descriptor draait. De resident TQ-context (BAR, NID) lijkt veel op een GPU-hardwarekanaal.

Het enige dat hier belangrijk is, is de 32-entry BAR-tabel (base address register). De gecompileerde ANE-command-stream verwijst naar virtuele adressen via relatieve offsets, en BAR biedt het basis IOVA (IOMMU peripheral virtual address) relocatie-adres. Een ANE virtuele adres-toegang heeft een hard-coded BAR-basisoffset nodig die bij compilatie is meegegeven, wat betekent dat het geen GPU-stijl load/store instructies heeft die dynamisch load/stores vanuit een virtueel adres kunnen uitvoeren.

Task Descriptor

De task manager loopt door en voert een keten van fixed-size task descriptors uit. Wat zit er in elke TD? Een TD is geen uitvoerbare instructiestroom; de ANE heeft geen ISA. Een TD is een reeks "ControlDMA" burst-write pakketten die naar de hardware-configuratieregisters van de ANE schrijven, zoals input-dimensie, input/output-adres en activatiefunctie. Elk ControlDMA-pakket bestaat uit een 32-bit transfer-woord gevolgd door \(N\) opeenvolgende 32-bit registerwaarden.

ControlDMA Word Formaat:

  • Bits 31-26: register count minus 1
  • Bits 25-2: first register base index
  • Bits 1-0: 00

ControlDMA is een flexibele unidirectionele DMA-engine die \(N\) 32-bit woorden van IOMMU virtueel DRAM naar de fysieke registerruimte van de ANE kopieert.

TD Secties:

Start byteSectieInformatie
0x000HeaderDependencies, chaining, en BAR selectors
0x028KernelDMASrc0xf401f800; 16 coefficient-DMA lanes
0x124Common0x3c000000; tensor en convolutie geometrie
0x168TileDMASrc0x6c013800; activatie-source DMA
0x1dcL20x44004800; lokale source/result configuratie
0x228Processing engine0x0c008800; PE configuratie
0x23cNeural engine0x1000c800; MAC en post-processing configuratie
0x254TileDMADst0x18017800; result-destination DMA
0x274EndTotaal 628 bytes

Omdat elke sectie naar één MMIO-registerblok schrijft, kan de TD netjes worden onderverdeeld in de datapath-secties van de ANE:

StartadresGrootteBloknaamFunctie
0x26bc000000x4000CommonBroadcast configuratie selector
0x26bc040000x4000L2L2 backing/register aperture
0x26bc080000x4000PEProcessing-element configuratie
0x26bc0c0000x4000NE / MACKernel format, MAC, bias, schaling, en nonlinear controls
0x26bc100000x3000OnbekendNiet geïdentificeerd registerbank
0x26bc130000x4000Tile DMA sourceInput-tile adressen, strides, formats, en DMA controls
0x26bc170000x4000Tile DMA destinationOutput-tile adressen, strides, formats, en DMA controls
0x26bc1b0000x4000OnbekendNiet geïdentificeerde configuratie- en tunable registers
0x26bc1f0000x4000KernelKernel backing / kernel DMA-source aperture
0x26bc230000x1000OnbekendNiet geïdentificeerd registerbank
0x26bc240000x1000Task ManagerTask submission, execution state, events, en completion
0x26bc250000x1000Task QueuesAcht queues met TD stacks, NIDs, priorities, en request pointers

Een TD is in feite een geserialiseerde register-file dump van de datapath-registers van de ANE. Elk "ANE-programma" is simpelweg de configuratie voor één pass door de datapath. We kunnen configureren hoe de vaste datapath opereert, maar niet welke operaties de datapath in staat is uit te voeren of hoe deze operaties worden gesequenced.

Wanneer het "magische" atomaire woord naar de task manager wordt geschreven om een TD uit te voeren, gebeurt er grofweg het volgende:

  1. ControlDMA kopieert de TD naar de configuratieregisters.
  2. KernelDMA kopieert kernel \(W\) naar het kernelgeheugen (KMem).
  3. TileDMA kopieert input \(X\) van DRAM naar L2.
  4. Elke MAC-core reduceert een rij met zijn weights, wat één rij van \(Y\) produceert.
  5. Stappen 2–3 herhalen voor alle rijen van \(X\).
  6. Post-processing wordt toegepast, en de voltooide resultaten worden in L2 opgeslagen.
  7. TileDMADst kopieert \(Y\) van L2 terug naar DRAM.

De ANE is een fixed-function dataflow engine, geen GPU die willekeurige instructies uitvoert. De TD configureert een domeinspecifieke datapath. Het beperken van de hardware-interface betekent meestal een kleiner oppervlak, deterministische beweging, lagere latentie en minder stroomverbruik. De compiler van de ANE kan expliciet schedulen wat de tensors doen, maar dat betekent ook dat de compiler alles expliciet moet regelen. Dit is een gerechtvaardigde tradeoff: we weten meestal bij compilatie al hoe het model eruitziet. Dynamische executie is niet wat de ANE beperkt. De processor-interface van de ANE is relatief generiek en start simpelweg taken, en die taken kunnen transformers beschrijven.

Wat de ANE echter echt vormgaf voor CNN's in plaats van transformers, is de geheugenverplaatsing.

3. Geheugen

Roofline

Het is altijd goed om de traagste schakel te identificeren. Apple's unified memory laat de ANE buffers benaderen uit de system DRAM-pool. Dit betekent echter niet dat de ANE direct zero-copy streamt uit die pool. De ANE moet eerst elk geheugen kopiëren naar zijn lokale "ANE-geheugen" of SRAM. Elke bandwidth-limited taak zal dus beperkt worden door de lokale geheugen-streaming throughput van de ANE.

De M1 ANE rapporteert \(11\text{ TOP/s}\) bij \(68\text{ GB/s}\) system DRAM-bandbreedte. Een MAC voert twee operaties uit maar verbruikt twee FP16-operanden (4 bytes):

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

Als elke MAC-operand uit DRAM zou streamen, zou \(11\text{ TOP/s}\) een streamingrate van \(22\text{ TB/s}\) vereisen, wat meer dan 300x de gerapporteerde \(68\text{ GB/s}\) is. De piek-MAC-throughput kan dus alleen worden bereikt door data uit een lokale ANE-geheugenreservoir te halen en deze te hergebruiken.

De M1 ANE's \(11\text{ TOP/s}\) bij \(68\text{ GB/s}\) bepaalt het roofline-ridge-punt:

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

Elke byte die uit DRAM wordt gehaald, moet gemiddeld ten minste 162 operaties ondersteunen voordat de DRAM-bandbreedte stopt de limiter te zijn. Onder deze 162:1 ratio zal het versnellen van de compute de decoded tokens/s niet verhogen.

Geheugenhiërarchie

Zelfs als de gemiddelde DRAM-bandbreedte voldoende zou zijn, leest de ANE DRAM niet direct. Dit komt door deterministische timing, fysieke routing en gedeeld verkeer. Als ANE-verkeer op de AXI concurreert met de CPU, GPU en display, kan het geen deterministische timing aan de MACs bieden. Bovendien willen we niet dat 16 cores die dezelfde input-tile consumeren, 16 identieke DRAM-transfers initiëren. "Lokaal ANE-geheugen" maakt het mogelijk dat intermediate activaties geproduceerd door één operatie direct door de volgende worden geconsumeerd, zonder een trip naar DRAM.

De geheugenhiërarchie is als volgt opgebouwd:

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: 16x per-core 64 KiB "L1" SRAM voor de kernel. Totaal 1 MiB.
  • L1: 16x per-core MAC input "L1" staging area.
  • L2: 1x shared 2 MiB L2 over alle cores.

DMA-engines

Er zijn drie DMA-engines naar/van de MACs:

  1. KernelDMASrc[0..15]: Zestien logische coefficient-lanes kopiëren weights van DRAM naar de bijbehorende per-core KMem-banks.
  2. TileDMASrc: Kopieert tiles van DRAM naar L2.
  3. TileDMADst: Kopieert tiles van L2 naar DRAM.

Een experiment met lane-scaling laat zien dat de logische lanes gelijktijdig vorderen: 32 geldige taken met 1 MiB aggregate coefficients per enabled lane kosten nagenoeg dezelfde tijd bij één als bij zestien lanes.

De aanwezigheid van drie DMA-engines onthult interessante aannames:

  • Kernel krijgt een dedicated pad: Dit is een commitment dat de kernel een aparte operand-klasse is met een andere levenscyclus.
  • KernelDMA is unidirectional load-only: Men ging ervan uit dat weights niet teruggeschreven zouden worden.
  • TileDMA is bidirectional en gedeeld: Intermediate tile-outputs kunnen worden teruggevoerd zonder een trip naar DRAM.
  • Kernelgeheugen is privé per core: Men verwachtte dat weights veelvuldig per-core hergebruikt zouden worden.
  • Nieuwe kernels moeten van DRAM worden geladen, niet van L2: Men verwachtte waarschijnlijk geen grote weights te streamen die meer dan \(64\text{ KiB} \times 16\) overschrijden.

Kernel L1

Een convolutie is geen symmetrische multiply-add van \(A\) en \(B\). Een convolutie schuift dezelfde kernel \(w[k]\) over de input:

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

Dit betekent dat de kernel één keer geladen kan worden en vervolgens hergebruikt kan worden over de gehele duur van de output, terwijl nieuwe inputs naar binnen schuiven. Dit is precies wat Apple's datapath-patent beschrijft.

De bedoelde steady state is:

  • DRAM → L2 → MAC (Activatie) → L2 → DRAM (Output)
  • DRAM → KMem (Coefficient)

Als je een ASIC ontwerpt om convoluties te verwerken — waarbij de kernel niet frequent en dynamisch wordt gestreamd (zoals bij KV caches) — is het logisch om de asymmetrie van de operanden te exploiteren en kernel-geheugenbewegingen te verminderen.

Maar waarom L2 → KMem niet toestaan?

  1. L2-bandbreedte: Ik betwijfel dat L2-contentie de reden was; resident kmem bestaat juist omdat kernels infrequent geladen zouden worden.
  2. Complexiteit: Apple distribueert L2-endpoints al naar elke core; een kernel-pad via dezelfde tile-route zou niet veel complexer zijn.

De fysieke layout van de ANE is gecentreerd rond de L2 SRAM-rechthoek, met cores aan de zijkanten. Mijn vermoeden is dat Apple simpelweg nooit verwachtte dat L2-resident tensors kernels zouden worden, wat een geldige aanname was in 2017. Bovendien houdt Apple van geïsoleerde modulaire modules; het was waarschijnlijk makkelijker om de 1 MiB kmem (read-only) en 2 MiB tile L2 volledig gescheiden te ontwikkelen.

4. Is het voorbij?

DRAM-doorvoersnelheid

Is de GPU sneller dan de ANE?

Transformer single-token decode is het ergste scenario voor weight-hergebruik en de ratio compute per geheugen, omdat we de volledige model-weights moeten streamen om één token te genereren. Als zowel de ANE als de GPU read-bandwidth bound zijn, zal degene met de hoogste leesbandbreedte meer tokens/s decoderen, ongeacht de piek-compute capaciteit.

Aangezien ANE en GPU hetzelfde DRAM delen, is de ANE niet noodzakelijkerwijs benadeeld in DRAM-toegang. Als de ANE langzamer is, komt dat omdat de DMA-controller niet genoeg parallelle requests kan onderhouden om de DQ te verzadigen.

Om de DRAM-leesbandbreedte te meten, kan een read-bandwidth-bound workload worden gegenereerd. De resultaten zijn:

  • ANE KernelDMA: 37.99 GB/s
  • ANE TileDMA: 59.08 GB/s
  • GPU: 77.70 GB/s

Kan de ANE 38 + 60 = 98 GB/s halen om het plafond van de M3 (100 GB/s) te raken? Nee. Experimenten tonen aan dat de gecombineerde runtime van kernel + tile overeenkomt met de som van de geïsoleerde runtimes. Dit betekent dat ANE's kernel- en tile-DMA requests serieel worden verzonden (één per keer). De ANE DRAM-throughput is dus dubbel benadeeld:

  1. Zowel geïsoleerde kernel als tile DMA zijn lager dan de GPU-leesbandbreedte.
  2. Kernel- en tile-DMA tijden zijn additief.

Met de ANE decode gepind aan de DRAM-roofline, kan een drastische verbetering (bijv. van 10 naar 25 tok/s) alleen komen van een $\sim 2.5\times$ hogere geheugen-streaming bandbreedte.

Maar... wat als we 50 GB/s aan extra kernelDMA-throughput zouden kunnen toevoegen?