AppleのNeural Engineを遡ってリバースエンジニアリングする

Retrospectively Reverse-Engineering Apple's Neural Engine Aug 10, 2026(5089 words) 3年前に、リバースエンジニアリングされたApple Neural Engine（ANE）ドライバーの開発を中止しました。ANEブロックがあまり有用ではなく、もっと有用なことができると悟り、他のより有用なブロックのアップストリーミングに移ったのです。ANEのアーキテクチャは、汎用アクセラレータプラットフォームを構築するにはあまりにも個性的すぎ、ANEハードウェアAPIアクセスを事実上開放するLinuxドライバーでは、実行できるワークロードのクラスを広げることはできませんでした。macOSでさえ、Finderでのアップサンプリングされたプレビュー画像の生成にANEを定期的に使用する程度です。 https://github.com/eiln/ane/tree/main M1ダイショット: https://mastodon.social/@dougall/115149886886125067 M5（2025年）の目玉機能は「LLMパフォーマンス」でしたが、ANEコアもGPUコア内に便利に折り込まれました。これは来るだろうと分かっていましたが、スタンドアロンNPUの終焉の始まりを公式に感じさせます。そこで、ANEの明らかな終焉に敬意を表して、さらに無用なことをしましょう。M1のANEを遡ってリバースエンジニアリングし、始めたことを完成させます。もう3年（クソ）経ちましたが、最初にこれを始めたときよりも多くのことを知っているはずです。 3年前の目標がANEをオペレーション実行可能にして有用にすることだったとすれば、今回は、計算、データフロー、スケジューラ、メモリ、実行モデルといった、ANEの完全な内部アーキテクチャのマッピングが目的です。なぜなら、これらの内部設計上の決定は、Appleが最初にA11 Bionic（2017年）でシリコンにコミットしたMLワークロードに関する仮定、そしてそれがCNN時代のNPUから今日のTransformerワークロードを実行するGPUへのシフトについて何を意味するかを明らかにしているからです。 1. 計算 16個の計算コアはおそらくANEの最も興味のない部分です。Appleは当初、予測可能な再利用を伴う密なテンソル削減で構成される、密な画像処理CNNワークロードをターゲットにしていました。M1 ANEの計算コアはMAC（Multiply-Accumulate）ユニットの大きな並列アレイですが、それだけでは、それが設計され、アクセラレーションされるワークロードについてほとんど何も語りません。 畳み込み層は、アクティベーションウィンドウと学習済みカーネル重みとの間の内積を実行し、アテンションはクエリベクトルとキーベクトルとの間の内積を実行します。内積は内積であり、MACはまさにそれを実行します。2017年のCNNモデルに対してANEを特殊化させたのはMACではなく、MACを取り巻くデータフローです。MACの入力と出力がいつ、どこに入り、留まり、移動するかです。Transformer、特に自己回帰デコードが壊れたという仮定は、ANEが電話で実行できるほど効率的なデータフローをアーキテクトするために利用した、予測可能な再利用パターンでした。M5の決定は、ANEの計算コアがTransformerに対しても依然として有用であったことを確認していますが、それは異なるデータフロー内でのことです。 それでも、16個の計算コアそれぞれの内部データフローは以下のようになります。 ┌────────────────────── core ─────────────────────┐ │ ┌───────── 256× MACs ─────────┐ ┌────────────┐ │ │ │ MAD ─► add ─► accumulator │─►│ activation │ │ │ │ ▲ │ │ └────────────┘ │ │ │ └──────────┘ │ └─────────────────────────────┘ └─────────────────────────────────────────────────┘ Multiply-Accumulate ANEは16個の並列計算コアを持っています。各計算コアには、128個のFP16（または256個のINT8）並列MACレーンがあります。各MACレーンは、以下の再帰を実行します。 \[ s leftarrow s+a\times b \] 2つのオペランドaとbを乗算し、その積を実行中の合計（アキュムレータ）に加算します。MAC操作をTサイクル繰り返すことで、T項の内積を計算します。 \[ s_T=s_0 + \sum_{t=0}^{T-1} a_t \, b_t. \] したがって、1つの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レーンは、それがどの行列またはテンソルの次元を削減しているかを知りません。内積、行列乗算、畳み込みが、オペランドがコアにどのようにマッピングされ、スケジューリングされるかから生じることは重要です。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サイクル間で外部メモリから遠く離れた場所からフェッチする必要がなくなります。 解像度に関して、読み出し時にFP16で固定小数点削減を行います。乗算器は16ビットで、32ビットレジスタにQ16.16として蓄積され、その後、符号拡張などによりFP16として読み出されます。整数（16進数）FP16表現で作業し、アキュムレータの範囲をプローブするために、CoreML ANEプログラムを構築して、すべての（1）のベクトルとの内積を計算します。これにより、各乗算の結果は有界vになりますが、アキュムレータ内の実行中の合計は増加し続けます。 \[ s=\sum_{i=0}^{255}v=256v. \] (v) CPU 16進数 CPU 値 ANE 16進数 CoreML 値 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合計がMAC後の活性化ブロックに直接フィードされ、中間メモリラウンドトリップを回避することです。これは、活性化が点ごとであるため可能です。スカラー削減が完了すると、その活性化はそのスカラーにのみ依存し、すぐに適用できます。 ANEがtanh()をどのように実装するかを判断するために、単一のTANH活性化レイヤーを含むCoreMLモデルをコンパイルし、結果のコンパイル済みハードウェアレジスタファイル（hwx）を調べます。係数領域には、0x4288から始まる33個の連続したFP16ワードが含まれています。 0x4270: 3120 3001 0000 0000 0000 0000 0000 0000 0000 0x4280: 0000 0044 0000 003c 0000 f52f d633 bc35 # 0.000000 0.124329 0.244873 0.358398 0x4290: 6537 7038 1539 a239 183a 793a c93a 0a3b # 0.462158 0.554688 0.635254 0.704102 0.761719 0.809082 0.848145 0.879883 0x42a0: 3e3b 673b 883b a23b b63b c63b d33b dd3b # 0.905273 0.925293 0.941406 0.954102 0.963867 0.971680 0.978027 0.982910 0x42b0: e53b eb3b ef3b f33b f63b f83b fa3b fb3b # 0.986816 0.989746 0.991699 0.993652 0.995117 0.996094 0.997070 0.997559 0x42c0: fc3b fd3b fe3b fe3b ff3b 0000 0000 0000 # 0.998047 0.998535 0.999023 0.999023 0.999512 0x42d0: 003c 0300 6000 0000 0000 0000 0000 0000 これらの33個のFP16ワードは、\(\tanh(x)\)の33個のIEEE LE FP16量子化サンプルに一致します。 \[ T_i=\operatorname{round}_{16}\!\left(\tanh(i/8)\right), \qquad i=0,1,\ldots,32. \] 次にRELU活性化レイヤーに切り替えます。 活性化プログラム 非線形モデル ルックアップ係数 アイデンティティ 0 なし ReLU 1 なし tanh 2 33 FP16ワード したがって、モード2はカスタム33エントリールックアップテーブルを選択します。33点では32個の間隔が定義されます。\(R=3\)の場合、ノットは \[ x_i=\frac{i}{8},\qquad i=0,\ldots,32, \] 間隔\(1/8\)で\([0,4]\)をカバーします。入力は\(u=2^R|x|\)としてテーブルにマッピングされるため、\(R\)はノット間隔を設定します。解像度は33ビンよりも滑らかです。隣接するエントリは線形補間されると推測されます。テストするために、単一のスパイクを持つインパルスLUTを構築します。 \[ T_8=1,\qquad T_k=0\ \text{for }k\ne8,\qquad R=3. \] 次に、\(T_8\)の周りの2つのセルにわたって入力をスイープします。測定された出力は三角形を形成します。大きさは\(|x|=0\)から線形に増加します。