行列エンジンと混合精度

この章でわかること:

• 12章の順位表の続き — ベンダーライブラリの実測(自作最良の11倍、1.47 TFLOP/s)
• 行列エンジン(テンソルコア、Apple AMXなど)というハードウェアの構造
• 混合精度 — 13章のf16/bf16が実戦投入される形
• Rustから行列計算能力を使う現実的な選択肢
• なぜ体系に必須か — 現代の計算ハードウェアの能力の重心は ここにあります。機械学習の時代の「計算が速い」の意味を 知らずに、計算機の体系は閉じません

この章の実験はリポジトリのexamples/ch25-gemm-libです(macOS専用。 他のOSではOpenBLAS等に読み替えてください)。

cd examples
cargo run --release -p ch25-gemm-lib

12章の順位表には続きがあった

12章で、n=1024の行列積を素朴な873ミリ秒から 11.3ミリ秒まで、77倍磨き上げました。あのとき 「実用レベルはライブラリを使うのが正解」と書き添えたのを 覚えているでしょうか。その言葉を実測にします。

macOSに標準搭載のAccelerateフレームワーク(BLASという線形代数 ライブラリ標準のApple実装)のsgemm(f32行列積)を、 20章のFFIで呼び出して同じ計算をします。

n = 1024 (筆者のApple M4)
自作 ikj+rayon (12章のCPU最良): 16.4ms  (  131 GFLOP/s)
自作 GPU blocked (12章のGPU最良): 11.3ms (  190 GFLOP/s)
Accelerate sgemm               :  1.5ms  ( 1470 GFLOP/s)

毎秒1.47兆回の浮動小数点演算。私たちが12章で到達した最良値の さらに8〜11倍です。素朴版から数えると約600倍—— そして誤差は10^-4のオーダーで自作版と一致しています(f32の 丸め順序の違い、13章による差です)。

この数字は、ソフトウェアの技巧だけでは説明できません。 AppleのCPUクラスタには行列演算専用のハードウェアが 搭載されており(初期世代は非公開の行列コプロセッサ、通称AMX。 M4世代ではArm標準のSME命令として公開)、Accelerateは これを利用していると考えられます。

行列エンジン — SIMDの次の段階

4章のSIMDは「1次元の並び」への同一演算でした。 行列エンジン(matrix engine)はその次の段階—— 「2次元のタイル」同士の乗算累積を1命令(相当)で行う専用回路です。

• NVIDIA テンソルコア(Tensor Core) — ワープ(9章) 単位で小さな行列の乗算累積(例: 16×16)を実行するGPU内ユニット
• Apple AMX / Arm SME — AppleのCPUクラスタに付属する行列演算 ユニット(AMXは非公開命令、SMEはArm標準)。今回のAccelerateの 性能はこの系統の仕事と考えられます
• Intel AMX — サーバCPUのタイルレジスタ+行列乗算命令

なぜこれほど速いのか。本書をここまで読んだ読者には、 2つの言葉で説明できます。

第一に、データ再利用の回路化です。15章の キャッシュブロッキング、12章のレジスタ ブロッキングで、私たちはソフトウェアで「読んだ値を使い回す」形を 作りました。行列エンジンは、タイル同士の乗算に必要な再利用を 配線として持っています。オペランドを1回読み込めば、 その値は回路の中でn回使われます。ルーフライン(10章)の 「演算とメモリの比率」を、命令セットのレベルで引き上げているのです。

第二に、制御の削減です。1命令が数百〜数千の乗加算に化けるため、 フェッチ・デコード(16章)のコストが 演算数千回ぶんに1回で済みます。9章で見た「制御を削って演算に 面積を回す」思想の極致です。

混合精度 — 13章の実戦投入

行列エンジンのカタログ値は、多くの場合f16やbf16 (13章)で最大になります。数値表現を 半分にすれば、同じ配線とメモリ帯域で2倍の要素を流せるからです。

ここで13章の知識がそのまま実戦になります。典型的な 行列エンジンの動作は「入力はf16/bf16、累積はf32で」です。 入力の精度は下げても、足し込む側を広い型に保つ—— 13章のKahanの実験で見た 「累積だけ広い型で」と同じ設計が、ハードウェアに焼き込まれています。 bf16がf32と指数幅を揃えている(=桁あふれの挙動が同じ)ことの 価値も、実際に使われる場面はここです。

さらに推論の世界では、重みをint8やint4に量子化 (quantization)して4〜8倍詰め込む技法が標準化しています。 「必要な精度だけに支払う」——数の表現はダイヤルであって 定数ではない、というのが現代の計算の常識です。

なぜMLがこのハードウェアを生んだか

行列エンジンの存在は、 「ハードウェアは主要ワークロードの写し鏡」という本書の隠れた テーマの最新の実例です。GPUがゲームの画面描画から生まれたように (9章)、行列エンジンは深層学習から 生まれました。ニューラルネットワークの計算時間は大部分が行列積で、 しかも学習は多少の丸め誤差に頑健です。「計算のほぼ全部が1種類の 演算で、精度を下げてよい」——ハードウェア屋にとってこれほど 専用化しがいのある相手はありません。

帰結として、行列積とそれ以外で、計算機の実効性能は 1桁以上違う時代になりました。あなたの問題を行列積の形に 寄せられるなら(多くのML・信号処理・物理計算はそうです)、 1.47 TFLOP/sの世界が使えます。寄せられないなら、 本書のPart I〜IVの世界——百数十GFLOP/sとメモリ帯域——が 主戦場です。この地形を知っていることが、設計の出発点になります。

Rustからの選択肢

行列計算能力への到達経路を整理します。

• OSのライブラリをFFIで — 今回のAccelerateの方法 (20章の実践)。LinuxならOpenBLAS/BLIS、 NVIDIAならcuBLASが相当します
• faer — 純Rustの 高性能線形代数ライブラリ。FFIなしでBLAS級を狙う注目株です
• candle / burn — MLフレームワーク。バックエンド(Metal/CUDA/wgpu)を切り替えて テンソルコア系の能力に到達します
• Metal Performance Shaders / cuBLAS — GPU側の行列ライブラリ。 wgpuの自作カーネル(12章)の上位互換です

指針は12章から変わりません——密行列の計算を自作しない。 自作が正当なのは、学びのため(本書がそうです)と、 ライブラリが対応しない特殊な形状・融合(カスタムカーネル)のみです。

まとめ

• ベンダーライブラリの行列積は1.47 TFLOP/s——12章の自作最良の 8〜11倍でした。正体は行列エンジン(AMX)というハードウェアです
• 行列エンジンは、ブロッキングによるデータ再利用を配線に焼き込み、 制御コストを演算数千回に1回へ薄めた、SIMDの次の段階です
• カタログ性能はf16/bf16の混合精度で最大化されます。 「積は狭く、累積は広く」は13章の原理の実戦形です
• 行列積に寄せられる問題は1桁上の性能地帯に入れます。 Rustからの入口はFFI、faer、candle/burnです

次章はPart VIの締めくくり——ここまで「壁時計」で測ってきた GPUの時間を、GPU自身のタイムスタンプで正確に測る方法です。