Gflop/s から Tflop/s へ: Swift における行列乗算の最適化
Large Language Model (LLM) の学習は本質的に巨大な行列乗算の演習です。その核心は、z += x * y を何兆回も繰り返すループです。Apple Silicon 上で開発する際の課題は、Swift の高レベルな安全性と、これらのワークロードに必要な生のパフォーマンスとのバランスを取ることにあります。
最近の取り組みで、開発者 zdw は Andrej Karpathy の llm.c（GPT‑2 互換モデルのプレーン C 実装）を Swift に書き直すことに挑みました。目標は単に C と同等の性能を出すことだけでなく、CPU と SIMD 命令、"シークレット" の AMX コプロセッサ、Metal 経由の GPU など、M シリーズチップが提供するすべてのツールを駆使して Swift の限界を押し上げることでした。
出発点: パフォーマンスギャップ
C から基本的な Swift へ matmul_forward 関数を移植したとき、最初の結果は顕著な差がありました。Release ビルドでランタイムアサートを除去したにもかかわらず、基本的な Swift 実装はプレーン C バージョンより 15〜20 倍遅くなっていました。
| モデル | Tokens/s | Training iterations/s | Training vs llm.c |
|---|---|---|---|
| llm.c | 0.926 | 0.175 | 100% |
| Basic Swift | 0.054 | 0.014 | 7.3% |
これはおよそ 2.8 Gflop/s の性能に相当します。1999 年ならば印象的だったかもしれませんが、現代の LLM ワークロードには受け入れがたい数字です。主な原因は _ArrayBuffer.beginCOWMutation() にありました。Swift の Copy‑on‑Write (COW) ユニークチェックが、配列がユニークであっても大きなオーバーヘッドを生んでいました。
ギャップを埋める: Swift レベルの最適化
COW のボトルネックを取り除く第一歩は、Swift 6.2 で導入された MutableSpan を採用することでした。これにより、ほぼゼロオーバーヘッドでメモリにアクセスできるようになりました。この改善だけでも学習速度は向上しましたが、Swift には C の -ffast-math フラグに相当する直接的なオプションがなく、Fused Multiply‑Add (FMA) 命令を有効にできませんでした。
緩和数学と SIMD の活用
Swift‑Numerics ライブラリとその Relaxed.multiplyAdd 関数を使用することで、実装はついに fmla（SIMD ベクトル化 FMA）命令を利用できるようになりました。この変更だけでトークン毎秒がほぼ 10 倍に向上しました。
ループ展開とインライン配列
最適化された C 実装は、コンパイラにループ展開を促すためにループを跨いで走査します。著者は Swift 6.2 の InlineArray を利用し、スタック上にバッファを確保してヒープ割り当ての高コストを回避しました。この段階で "Fast Swift" は C と同等、むしろ若干上回るトレーニングイテレーション速度（llm.c の 106.6%）を達成しました。
スケールアップ: マルチスレッドと AMX
シングルスレッドの性能は解決したものの、次の飛躍はすべての CPU コアを活用することでした。DispatchQueue.concurrentPerform を使って M3 Max の 16 コアにワークロードを分散させました。しかし、これによりコードが大幅に "視覚的に散らかる" ため、withUnsafeMutableBufferPointer と @unchecked Sendable ラッパーで Swift の並行安全チェックを回避する必要がありました。
"シークレット" の武器: AMX
標準 SIMD に加えて、Apple Silicon には AMX（Apple Matrix Coprocessor）があります。Apple は公式には Accelerate フレームワーク経由でのみ公開していますが、AMX_MATFP のような逆コンパイルされた命令を使えば 16×16 タイルを直接操作できます。
警告: 本番環境での AMX 命令の直接使用は推奨されません。未文書化であり、バイナリ互換性が壊れる可能性があります。Accelerate フレームワークの利用が推奨パスです。
AMX 命令を実装した結果、トレーニング性能は元の llm.c の 958.8% にまで向上しました。
最終フロンティア: Metal と GPU
Tflop/s 領域に到達するため、ワークロードを Metal を使って GPU に移行しました。移行は Metal/C++ でコンピュートカーネルを書き、Swift で呼び出し層を実装する形で行われました。
- Basic Metal: 素朴なカーネルは AMX に比べて僅かな向上しかもたらさなかった。
- Threaded Metal:
threadsPerThreadgroupを最適化することで大幅に跳ね上がり（llm.cの 2204.6%）。 - Tiled Metal: メモリ局所性を改善するタイル化カーネルを実装し（長い行の走査を減らす）、ついに 1 Tflop/s の壁を突破した。
最終性能比較
| モデル | Tokens/s | Training iterations/s | Training vs llm.c |
|---|---|---|---|
| llm.c | 0.926 | 0.175 | 100% |
| Multithreaded Swift | 4.356 | 1.014 | 558.5% |
| AMX | 5.884 | 1.678 | 958.8% |
| Tiled Metal | 11.123 | 5.351 | 3057.7% |
技術的洞察と反論
FMA 議論
コミュニティで取り上げられた重要な点は -ffast-math の使用です。著者は FMA を有効にするためにこのフラグを使用しましたが、専門家の中には -ffast-math が広範すぎて数値的な不正確さを招く可能性があると指摘する人もいます。FMA のみを安全に有効にしたい場合は、-ffp-contract=fast の使用が推奨されます。
GPU ソフトウェアの堀
"Basic Metal" から "Tiled Metal" への移行の難しさは、NVIDIA の CUDA のようなソフトウェアエコシステムが依然として支配的である理由を浮き彫りにしています。GPU のピーク性能はハードウェアだけで決まるわけではなく、特定のデータ形状に最適化された膨大なカーネルライブラリが不可欠です。
結論
2.8 Gflop/s の素朴な実装から出発し、著者は 1.1 Tflop/s（382 倍）の性能を達成しました。この旅路は、Swift が C と同等、あるいはそれ以上の速度を出せることを示す一方で、コードが unsafe ポインタや手動メモリ管理に陥り、言語本来のエレガンスが失われる代償が伴うことを示しています。実運用では、既に何年もかけてカーネルを最適化してきた Accelerate、CoreML、MPSGraph といった確立されたフレームワークを利用することが最善策です。