RustではじめるCPUとGPU

行列エンジンと混合精度

この章でわかること:

  • 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倍高速化しました。12章では 「実用レベルではライブラリを使うのが正解」と書きました。 この章では、その主張を実測で確かめます。

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)

Accelerateのsgemmは、毎秒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倍の要素を格納する技法が 標準化しています。数の表現は用途に応じて選ぶ設定値であり、 固定の前提ではありません。

機械学習が行列エンジンを必要とした理由

行列エンジンは、「ハードウェアの設計は主要なワークロードに 合わせて決まる」という、本書で繰り返し述べた事実の最新の例です。 GPUがゲームの画面描画のために設計されたように (9章)、行列エンジンは深層学習のために 設計されました。ニューラルネットワークの計算時間は大部分が行列積で、 しかも学習は多少の丸め誤差の影響を受けにくい性質があります。 計算のほぼ全部が1種類の演算で、精度を下げてよいという条件は、 専用ハードウェア化に最も適しています。

このため、行列積とそれ以外とで、計算機の実効性能は 1桁以上異なります。解きたい問題を行列積の形に 変換できるなら(多くの機械学習・信号処理・物理計算が該当します)、 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自身のタイムスタンプで正確に測る方法を扱います。