Apple Neural Engine (ANE)アーキテクチャの逆解析

Apple Neural Engineは固定機能のCNNアクセラレータである

Apple Neural Engine (ANE)は、独自の命令セットアーキテクチャ (ISA) を持つ汎用プロセッサではなく、畳み込みニューラルネットワーク (CNN) の密なテンソル還元と予測可能な再利用パターンに特化した固定機能のデータフロー・エンジンである。Transformerワークロードを実行できるが、2017年当時の画像処理に最適化されたデータパスにマッピングすることで実行するため、現代のLLM推論では顕著なパフォーマンスボトルネックが生じる。

計算アーキテクチャとMACデータパス

ANEは16の計算コアから構成され、各コアには128のFP16(または256のINT8)並列乗算累積(MAC)レーンが含まれている。

時間方向のスカラ還元

各MACレーンは時間方向にスカラ還元を実行し、2つのオペランドを乗算して32ビットの累積レジスタに加算する。16コアのANEは合計2,048の並列MACレーンを提供し、チップ全体にわたって空間的に還元を実行する。ハードウェアは特定のレイヤータイプ(例:4チャネルCNN)をエンコードしない。代わりに、オペランドのマッピングとタスクのスケジューリングが、ドット積、行列乗算、または畳み込みのいずれであるかを決定する。

精度と累積

内部還元は固定小数点演算で行われる。乗算器は16ビットであり、結果はQ16.16形式で32ビットレジスタに累積される。その後、FP16として読み出される。累積レジスタは$\pm 2^{15}$で飽和し、これは16小数ビットを持つ符号付き32ビット固定小数点値の範囲に相当する。

非線形活性化とLUT

ANEはtanhなどの非線形活性化を、33エントリの区分線形ルックアップテーブル(LUT)を使って実装している。ハードウェアは入力をLUT座標にスケーリングし、隣接するエントリ間で線形補間を行い、滑らかな出力を生成する。この活性化ブロックはMAC後のパスに直接統合されており、中間メモリの往復を回避している。

スケジューラとタスク記述子

GPUとは異なり、変長の命令ストリームをパースするのではなく、ANEはタスク記述子(TD)によって制御される。これらのTDは実行可能なコードではなく、1回のデータパス通過用にANEのハードウェアレジスタを設定するためのシリアル化されたレジスタファイルのダンプである。

コマンド送信

ドライバは固定サイズのTDのコマンドストリームをタスクマネージャーに送信する。その後、ハードウェアは「ControlDMA」エンジンを使って、これらの設定値をDRAMからANEの物理レジスタ空間にコピーする。設定が完了すると、ANEはタスクを実行し、完了時にARM64コアに割り込みを発生させる。

タスクキュー(TQ)とBAR

ANEは8つのタスクキュー(TQ)を用い、これらは起動コンテキストとして機能する。このシステムの重要な構成要素は32エントリのベースアドレスレジスタ(BAR)テーブルである。ANEにはGPU風の動的ロード/ストア命令がないため、すべての仮想アドレスアクセスは、コンパイル時にハードコードされたBARベースオフセットからの相対オフセットとして実行されなければならない。

メモリ階層と帯域幅ボトルネック

ANEのパフォーマンスは、CNNカーネルのDRAMトラフィックを最小限に抑えるために設計された厳格なメモリ階層によって制御される。

ローカルメモリ構成

  • KMem: コアごとに16個の64 KiB SRAM(合計1 MiB)で、カーネル用に使用。
  • L1: コアごとに16個のMAC入力ステージング領域。
  • L2: 全コアで共有される1つの2 MiB L2 SRAM。

DRAMルーフライン

M1 ANEの場合、DRAM帯域幅制限を回避するための算術強度は1バイトあたり162回の演算が必要である。この比率よりもオンチップ再利用が少ないワークロードでは、計算能力を増やしてもパフォーマンスは向上しない。

DMAエンジンの制限

ANEは3つの異なるDMAエンジンを使用している:KernelDMASrc(重み → KMem)、TileDMASrc(DRAM → L2)、TileDMADst(L2 → DRAM)。

逆解析により、現代のワークロードに対して2つの重大な欠陥が明らかになった:

  1. 非対称なパス:カーネルメモリは読み込み専用であり、L2からポピュレートできない。これにより、すべての新しいカーネルはDRAMからロードしなければならない。これはCNNでは妥当な仮定だったが、トランスフォーマーではテンソルが頻繁に役割を変えるため、悪影響を及ぼす。
  2. 逐次実行:実験データによると、KernelDMAとTileDMAのリクエストは逐次的に送信される。実行時間は加算されるため、重みと活性化のフェッチを重ねてシステムDRAM帯域幅を満たすことはANEでは不可能である。

洞察の統合

技術的分析とコミュニティの議論から、CNN時代のANEから現代のAIハードウェアへの移行が浮き彫りになった:

  • CNNとトランスフォーマーのマッピング:ユーザーは、トランスフォーマーをANEに移植する際に「それがCNNだったかのように振る舞う」必要があると指摘しており、1x1畳み込みを使用して4Dテンソルで行列乗算をシミュレートしている。
  • ハードウェアの進化:M1 ANEは独立したブロックだったが、最新のアーキテクチャ(例:M5)ではANEコアがGPUコアに統合されたと報告されており、トランスフォーマーベースのLLMに必要なより柔軟なデータフローへのシフトを示している。
  • 初期のイノベーション:現在のボトルネックにもかかわらず、ANEは専用AIシリコンへの重要な初期投資を示しており、Aシリーズチップに2017年から搭載されている。

"ANEのアーキテクチャは、一般用途のアクセラレータプラットフォームを構築するにはあまりにも偏見が強すぎた……[それは] AppleがA11 Bionic(2017年)で最初にシリコンにコミットしたMLワークロードに関する仮定を示しており、CNN時代のNPUから今日のGPUがトランスフォーマーワークロードを実行する時代へのシフトを物語っている。"

DRAM読み取り帯域幅比較

デバイス/エンジン 持続的読み取り帯域幅
ANE KernelDMA 37.99 GB/s
ANE TileDMA 59.08 GB/s
GPU (Metal) 77.70 GB/s

ANEのDMAリクエストは逐次的であるため、ANEはピークMACスループットにかかわらず、LLM推論で最もメモリ制限の厳しい単トークンデコードにおいて、GPUよりも根本的に遅い。

Sources

関連

  • プロジェクト
  • Dispatch
  • Dispatch
  • プロジェクト
  • Dispatch