3年前、エイリーン・ヨンはアップルのMシリーズチップに搭載された専用機械学習アクセラレータであるApple Neural Engine(ANE)のリバースエンジニアリングを開始し、Linuxドライバーを公開しました。彼女は一旦、このブロックは汎用性が高すぎると考え、作業を中断しました。そして今回、ANEを新しいワークロードに活用するためではなく、2017年のA11 Bionicにニューラルアクセラレーションが初めて導入された際の、アップルがとった設計戦略を明らかにするために、再び作業に取り組んでいます。

本研究の主な発見は、ANEがCNN(Convolutional Neural Networks、画像処理アーキテクチャの2018年以前の主流)向けに設計されたのに対し、現在のLLMを駆動するトランスフォーマー・モデルには適していないということです。CNNの密度の高い、予測可能なデータ再利用は、ANEのデータフローとよくマッチしますが、トランスフォーマーの自己回帰的デコードには適していません。ヨン氏によると、アップルの対応策は、M5(2025年)のGPUクラスタにANEコアを統合することであり、これにより、スタンドアロンのNPUは事実上、廃止されたことになります。

M1 ANEの各計算コアには128の並列乗算累積(MAC)レイがあり、チップ全体では2048となります。これらのMACは、再帰式s ← s + a × bを実行します。MACはスカラー減算ユニットであり、その演算対象のテンソルの次元は知りません。これらの意味は、オペランドがコアにどのようにスケジュールされるかによって決まります。蓄積器は32ビットの固定小数点レジスタで、Q16.16形式、2^15で飽和し、積はFP16で出力されます。この選択—16ビットの乗算、32ビットの蓄積、16ビットの出力—は、アップルがターゲットとしたCNNワークロードの結果であり、アクティベーションは有界で、範囲は予測可能で、固定小数点に優れています。

非線形性については、ANEはハードウェアでtanhを計算しません。33エントリのピースワイズ線形ルックアップテーブルを使用し、線形補間を行います。区間は1/8の間隔で配置されます。ReLUはモード1(恒等写像)、tanhはモード2(ルックアップテーブル)です。CoreMLのコンパイラは、バイアスとアクティベーションを同じMAC後のパスにコンパイル時に組み込みます。バイアスレジスタと非線形モードはレイごとに設定され、MAC出力は中間メモリ往復なしで直接ルックアップにフィードされます。

Diagram showing Core ML folding constant scale and offset into convolution weights and bias before the ReLU activation
Core ML compiler folds bias and activation into the post-MAC path at compile time, avoiding an intermediate memory round-tripeiln.github.io

The article explains that CoreML's compiler folds bias and activation into the same post-MAC path at compile time, with the bias register and nonlinear mode set per-layer, and the MAC output feeds directly into the lookup without an intermediate memory round-trip

One nonzero lookup-table entry produces two straight-line segments on the ANE's piecewise-linear tanh approximation
A single nonzero LUT entry generates two straight-line segments, illustrating the ANE's piecewise-linear interpolation schemeeiln.github.io

The article describes the ANE's nonlinearity approach: a 33-entry piecewise-linear lookup table with linear interpolation, spacing knots at intervals of 1/8, where tanh is mode 2 and ReLU is mode 1

Core ML tanh overlaid with double-precision tanh, followed by signed error plot
Core ML tanh output overlaid against double-precision tanh, with the signed error beloweiln.github.io

The article discusses the ANE's tanh implementation as mode 2 of the piecewise-linear lookup table; this figure validates the approximation's accuracy against a double-precision reference

スケジューラは、ヨン氏の説明によると、非常に単純です。ドライバはCONVやMATMULのコードを送信しません。すべてのニューラル演算は、固定サイズのタスク記述子の命令ストリームに事前にコンパイルされます。ドライバは、このストリームをメモリに読み込み、TM_ADDR/TM_SIZEをそれに向け、TM_PUSHを発行します。ハードウェアは、完了するまでこのコマンドを所有し、完了するとARM64コアに割り込みを上げます。8つのタスクキュー(qid 0-7)はピングポンスタージを提供し、1つの記述子が実行されている間、ソフトウェアがもう1つを埋めます。各キューは、タイルレベルの作業をマッピングする32エントリのベースアドレスレジスタテーブルを持っています。

ヨン氏のGitHubリポジトリ(eiln/ane)には、カーネルモジュール、ユーザースペースライブラリ、Pythonバインディングが含まれています。193のコミット、537のスターです。この研究は、現在のANEを有用にすることを目的としたものではありません。2017年にアップルがシリコンに刻み込んだアーキテクチャの仮定を読み取り、それらがトランスフォーマー時代に直面したときに何が起こったかを追跡するためのものです。

不明な点も残っています。M4以降のANEブロックが、ヨンが文書化したM1のANEの設計を超えた機能を提供しているかどうか(議論のきっかけとして、AppleはM6/A20にANEブロックを2倍に増設し、M5のGPUコア内の「ニューラルアクセラレータ」は別のサブシステムである可能性があります)、記事には独立した検証がなく、ヨンのソースはM1 ANEのみを扱っています。また、独立した開発者がANEにアクセスできるかどうか(Appleがコンピュテーショナル写真処理パイプラインに限定しているかどうか)も、この供給された資料では解決されていません。1人のコメント者はその制限を主張していますが、供給されたドライバー自体が、逆アーキテクチャによってLinuxアクセスが存在することを示しています。