概要
Swiftを用いて、Apple SiliconのCPU、AMX、GPUを最大限に活用し、行列演算を標準実装の380倍以上高速化してTflop/s級のパフォーマンスを実現するプロセスを解説しています。
詳細内容
この記事は、Andrej Karpathy氏の「llm.c」をモデルに、外部ライブラリを使わずSwiftでLLM（大規模言語モデル）の訓練を最適化するシリーズの第1回です。主な内容は、行列演算の性能を2.8 Gflop/sから1.1 Tflop/sまで引き上げる具体的な手法の紹介です。まず、SwiftのArray固有のオーバーヘッド（COWチェックなど）を回避するためにSwift 6.2の`MutableSpan`を導入し、次に`Swift Numerics`の`Relaxed.multiplyAdd`を用いてSIMD命令を活用したFused Multiply-Add（FMA）を実装します。さらに、`InlineArray`によるスタック割り当てとマルチスレッド化（DispatchQueue.concurrentPerform）を組み合わせ、最終的には非公開のAMX（Apple Matrix Coprocessor）命令や、Metalによるタイリング（Tiling）演算を実装することで、GPUのポテンシャルを引き出しています。理論的な解説だけでなく、Apple Siliconの各演算ユニットの特性と、Swiftでの低レイヤー最適化の可能性を実証する内容となっています。