2
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

AMX (Apple Matrix Coprocessor) ってなんだ?〜Appleが隠し続ける最強コプロセッサ〜

2
Posted at

はじめに:Appleが公開しない「秘密兵器」

Apple Siliconには、公式ドキュメントに載っていない秘密のコプロセッサが存在する。

その名は「AMX」。Apple Matrix Coprocessorの略だ。

Neural Engineじゃない。GPUでもない。CPUコアの中に潜む、行列演算専用の隠しユニット。

「え、そんなのあるの?」

ある。しかも2019年のA13 Bionic(iPhone 11)から搭載されている。M1、M2、M3、M4、すべてのApple Siliconに入っている。

なのにAppleは、このAMXについて公式ドキュメントを一切出していない。WWDC(開発者会議)でも、ほんの30秒触れただけ。API も公開されていない。

なぜ隠すのか? そして、どうすれば使えるのか?

今回は、リバースエンジニアリングによって明らかになったAMXの全貌を解説する。


Neural Engineとは何が違う?

まず混乱しがちなポイントを整理しよう。

Apple Siliconには「Neural Engine」という機械学習アクセラレータが搭載されている。M4なら16コア、38 TOPS(毎秒38兆回の演算)という化け物スペック。

AMXはこれとは全く別物だ。

Neural Engine

独立したハードウェアアクセラレータ。CPUやGPUとは別の「部署」にいる。機械学習の推論に特化。INT8やFP16といった低精度の演算が得意。使うにはCore ML経由が基本。

AMX(Apple Matrix Coprocessor)

CPUコアに付属するコプロセッサ。CPUの命令ストリームから直接制御される。汎用的な行列演算が可能。FP32やFP64の高精度演算もサポート。AccelerateフレームワークのBLAS/LAPACKが内部で使用。

イメージで言うと、Neural Engineは「外注先の専門業者」で、AMXは「社内の精鋭部隊」だ。Neural Engineに仕事を投げるには依頼書(Core MLモデル)が必要だが、AMXはCPUが直接指示を出せる。


どれくらい速いのか?

ここからが本題。AMXの性能は、NEONベクトル命令と比べてどれくらい速いのか?

ある研究者の検証によると、M1 Maxの単一コアでの行列乗算(SGEMM)性能は以下の通り。

ARM NEON(128bit SIMD): 約102 GFLOPS
AMX: 約1,475 GFLOPS

14倍以上の差がある。

これ、同じCPUコアの話だ。NEONを使うかAMXを使うかで、同じ計算が14倍速くなる。

なぜこんな差が出るのか?

NEONは128ビット幅のSIMD(Single Instruction Multiple Data)ユニット。4つのFP32を同時に処理できる。

AMXは、32x32のマトリクス演算グリッドを持つ。外積(outer product)演算なら、1命令で1024個のFP16乗算を実行できる。FP32でも256個。

並列度が桁違いなのだ。


AMXのアーキテクチャ

リバースエンジニアリングによって判明したAMXの構造を見てみよう。

レジスタ構成

AMXには3種類のレジスタプールがある。

Xレジスタ(入力用)
8本 x 64バイト = 512バイト
行列演算の「左側」の入力を格納

Yレジスタ(入力用)
8本 x 64バイト = 512バイト
行列演算の「右側」の入力を格納

Zレジスタ(出力/累積用)
64本 x 64バイト = 4,096バイト
演算結果を格納。累積加算(accumulate)にも使用

これらのレジスタは、通常のCPUレジスタとは完全に独立している。AMX専用の作業スペースだ。

演算モード

AMXは主に2つのモードで動作する。

マトリクスモード(外積演算)
XレジスタとYレジスタの外積を計算し、Zレジスタに累積。行列乗算の基本演算。

ベクトルモード(内積演算)
XレジスタとYレジスタの内積を計算。32x32グリッドの1行だけを使用。

対応データ型は、FP16、FP32、FP64、INT8、INT16、INT32と幅広い。


なぜAppleは公開しないのか

さて、これだけ強力なコプロセッサを、なぜAppleは公式サポートしないのか?

推測だが、いくつかの理由が考えられる。

ARMライセンスの問題

Appleは ARM アーキテクチャのライセンスを受けて Apple Silicon を設計している。ARM の方針として、独自拡張命令を公開してアーキテクチャを「分断」することは推奨されていない。

Appleが独自にAMX命令を公開すると、ARM ISA(Instruction Set Architecture)の統一性が損なわれる。他のARM実装との互換性問題も発生しうる。

将来の柔軟性

公式APIを出してしまうと、後から仕様変更しにくくなる。AMXの実装詳細を隠しておけば、M5やM6でアーキテクチャを大幅に変更しても、既存アプリには影響しない。

AccelerateフレームワークというラッパーがAMXの詳細を吸収してくれるから、開発者は気にしなくていい——というのがAppleの設計思想だろう。

競争優位性

AMXの詳細を公開すると、競合他社(Qualcomm、Samsungなど)に手の内を見せることになる。「謎のコプロセッサで速い」という状態を維持したいのかもしれない。


じゃあどうやって使うの?

「非公開なら使えないじゃん」

いや、実は使える。間接的に。

Accelerateフレームワーク経由

Appleが提供するAccelerateフレームワークには、BLAS(Basic Linear Algebra Subprograms)やLAPACK(Linear Algebra Package)といった線形代数ライブラリが含まれている。

これらのライブラリは、内部でAMXを使用している。

import Accelerate

// 行列乗算(SGEMM)
// 内部でAMXが使われる
var C = [Float](repeating: 0, count: m * n)
cblas_sgemm(
    CblasRowMajor,      // 行優先
    CblasNoTrans,       // Aは転置しない
    CblasNoTrans,       // Bは転置しない
    Int32(m),           // Aの行数
    Int32(n),           // Bの列数
    Int32(k),           // Aの列数 = Bの行数
    1.0,                // alpha
    A, Int32(k),        // 行列A
    B, Int32(n),        // 行列B
    0.0,                // beta
    &C, Int32(n)        // 結果行列C
)

このコードを実行すると、Accelerateが内部でAMXを呼び出す。開発者はAMXの存在を意識する必要がない。

Pythonの場合、NumPyやSciPyがAccelerateにリンクされていれば、同様にAMXの恩恵を受けられる。

import numpy as np

# NumPyの行列乗算
# macOSではAccelerate経由でAMXが使われる
A = np.random.randn(1000, 1000).astype(np.float32)
B = np.random.randn(1000, 1000).astype(np.float32)
C = A @ B  # この演算でAMXが動く

vDSP / BNNSフレームワーク

信号処理(vDSP)やニューラルネットワーク(BNNS)のフレームワークも、適切な場面でAMXを活用する。

import Accelerate

// 畳み込み演算(BNNSFilter)
// 適切なサイズなら内部でAMXが使われる可能性がある
let convolution = BNNSFilterCreateLayerConvolution(
    &inputDesc,
    &outputDesc,
    &weightsDesc,
    &biasDesc,
    &params,
    nil
)

直接使用(非推奨だが可能)

リバースエンジニアリングによって、AMXの命令セットは解析されている。直接アセンブリを書けば使えなくはない。

// 警告: 非公式・非サポート・将来動かなくなる可能性大
// AMXを有効化(op=17, operand=0)
__asm__ volatile(
    "nop\n"
    "nop\n"
    "nop\n"
    ".word 0x00201220"  // AMX enable
);

// 以下、AMX命令を直接発行...

ただし、これは絶対に推奨しない

理由は明白で、Appleがいつ仕様を変更するかわからないから。ある日のmacOSアップデートで突然動かなくなっても、文句は言えない。

Accelerate経由で使うのが唯一の「正しい」方法だ。


AMXが効く場面、効かない場面

効く場面

大きな行列乗算
行列サイズが大きいほど、AMXのスループットを活かせる。BLASのSGEMM、DGEMM など。

バッチ処理
同じ演算を大量のデータに適用する場面。画像のバッチ処理、信号処理など。

密行列演算
要素がほとんど埋まっている行列の演算。疎行列ではオーバーヘッドが相対的に大きくなる。

効かない場面

小さな演算
AMXの初期化オーバーヘッドがあるため、小さな行列では逆に遅くなる可能性がある。

スカラー演算
単一の数値に対する演算には向かない。NEONや通常のALUのほうが適切。

分岐の多い処理
AMXはストリーミング処理に最適化されている。if文だらけのコードには不向き。


Neural Engine vs AMX:どっちを使うべき?

「両方行列演算できるなら、どっちがいいの?」

用途による。

Neural Engineを選ぶべきケース

  • 機械学習の推論(特にCore MLモデル)
  • INT8/FP16で十分な精度の処理
  • バッテリー効率を重視
  • モデルサイズが大きい(Neural Engineのほうがスループットが高い)

AMXを選ぶべきケース(Accelerate経由)

  • 高精度演算(FP32/FP64)が必要
  • 科学計算、数値シミュレーション
  • BLASやLAPACKを使う既存コード
  • Neural Engineを他のタスクに使いたい

実際のところ、Core MLを使えば、フレームワークが自動的に最適なユニット(CPU/AMX、GPU、Neural Engine)を選んでくれる。開発者が細かく指定する必要はない。


ベンチマーク:実際に測ってみる

AMXの効果を確認するには、Accelerateの行列演算をベンチマークしてみるといい。

import Accelerate
import Foundation

func benchmarkSGEMM(size: Int, iterations: Int) -> Double {
    let count = size * size
    var A = [Float](repeating: 0, count: count)
    var B = [Float](repeating: 0, count: count)
    var C = [Float](repeating: 0, count: count)
    
    // ランダム初期化
    for i in 0..<count {
        A[i] = Float.random(in: -1...1)
        B[i] = Float.random(in: -1...1)
    }
    
    let start = CFAbsoluteTimeGetCurrent()
    
    for _ in 0..<iterations {
        cblas_sgemm(
            CblasRowMajor,
            CblasNoTrans,
            CblasNoTrans,
            Int32(size),
            Int32(size),
            Int32(size),
            1.0,
            A, Int32(size),
            B, Int32(size),
            0.0,
            &C, Int32(size)
        )
    }
    
    let elapsed = CFAbsoluteTimeGetCurrent() - start
    
    // GFLOPS計算: 2 * n^3 演算 / 時間
    let flops = Double(2 * size * size * size * iterations)
    let gflops = flops / elapsed / 1e9
    
    return gflops
}

// 実行
let gflops = benchmarkSGEMM(size: 1024, iterations: 100)
print("Performance: \(gflops) GFLOPS")

M1 Maxで実行すると、シングルスレッドで1000 GFLOPS以上が出るはずだ。NEONだけでは絶対に到達できない数字。これがAMXの威力。


まとめ:存在するが、見えない力

AMX(Apple Matrix Coprocessor)は、Apple Siliconに搭載された「公式には存在しない」コプロセッサだ。

ポイント

  • Neural Engineとは別物。CPUコアに付属する行列演算ユニット
  • NEON比で14倍以上の行列演算性能
  • Accelerateフレームワーク経由で自動的に使われる
  • 直接APIは非公開。リバースエンジニアリングでのみ判明
  • FP32/FP64の高精度演算に対応(Neural EngineはFP16/INT8中心)

使い方

  • Accelerateフレームワーク(BLAS、LAPACK、vDSP、BNNS)を使う
  • NumPy/SciPyをmacOSで使う(Accelerateにリンクされていれば自動)
  • Core MLを使う(最適なユニットが自動選択される)

使っちゃダメな方法

  • 非公開命令を直接叩く(将来動かなくなるリスク大)

Appleが隠し続ける最強コプロセッサ。その存在を知っているだけで、「なぜApple Siliconはこんなに速いのか」の理解が深まる。

そして、Accelerateを使うだけで、その恩恵は自動的に受けられる。

公式が隠しても、リバースエンジニアリングの光は届く。技術者の好奇心、侮るなかれ。


参考リンク

2
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
2
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?