150行のCコードでNumPyの行列積を上回る
(salykova.github.io)- NumPyの行列積は外部のBLASライブラリに依存するが、この実装は純粋なCとFMA3・AVX2だけで、シングルスレッドとマルチスレッドの性能をBLAS級まで引き上げることを目指している
- 性能の中核は、$C$を小さなブロックに分割し、16×6マイクロカーネルが
YMMレジスタ内でrank-1 updateを繰り返してメモリアクセスを減らす構造にある - 任意サイズの行列では境界処理がボトルネックになりやすいため、マスクストアとゼロ埋めバッファを組み合わせ、マスクロードによる性能低下を回避している
- キャッシュ再利用は
k_c,m_c,n_cブロッキングで確保し、実際の最高性能はスレッド数・カーネルサイズ・タイルサイズのチューニングに大きく左右される - AVX-512はより広いCPUサポートのため除外されているので、AVX-512対応CPUではBLASのほうが速い可能性があり、OpenBLASとの比較もAVX-512を無効化した条件で行われている
実装目標と比較対象
- 実装コードはsgemm.cで公開されており、最新プロセッサでマルチスレッドFP32行列積を最適化している
- NumPyは行列積のような線形代数演算を外部のBLASライブラリに依存している
- 例としてIntel MKL、Accelerate、BLIS、GotoBLAS、OpenBLASがある
- OpenBLAS、GotoBLAS、BLISはC/FORTRAN/Assemblyで書かれており、CPUマイクロアーキテクチャごとの手動最適化された行列積実装を含む
- 目標は、低レベルのアセンブリを使わず純粋なCで書きつつ、次の条件を満たす行列積実装である
- 任意の行列サイズで動作する
- 最新のx86-64プロセッサで実行できる
- 既存のBLASライブラリと競合できる
- コードがシンプルで拡張しやすい
- 参考資料は、Simon BoehmのFast Multidimensional Matrix Multiplication on CPU from Scratch、Sergey SlotinのMatrix Multiplication、GeohotのCan you multiply a matrix?、GotoBLAS・BLIS関連論文である
ベンチマーク条件とFLOPS計算
- テスト環境はAMD Ryzen 7 9700X、32GB DDR5 6000 MHz CL36、OpenBLAS 0.3.26、GCC 13.3、Ubuntu 24.04.1 LTS
- コンパイルフラグは
-O3 -march=native -mno-avx512f -fopenmpを使用する - 公平な比較のため、OpenBLASのインストール時に適切な
TARGETを設定し、AVX-512命令を無効化する必要がある- Zen4/5プロセッサは
make TARGET=ZENでコンパイルする - そうしないとOpenBLASがデフォルトでAVX-512命令を使う
- Zen4/5プロセッサは
- OpenBLASのFP32行列積は
cblas_sgemmAPIで実行する - ベンチマークは正方行列を対象とする
m=n=k=200からm=n=k=10000まで200刻みで評価する- 行列積を
n_iter回繰り返し、中央値の実行時間を性能測定に使う
- $M \times K$行列$A$と$K \times N$行列$B$を掛けると、総演算量は**$2MNK$ FLOP**となる
- 性能は
FLOPS=(2*m*n*k)/exec_timeで計算する
- 性能は
理論限界とSIMDベース
- 最新のx86-64 CPUはSIMD拡張によって複数データを並列処理する
- 主な命令はAVX2とFMA
- どちらも256ビットの
YMMレジスタを使う - 各
YMMレジスタには32ビットfloatを8個格納できる
- どちらも256ビットの
- FMA命令
VFMADD231PSはYMM1 = YMM2 * YMM3 + YMM1形式のpacked single演算を行う - Ryzen 9700Xではfused multiply-addのスループットは0.5 cycles/instruction、つまり1サイクルあたり2命令
- 理論上、Ryzen 9700Xは単一コアで1サイクルあたり32 FLOPを実行できる
- 計算式は
8 floats × 2(add+mul) × 2(1/TP) - 8コアで4.7GHzの持続クロックを仮定すると、マルチスレッドの理論ピークは1203 FLOPSと見積もられる
- 計算式は
基本実装とマイクロカーネル
- 行列はcolumn-major順で格納する
A[row][col]にはCポインタでptr[col*M + row]としてアクセスする
- 最も単純な実装は、$C$のすべての行と列を走査し、各要素ごとに$A$の行と$B$の列の内積を計算する
- 高性能実装の要は、$C$を$m_R \times n_R$部分行列に分割し、各部分行列を効率よく計算するマイクロカーネルである
- カーネルは$\bar{C}$をレジスタ上で0に初期化した後、$K$次元に沿って反復する
- $\bar{A}$の列ベクトルと$\bar{B}$の行ベクトルをレジスタに読み込む
- 2つのベクトルの外積を計算して$\bar{C}$アキュムレータに加える
- 各ステップはrank-1 updateである
- この方法では、naive方式のメモリアクセス量$2K m_R n_R$に比べ、レジスタへ読み込む要素数を$(m_R+n_R)K$まで減らせる
- AVX CPUには16個のYMMレジスタがあるため、カーネルサイズは次の制約を満たす必要がある
- $(m_R/8) \cdot n_R + m_R/8 + 1 \le 16$
- $m_R$は8の倍数でなければならない
- 理論的には$m_R$と$n_R$が大きく、かつ同程度であるほどメモリアクセス削減効果は大きいが、実際のRyzen 9700Xでは16×6カーネルが最も高い性能を示した
- 実装では
immintrin.hのintrinsicを使用する__m256は256ビットベクタ型で、YMMレジスタの内容を表す_mm256_loadu_psでAの列ベクトルをロードする_mm256_broadcast_ssでBのスカラ値を8個のfloatベクタへブロードキャストする_mm256_fmadd_psでアキュムレータを更新する_mm256_storeu_psで結果をメモリへ保存する
- 生成されたアセンブリには
vfmadd231psやvbroadcastssのようなSIMD FMA命令が含まれる
任意サイズ行列のためのパディング
- 基本の16×6カーネルは、$M$と$N$がそれぞれ16と6の倍数のときにそのまま動作する
- 境界領域で列数$n$が6より小さい場合は、保存ループを
j < nまでだけ実行する - 行数$m$が16より小さいときは、
_mm256_storeu_psが一度に8要素を保存するため、マスクストアが必要になる_mm256_maskstore_psは、マスクビットが有効な要素だけをメモリに保存する- マスクは重なっている行数$m$に応じて生成する
- 境界でのロードまで
_mm256_maskload_psで処理すると、カーネル性能が大きく低下する可能性がある- マスク計算の追加命令がオーバーヘッドになる
- $n$がコンパイル時定数ではないため、コンパイラがループを効率よくアンロールしにくい
- そこで$m \neq m_R$なら$\bar{A}$をバッファにコピーしてゼロ埋めし、$n \neq n_R$なら$\bar{B}$もバッファにコピーしてゼロで埋める
- 関連実装はmatmul_pad.hにある
キャッシュブロッキングとデータ再利用
- レジスタとDRAMの間にはCPUキャッシュ階層があり、最新のデスクトップCPUは通常L1、L2、L3キャッシュを使う
- キャッシュはDRAMより高速だが容量が限られるため、$A$、$B$、$C$全体をすべてキャッシュに載せる方法は不可能である
- 行列を小さなブロックに分けてキャッシュに載せ、同じデータを複数のrank-1 updateで再利用する方法がキャッシュブロッキングまたはタイリングである
- シングルスレッドのキャッシュブロッキングはBLISの構造に似た5重ループ形式である
- 最外ループは$N$次元に沿って$C_j$と$B_j$ブロックを作る
- 次のループは$K$次元に沿って$A_j$と$B_p$ブロックを作る
- $B_p$はパッキングされて$\tilde{B}_p$となり、必要に応じてゼロ埋めされてL3キャッシュ再利用を狙う
- 次のループは$M$次元に沿って$C_i$と$A_j$ブロックを作り、$A_j$はパッキングされて$\tilde{A}_j$となる
- 最後の2つのループはキャッシュブロックを$m_R \times k_c$、$k_c \times n_R$パネルに分けてカーネルへ渡す
- パッキングされた$\tilde{A}_j$と$\tilde{B}_p$は異なる形で格納される
- $\tilde{A}_j$内部のパネルはcolumn-majorで格納される
- $\tilde{B}_p$内部のパネルはrow-majorで格納される
- キャッシュブロッキングのパラメータはCPUモデルごとのキャッシュサイズに合わせて調整する必要がある
- $k_c \times n_c$はL3キャッシュを満たす出発点になる
- $m_c \times k_c$はL2キャッシュを満たす出発点になる
- $k_c \times n_R$はL1キャッシュを満たす出発点になる
- 実際には理論値より大きな値のほうが高性能になる場合が多く、CPUがキャッシュ配置を自動管理するため、アルゴリズムレベルではループとアクセスパターンを設計する必要がある
- 実装はmatmul_cache.hにある
カーネルの微細最適化
__m256 C_buffer[6][2]のように配列でアキュムレータを定義する代わりに、アキュムレータ変数を明示的に展開して宣言する- この方法はGCCがコードをより適切に最適化し、レジスタスピルを避けるのに役立つ
- マスク計算もベクトル命令を使うように変更する
mask[32]静的配列を用意し、_mm256_cvtepi8_epi32と_mm_loadu_si64を使う
- 該当実装はmatmul_micro.hにある
マルチスレッド戦略
- 並列化の対象は算術演算とパッキングの両方である
- マイクロカーネルの外側にある5番目、4番目、3番目のループは、キャッシュブロックサイズ単位で反復する
- すべてのスレッドを忙しく保つには、反復回数がスレッド数以上でなければならない
- 入力行列の次元は概ね
スレッド数 × キャッシュブロックサイズ以上である必要がある
- Ryzen 9700Xのシングルスレッドで高性能だったキャッシュブロックサイズは$n_c=1535$、$m_c=1024$
- 8コアをすべて活用するには、少なくとも$\max(m_c,n_c) \times 8 = 1535 \times 8 = 12280$の次元サイズが必要になる
- 一方、最後の2つのループは小さな$m_R$、$n_R$ブロックを反復するため、並列化に向いている
- 一般に$m_R$、$n_R$は20未満である
- $m_c$、$n_c$をコア数の倍数に選ぶと作業を均等に分配できる
- Ryzen 9700Xでは、2つの内側ループを
#pragma omp parallel for collapse(2) num_threads(NTHREADS)でまとめて並列化する方法が最も高い性能を示した - 多コアのプロセッサ、特に16コア超の環境では、ネストした並列化や2〜3本のループ並列化を検討できる
- $\tilde{A}$と$\tilde{B}$のパッキングもOpenMPで並列化する
pack_blockAはmcをMR単位で走査しながら並列化するpack_blockBはncをNR単位で走査しながら並列化する
- マルチスレッド実装でRyzen 9700X上で高性能だったパラメータは次の通り
- $m_c = m_R \times \text{number of threads} \times 5$
- $n_c = n_R \times \text{number of threads} \times 50$
- 最終的なマルチスレッド実装はmatmul_parallel.hにある
1件のコメント
Hacker News のコメント
この記事の要点が、一般に性能改善の余地は残っているということなら、むしろ改善幅を過小評価している。行列乗算ライブラリに注がれる労力は、ほとんどのソフトウェアよりはるかに大きいにもかかわらずだ。
すでに強力に最適化されたコードでなければ、大きな労力なしに既存コードから10〜1000倍以上改善されることは珍しくない。おおよそ重要度順に見ると、アルゴリズム選択が適切かどうか、作業そのものをなくせるかどうかが最も重要で、カーネルとの往復や
mallocのような重い処理を減らせるかどうかも大きい。ベクトル化は明示的なベクトル組み込み関数もよいが、構造体配列の代わりに配列／構造体の配列へデータを再構成するだけで、同じ機械語が出る場合も多い。キャッシュ効率も重要で、並列コードでは偽共有のようにスレッドごとのデータ分離ができていない場合、さらに複雑になる。最後に、組み込み関数や手書きアセンブリのようなハードウェア別の最適化も可能だ。
結合がリモートサーバー上で行われるようにクエリを変え、ネットワークには5〜10行だけ送るようにしたら、すぐに速くなった。固定オーバーヘッドとレイテンシは常にあるが、必要量をはるかに超えるデータをネットワーク接続で送ると、結局性能は破綻する。レイテンシの影響を扱った “It's the latency, stupid” も読む価値がある: http://www.stuartcheshire.org/rants/latency.html
全体としては、上の考慮事項とおおよその順序に同意する。
多くの面接は、実装がなぜ遅いのかを推論し、ベンチマークして修正する方法を見る代わりに、「Google がそうしているから」式の、難解なアルゴリズム暗記クイズになってしまっている。
よくあるコーディングパターンはハードウェアに十分特化しておらず、多くの性能を取り残している。この記事は興味深い例で、別の古典的なデモとして “There's plenty of room at the top” がある。
https://www.science.org/doi/10.1126/science.aam9744
この内容を理解するには、BLIS リポジトリの論文群がほぼ正典に近い。最適化された BLAS で性能が出ないと考える理由が分からないし、十分大きな行列なら CPU ピークの90%以上を期待すべきだ。
最後に見たとき、シリアル OpenBLAS は概ね MKL と同程度で、BLAS は基本線形代数ブロックとして
matmulではなく GEMM を実装する。通常のベンチマークフレームワークの代わりに numpy を使うのも理解できないし、Zen では AMD の BLAS、つまり BLIS ベースの実装と比較すべきだと思う。BLIS は以前、OpenBLAS より並列化の面で話がよく、AMD BLIS には「小さい」次元向けに実装を切り替える仕組みもあるが、現在の OpenBLAS にあるかは分からない。マイクロカーネルのベクトル化に SIMD 組み込み関数が必須というわけではなく、まともな C コンパイラなら完全にベクトル化し、ループも展開してくれる。BLIS の純 C マイクロカーネルは、適切なブロックサイズなら Haswell の手作業で最適化された実装に対して80%以上の性能を出す。差はおそらくプリフェッチのせいだろうが、正確には理解していない。
最新のコンパイラでも、自動ベクトル化とループ展開を100%の成功率で適切にやり切れるわけではない。
記事と実装は良さそうに見えるが、「秘訣」が何なのか気になる。OpenBLASはまさにこの問題のために、何十年も アセンブリ+Cで最適化されてきたのに、どうやって勝てるのだろうか
キャッシュなどを詳しく扱っているが、BLASはこうしたものを活用していないのか、それとも特定のプロセッサによりうまく合わせているのか気になる
性能差は平均スループットよりピークスループットでずっと顕著だったが、ピークが重要なアプリケーションはほとんどない。示されているベンチマークコードでは、numpy側はPythonのアロケータを通っていて、C実装はアロケータを通っていないようなので、マイクロベンチマークの誤りや不一致をまず確認すべき箇所だ。多くのnumpyルーチンはインプレース演算をサポートしているため、両方ともインプレース版のベンチマークを明示的に見る必要がありそうだ
numpyには下位実装とは無関係に実行される境界チェックとエラー処理もあり、小さな行列では純粋なPythonリストよりも非常に遅くなる理由になる。数千サイクルの純粋なオーバーヘッドを足すと、高速にするのは難しい
この実装は関連するキャッシュを飽和させようとする、かなり原則的なアプローチで、ある意味では当たり前だが、明確なエンジニアリング上の改善はこうした議論で強調する価値がある。OpenBLASにも多くの人手が投入されているが、すべてを考え尽くしている可能性は低い。きちんと説明するには、双方のコードに対する深い分析が必要だ
C++とCの実装はメタプログラミングのアプローチ [2], [3] を見ればよい。本当に驚くべきなのは、Matlab、Julia、Mojoのような多くの現代的な言語が今でもOpenBLASに依存している点だが、もちろんそれぞれ理由はあるのだろう
[1] Numeric age for D: Mir GLAS is faster than OpenBLAS and Eigen (2016):
http://blog.mir.dlang.io/glas/benchmark/openblas/2016/09/23/...
[2] Vastly outperforming LAPACK with C++ metaprogramming (2018):
https://wordsandbuttons.online/vastly_outperforming_lapack_w...
[3] Outperforming LAPACK with C metaprogramming (2018):
https://wordsandbuttons.online/outperforming_lapack_with_c_m...
-march=nativeは正確なCPUモデルに合わせてコンパイルするので、利点があるかもしれない。numpyはより汎用的で古いx86-64向けにコンパイルされている可能性が高いRyzen CPUでは
-march=nativeがたぶんv4を使い、numpyはv1かv2をターゲットにしていそうだhttps://en.wikipedia.org/wiki/X86-64#Microarchitecture_level...
記事も良いし、ベンチマークを簡単に再現できるようにしている点も素晴らしい。私の 16コア Xeon W-2245 3.90GHz では、
matmul.cが8192x8192の行列乗算をgcc -O3で1.41秒、clang -O2で1.47秒で実行し、NumPyは1.07秒だったAVX-512カーネルならずっと速くなると思う。性能が物足りないもう一つの理由はOpenMPかもしれず、経験上、
pthreadsでスレッドプールを明示的に管理するとオーバーヘッドを減らせる。CPU数もハードコードではなくsysconf(_SC_NPROCESSORS_ONLN)を使うほうがよい片方はPythonで、もう片方はCというように負荷を変える理由はない。両方ともCで書き、片方は BLASライブラリ を呼び、もう片方はこの実装を呼ぶ形にすれば、条件をそろえた比較ができたはずだ
オーバーヘッドはそれほど大きくないが、このスレッドの他の場所でも述べられているように、正しく呼び出すことが重要だ。素朴なnumpyコードと調整済みのCコードを戦わせるのは、明らかに公正な比較ではない
ホットパスではないが、マスク生成の非効率さ、つまり
bit_maskの使用が気になる。より効率的な方法としては、{-1,-1,...,0,0,...}という形のグローバル定数配列を作り、要素オフセット16-m、8-mからロードするか、定数ベクトル{0,1,2,3,4,...}をブロードキャストされたmおよびm-8と比較する方法があるただし、行列の1列にだけ関係する部分で、その後に続く
maskload/maskstoreループのほうがはるかに時間がかかるので、本当に些細な指摘だ。特にストアは Zen 4 でも依然として遅く[1]、AVX-512 命令はマスクをマスクレジスタから受け取るという違いしかないのに6倍速い。clang はいずれにせよシフトを自動ベクトル化するので、私の提案より2〜3倍遅い程度だろう[1]: https://uops.info/table.html?search=vmaskmovps&cb_lat=on&cb_...
新しい視点を与えてくれるフィードバックには本当に感謝しています。「定数のグローバル配列を作ってロードする」方法は、記憶では試したときにビットマスクのシフトより少し遅かった気がしますが、確かめるためにもう一度テストしてみます。「定数ベクトル
{0, 1, 2, 3, 4, ...}をブロードキャストされたmとm-8と比較する」方法は良いアイデアなので試してみますint8_tにしておき、ロード時にバイトをint32_tに符号拡張できます。_mm_loadu_si64/_mm256_cvtepi8_epi32の組み合わせは、メモリオペランドを使う単一のvpmovsxbd命令にコンパイルされるはずですこうすれば、
alignas(32)で適切にアラインした場合、定数配列全体が1本のキャッシュラインに収まります。元記事のユースケースではマスクが2つ必要なので、2つ目のvpmovsxbd命令は確実な L1D キャッシュヒットになり、うまく合いますjart の tinyBLAS はどうだろう
https://hacks.mozilla.org/2024/04/llamafiles-progress-four-m...
それから https://justine.lol/matmul/
ベンチマーク以外で、行列乗算そのものをマルチスレッド化する理由は何だろう。実際には、乗算を使うアルゴリズム側でマルチスレッドを使うほうが有利ではないだろうか
とはいえ、HPC コードではたいてい GEMM はボトルネックではない
まだざっと見ただけだが、この記事は細部と説明が多い。高速な行列乗算がアーキテクチャ上の考慮事項を反映してどのように実装されるかをかなりよく説明している記事に見えるので、読むリストに入れておいた