0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

matmulってなんだ? — 素朴な3重ループが1000倍遅くなる理由

0
Posted at

この記事の対象読者

  • PythonNumPyを使ったことがあり、A @ B と書いたことはあるが、その中で何が起きているかは説明できない方

対象読者は1レベルに固定しています。線形代数の講義を受けたことがなくても読めるように書きました。

この記事で得られること

  • 行列積の計算量が 2MNK FLOPsで数えられること、そしてその数え方の使いどころ
  • 同じ計算式でも実装によって1000倍の速度差が出る理由
  • 演算強度という指標と、なぜ深層学習が行列積という形に計算を寄せているのか
  • 手元で再現できる計測スクリプト一式

この記事で扱わないこと

  • Strassenアルゴリズムなど、計算量そのものを下げる手法の理論
  • CUDAカーネルの実装(本記事の計測はすべてCPUです)
  • 行列の数学的性質(固有値、階数など)

導入:厨房に立った日

自作のミニ言語モデルを書いていたとき、線形層をPythonの3重ループで実装しました。動きます。テストも通ります。ただ、学習が1エポックで終わりませんでした。翌朝に見たらまだ回っていました。

np.dot に置き換えたら、同じ計算が数百ミリ秒で終わりました。式は1文字も変えていません。

この記事は、そのときに調べたことの整理です。全体を通して、行列積を大量注文をさばく厨房にたとえて説明します。

  • FLOPs = 作る料理の総量
  • メモリ転送 = 食材の搬入
  • 演算強度 = 食材1kgあたり何皿つくれるか
  • BLAS = 動線を設計し尽くしたプロの厨房
  • 素朴な3重ループ = 注文が来るたびに冷蔵庫まで走る素人

1. matmulの定義とFLOPsの数え方

このセクションで分かること:行列積の定義と、計算量を 2MNK で数える方法。これ以降の全ての計測はこの数式が基準になります。

matmulは matrix multiplication の略で、行列積のことです。NumPynp.matmul 、演算子の @PyTorchtorch.matmul はすべて同じ演算を指します。

$M \times K$ の行列 $A$ と $K \times N$ の行列 $B$ の積 $C = AB$ は、次のように定義されます。

C_{ij} = \sum_{k=1}^{K} A_{ik} B_{kj}

出力は $M \times N$ 個の要素を持ち、その各要素は長さ $K$ のベクトル同士の内積です。内積1回につき $K$ 回の乗算と $K$ 回の加算が必要なので、総演算回数は次のようになります。

\mathrm{FLOPs} = 2MNK

係数の2は、乗算と加算を1回ずつ数えているためです。この数え方はNVIDIAの公式ガイドでも採用されています。積和演算が M・N・K 回必要で、1回あたり2演算なので合計 2MNKという数え方が示されています。

厨房のたとえで言えば、$2MNK$ はその日に作る料理の総量です。厨房の設計がどうであれ、作る量そのものは変わりません。変わるのは、その量をさばくのにかかる時間です。

計算のしかたは1通りではありません。同じ $2MNK$ を、どの順番でさばくかという選択肢があります。


2. 実測:素朴な3重ループ と NumPy

このセクションで分かること:同じ式でも実装で何倍違うのか、実際の数字。計測環境と再現手順も示します。

計測はすべて、Xeon 2.10GHz の1コア、メモリ3GB、Python 3.12.3、NumPy 2.4.4(OpenBLAS 0.3.31)の環境で行いました。GPUは使っていません。時間は3〜5回試行の最小値です。

厨房に入る前に、まず素人の動きを見ます。ここからは実際の走り方、つまりコードです。先ほどの $2MNK$ が、そのまま flops 変数として登場します。

クリックで計測コードを展開
import time
import numpy as np

def naive_matmul(A, B):
    n = len(A); k = len(A[0]); m = len(B[0])
    C = [[0.0] * m for _ in range(n)]
    for i in range(n):
        Ai, Ci = A[i], C[i]
        for p in range(k):
            a, Bp = Ai[p], B[p]
            for j in range(m):
                Ci[j] += a * Bp[j]
    return C

def timeit(fn, repeat=3):
    best = float("inf")
    for _ in range(repeat):
        t0 = time.perf_counter()
        fn()
        best = min(best, time.perf_counter() - t0)
    return best

for n in (32, 64, 128, 256):
    A, B = np.random.rand(n, n), np.random.rand(n, n)
    Al, Bl = A.tolist(), B.tolist()
    t_naive = timeit(lambda: naive_matmul(Al, Bl), 1)
    t_np = timeit(lambda: A @ B, 5)
    flops = 2 * n ** 3
    print(n, t_naive / t_np, flops / t_naive / 1e9, flops / t_np / 1e9)

結果です。

N 素朴ループ NumPy 倍率 素朴 GFLOPS NumPy GFLOPS
32 1.33 ms 0.0024 ms 550倍 0.049 27.1
64 9.58 ms 0.010 ms 916倍 0.055 50.1
128 72.5 ms 0.077 ms 939倍 0.058 54.3
256 574 ms 0.53 ms 1079倍 0.058 63.1

注目すべきは、素朴ループのGFLOPSがサイズによらず 0.05 前後で頭打ちになっている点です。つまり素朴ループの遅さは「Pythonが遅いから」だけでは説明しきれません。行列が大きくなってもスループットが1ミリも改善していない、という別の事実がここにあります。

一方でNumPy側は $N$ が大きくなるほどGFLOPSが伸びています。厨房で言えば、注文が増えるほど段取りの良さが効いてくる状態です。

さらに大きなサイズで、float32のスループットの上限を見ます。

N 時間 GFLOPS 前サイズからの時間比
128 0.035 ms 119.1
256 0.236 ms 142.1 6.70倍
512 1.92 ms 139.7 8.14倍
1024 14.9 ms 144.1 7.76倍
2048 119 ms 144.4 7.98倍

$N$ を2倍にすると時間はおよそ8倍。$2^3 = 8$ なので、実測が $O(N^3)$ に一致しています。理論式が机上の空論ではないことが、この列だけで確認できます。


3. なぜ1000倍も違うのか

このセクションで分かること:速度差の正体がメモリ階層であること、そして「同じ食材を何度使い回せるか」という視点。

素朴ループは、$C_{ij}$ を1つ計算するために $A$ の1行と $B$ の1列を読みます。次の $C_{i,j+1}$ でも $A$ の同じ行を読み直します。厨房で言えば、1皿つくるたびに冷蔵庫まで走って同じ食材を取りに行っている状態です。

BLASは行列をタイルに割り、タイルをキャッシュに載せてから使い切ります。冷蔵庫から一度に運んだ食材で、作業台の上で何皿分もまとめて仕込む動きです。

ここで、うっかり信じていた仮説が1つ外れました。メモリレイアウトを揃えれば速くなるだろうという予想です。C順序とFortran順序を組み合わせて計測しました。

配置 時間 GFLOPS
C @ C 14.36 ms 149.5
C @ F 14.63 ms 146.8
F @ C 14.79 ms 145.2
F @ F 14.90 ms 144.2

差は最大でも4%弱でした。OpenBLASは内部でタイルへの詰め替えを行うため、呼び出し側のレイアウトはほぼ吸収されます。手元で確かめるまでは、ここを気にして np.ascontiguousarray を挟むコードを書いていました。少なくとも大きな正方行列については、その手間は無駄でした。


ここまでのまとめ

  1. 行列積の演算量は 2MNK FLOPsで数える。実装を変えてもこの量は減らない。
  2. 素朴な3重ループとBLASでは、$N=256$ で1079倍の差が出た。実測は $O(N^3)$ に一致した。
  3. 差の正体はメモリ階層の使い方。ただし呼び出し側のメモリレイアウトはBLASがほぼ吸収する。

4. 演算強度 — matmulだけが特別な理由

このセクションで分かること:行列積が「食材1kgあたり何皿つくれるか」で圧倒的に有利であること、そしてその指標の計算方法。

演算強度は、転送1バイトあたり何FLOPs実行できるかを表す指標です。NVIDIAのガイドは計算回数とメモリアクセス回数の比率を演算強度と呼ぶと定義しています。

厨房でいえば、搬入した食材1kgから何皿つくれるかです。この値が小さいと、コンロがいくら強力でも搬入待ちで手が止まります。

同じサイズの2048×2048のfloat32行列に対して、行列積と要素ごとの積を計測しました。

演算 FLOPs 転送バイト 演算強度 実測GFLOPS
行列積 1.72e10 5.03e7 341.3 148.7
要素ごとの積 4.19e6 5.03e7 0.083 2.1

演算強度で約4100倍、実測スループットで約70倍の差です。要素ごとの積は、読んだ値を1回掛けたら捨てるので、完全にメモリ帯域で律速します。行列積だけが、読んだ値を $N$ 回使い回せます。

これが、深層学習の設計が徹底して行列積に寄せられている理由です。試しに、Transformer1層のFLOPsを自分で数えてみました。$d_{model}=2048$、系列長 $L=2048$、ヘッド16、FFN幅 $4d$ の場合です。

FLOPsの99.78%が行列積でした。softmaxもLayerNormもGELUも、合わせて0.22%です。Self-Attentionの本質は注意の重み付けにありますが、計算資源の観点では、Transformerはほぼ行列積を並べたものだと言えます。


5. 実測:dtypeと形状の罠

このセクションで分かること:精度を下げれば速くなるとは限らないこと、そして総FLOPsが同じでも形状で2割変わること。

ここで、はっきり失敗した計測を載せます。float16にすれば速くなるはずと思って回した結果です。

dtype 1行列のサイズ 時間 GFLOPS
float64 8.0 MB 31.4 ms 68.4
float32 4.0 MB 15.1 ms 142.5
float16 2.0 MB 5691 ms 0.38

float16はfloat32より約378倍遅いという結果になりました。素朴な3重ループより遅いです。

理由は単純で、OpenBLASにhalf精度のGEMMカーネルがなく、NumPyがループでソフトウェア実装に落ちるためです。半精度が速いのはTensor Coreのような専用ハードウェアがある場合の話で、CPUのNumPyでは逆効果になります。厨房で言えば、小さい包丁に持ち替えたのに、その包丁用のまな板が厨房になかった、という状況です。

形状の影響も測りました。いずれも総FLOPsは 2.68e8 で同一です。

形状 時間 GFLOPS
(512,512) × (512,512) 1.87 ms 143.4
(4096,64) × (64,512) 2.21 ms 121.3
(512,64) × (64,4096) 2.30 ms 116.9
(64,4096) × (4096,512) 2.42 ms 110.8

同じ仕事量なのに、細長い行列では最大23%のスループット低下が出ました。$K$ が小さいと、タイル内で使い回せる回数が減り、演算強度が落ちるためです。バッチサイズを小さくすると学習効率が落ちる現象の一部は、ここに起因します。


6. 行列積は次の世代のモデルにも残る

このセクションで分かること:行列積という視点が、Transformer以降のアーキテクチャを読むときにも使えること。

Transformerの代替として提案されている状態空間モデル系のアーキテクチャは、系列長に対する計算量を $O(L^2)$ から $O(L)$ に下げます。ただし、それだけでは実測速度は上がりません。

理由は本記事で見たとおりです。演算量が減っても、演算強度が下がって行列積の形にできなければ、ハードウェアの性能を引き出せません。厨房の総調理量を半分にしても、動線が悪ければ営業時間は縮まないのと同じです。

実際、Mamba-2で導入されたSSD: Structured State Space Dualityという枠組みの主眼は、計算量を下げることではなく再帰計算を行列積の形に書き換えることにあります。行列積を単位に考える視点は、Transformer以降のモデルを読むときにこそ効いてきます。


トラブルシューティング

症状 原因 対処
行列積が異様に遅い dtypeがfloat16やobject arr.dtype を確認し、float32かfloat64に変換する
@ でエラー: shapes not aligned 内側の次元 $K$ が不一致 A.shape[-1] == B.shape[-2] を確認する
期待の1/10しか速度が出ない BLASがスレッドを使えていない numpy.show_config() でBLAS名を確認、OMP_NUM_THREADS を見直す
バッチが小さいと極端に遅い $K$ や $M$ が小さく演算強度が不足 バッチをまとめる、あるいは複数の小行列を1回の呼び出しに束ねる
メモリ不足で落ちる 中間行列 $M \times N$ が巨大 出力をブロックに分けて計算する
np.dot@ で結果が違う 3次元以上でブロードキャスト規則が異なる 高次元では np.matmul の仕様を確認する

用語集

用語 説明
matmul 行列積。matrix multiplicationの略
GEMM General Matrix Multiplyの略。BLASにおける行列積の名称
BLAS 線形代数の基本演算を定めたインターフェース仕様。OpenBLASなどが実装
FLOPs 浮動小数点演算の回数。行列積では 2MNK
GFLOPS 1秒あたり10億回の浮動小数点演算。スループットの単位
演算強度 転送1バイトあたりのFLOPs。大きいほど計算律速に近づく
メモリ律速 演算器ではなくメモリ帯域が性能を決めている状態
タイル化 行列を小ブロックに分割し、キャッシュ上で使い切る最適化
C順序 / Fortran順序 多次元配列を1次元メモリに並べる際の順序。行優先か列優先か

学習ロードマップ

入門レベル

  1. 本記事の計測スクリプトを手元で回し、自分の環境のGFLOPSを測る
  2. numpy.show_config() で使用中のBLAS実装を確認する
  3. np.einsum で行列積を書き、@ と速度を比較する

中級レベル

  1. NVIDIAの Matrix Multiplication Background を読み、演算強度とops:byte比の関係を理解する
  2. タイルサイズを変えた自作のブロック行列積をNumPyで書き、キャッシュ効果を測る
  3. PyTorchtorch.matmul のCPU版とGPU版を比較する

発展レベル

  1. CUDAで素朴なGEMMカーネルを書き、共有メモリでタイル化して段階的に改善する
  2. Tensor Coreを使う実装を試し、float16やbfloat16の効果を測る
  3. 状態空間モデル系のアーキテクチャが、再帰計算をどうやって行列積に落としているかを追う

参考文献

NVIDIA Deep Learning Performance / Matrix Multiplication Background User's Guide — GEMMの定義、2MNK というFLOPsの数え方、演算強度とタイル量子化の解説。本記事のFLOPs計算はこの定義に従っています。

NVIDIA Deep Learning Performance / Get Started With Deep Learning Performance — 演算強度の定義と、演算がメモリ律速か計算律速かを判定する手順。

NVIDIA Deep Learning Performance / Linear/Fully-Connected Layers User's Guide — 全結合層をGEMMの3次元 M・N・K に対応づける方法と、バッチサイズが小さいときにメモリ律速になる具体例。

OpenBLAS 公式リポジトリ — 本記事の計測で使われたBLAS実装。対応アーキテクチャ別のカーネル一覧が確認できます。


関連記事

本記事は、Mamba系アーキテクチャを理解するための土台として書いた4本組の1本目です。

親記事にあたるのが以下の2本です。

NumPy / Python / GPU / CUDA の各記事もあわせてどうぞ。


計測スクリプトはすべて本文中に掲載しています。同じ環境でなくても、倍率と傾向は再現するはずです。もし手元で違う結果が出たら、その差分こそが環境固有の情報です。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?