行列積 — 定義からキャッシュラインまで

私がこれまで本番投入してきたモデル — トランスフォーマーの各層も、value ヘッドも、音声エンコーダも — その FLOP の大半をたったひとつの演算の中で使っている。 $A \in \R^{m \times k}$ 、 $B \in \R^{k \times n}$ に対し、積 $C = AB$ は成分ごとに次のように定義される:

c_{ij} = \sum_{p=1}^{k} a_{ip}\, b_{pj}, \qquad 1 \le i \le m,\ \; 1 \le j \le n.

これは $mnk$ 回の積和 — $2mnk$ FLOP — であり、定義はそのまま誰もが最初に書く行列積のコードになる。定義は3重ループにすぎない。性能はメモリの物語であり、この記事はそれをキャッシュラインまで掘り下げる。

2億6800万FLOPの使い方、3通り

$m = n = k = 512$ のとき、積のコストは $2 \cdot 512^3 \approx 2.7 \times 10^8$ FLOP。以下が定義そのままのコードで、ベンチマークスクリプトが計測しているのはまさにこの関数である:

def matmul_naive(A, B, n):
    C = [[0.0] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            acc = 0.0
            for p in range(n):
                acc += A[i][p] * B[p][j]
            C[i][j] = acc
    return C

私のデスクトップでは、Python のリスト上でこれが 9.3秒 かかる。一見「最適化」に見える書き換え — 同じ3重ループを NumPy 配列のスカラー添字 (A[i, p] * B[p, j]) で回す — は 41.9秒、素のリストの4.5倍も遅い。そして OpenBLAS に委譲される A @ B は、1スレッドで 5.1ミリ秒。同じ行列、同じ2億6800万FLOP、答えは $2.3 \times 10^{-13}$ の精度で一致する:

512³の乗算1回の対数スケール比較: NumPyスカラー添字41.9秒、純Pythonリスト9.3秒、OpenBLAS 5.1ミリ秒 — 図1 — 同じ乗算、3通りの実装。対数スケールに注意

このグラフには説明すべき点がふたつある。スキャンダルのほう — NumPy が素のリストに4.5倍負ける — の種明かしは簡単だ。A[i, p] のたびに C-API の境界を越えて Python の float オブジェクトをひとつ箱詰めする。1アクセスあたり約 120ナノ秒のオーバーヘッドが、2億6800万回。NumPy の契約は配列まるごとの演算であり、リストのように添字で触れば、機構の代金だけ払って機構を一度も使わないことになる。

興味深いのはもうひとつの数字だ。愚直なループと A @ B の間の 1800倍 — OpenBLAS は定義が要求する $2mnk$ FLOP を、同じ倍精度で、過不足なく実行している。このうちインタープリタの寄与はせいぜい2桁分 — 同じ3重ループを -O3 でコンパイルすれば、このクラスのコアではだいたい1〜2 GFLOP/s になる。残り、コンパイルしても消えない30倍余りの因子 — それがメモリである。

ベンチマーク環境

Intel i7-6700K(Skylake、4コア/8スレッド、ベース4.0 GHz/シングルコアターボ4.2 GHz)、キャッシュはコアあたり L1d 32 KB + L2 256 KB、共有 L3 8 MB、デュアルチャネル DDR4。Python 3.11.7、NumPy 2.4.2(scipy-openblas)、すべて float64。BLAS は OPENBLAS_NUM_THREADS=1 で1スレッドに固定(末尾のマルチスレッド計測を除く)。遅い2つのループは1回実行、それ以外は固定シードでベスト・オブ・5。結果は A @ B と照合済み: 最大絶対偏差 ≈ 2.3 × 10⁻¹³。

64バイトの真実

CPU がメモリから float64 を1個だけ読むことはない。読むのはキャッシュライン — 64バイト、double 8個 — であり、それをキャッシュ階層に保持する。このマシンではコアあたり L1d 32 KB、L2 256 KB、全コア共有の L3 8 MB。L1 ヒットのロードは約4サイクル、DRAM までの往復は数百サイクル。速い数値計算コードのすべては、ひとつの規則から導かれる: ラインを引いてきたら、そこに載っているものを使い切れ — そして追い出される前に再利用しろ。

NumPy は行列を行優先(C順)で格納する。行 $i$ は連続で、要素 $b_{pj}$ は $b_{p,j-1}$ の8バイト先にあるが、 $b_{p-1,j}$ からは行1本ぶん先にある。ここで定義の内側ループを見てほしい。A[i][p] は行に沿って進む — 逐次アクセスで、フェッチしたライン1本につき8個の double が全部役に立ち、ハードウェアプリフェッチャが得意とするパターンで先回りしてくれる。一方 B[p][j] は列を下へ、1ステップごとに4 KB ずつ跳ぶ。この下降は毎ステップ新しいキャッシュラインと新しいメモリページを開く。列を1本下りきるだけで512本の異なるライン(4 KB の有効データのために32 KB のトラフィック)に触れ、プリフェッチャには取り付く島もない。

この効果は行列積なしでも簡単に分離できる。8192×8192 の行列1枚 — 537 MB、どのキャッシュよりはるかに大きい — を、行沿いと列沿いで2回合計してみる:

M = rng.standard_normal((8192, 8192))

s = 0.0
for i in range(8192): s += M[i, :].sum()   # 行に沿って:  44 ms

s = 0.0
for j in range(8192): s += M[:, j].sum()   # 列を下って: 597 ms

同じバイト列、同じ加算で、13.6倍の差。行パスはプリフェッチャの後ろでキャッシュラインを丸ごとストリームする。列パスはそのパスで開いたライン1本につき double を1個しか使わず、しかも 64 KB ストライドが列1本につき8192の異なるページに触れる — TLB の容量をはるかに超え、キャッシュミスの上にページウォークの代金まで払う。(隣の列があとで各ラインの残り7個を L3 から回収する — ペナルティが13.6倍で済んでいるのはそのためだ。)行パスは 537 MB を44 msで運ぶ — 約 12.2 GB/s、このコアの実効読み取り帯域である。この数字を覚えておいてほしい。

演算強度 — 素朴なループが速くなれない理由

コアを縛る天井はふたつ: 計算の速さと、供給の速さだ。このコアは 4.2 GHz、256ビット FMA ポート2本で、倍精度のピークは

4.2\ \text{GHz} \times 2\ \tfrac{\text{FMA}}{\text{cycle}} \times 4\ \tfrac{\text{double}}{\text{FMA}} \times 2\ \tfrac{\text{FLOP}}{\text{double}} \;\approx\; 67\ \tfrac{\text{GFLOP}}{\text{s}}

— だが持続できる読み取りは約 12.2 GB/s しかない。どちらの天井が効くかを決めるのが演算強度: 移動1バイトあたりに実行する FLOP 数である。

素朴なループは — オペランドがキャッシュを超えたサイズでは — 両オペランドを使うたびにコアへ流し直す。積和1回につき8バイトの読み取りが2回、強度は $q = \tfrac{2\ \text{FLOP}}{16\ \text{B}} = 0.125$ 。12.2 GB/s のもとでは、このアクセスパターンである限り完全にベクトル化されたループでさえ $0.125 \times 12.2 \approx 1.5$ GFLOP/s 前後で頭打ちになる — 計算ピークの40分の1以下、インタープリタの1サイクルを数える前の話だ。この天井はコンパイルでは持ち上がらない。ループそのものが、構造からしてメモリ律速なのである。

速くなるために必要なのは FLOP を減らすことではない。1バイトあたりの FLOP を増やすことだ。

ブロッキング: √M のアイデア

処方箋は何十年も前からあり、今もあらゆる BLAS とあらゆる GPU 行列積カーネルを支えている: 小さなワーキングセットをキャッシュに常駐させ、追い出される前に再利用されるよう、問題をタイルに切れ。

行列を $b \times b$ のタイルに切り、 $C$ の各タイルを小さなタイル積の和として累積する:

for i0 in range(0, n, b):
    for j0 in range(0, n, b):
        # この C タイルは p ループの間ずっと再利用される
        for p0 in range(0, n, b):
            C[i0:i0+b, j0:j0+b] += A[i0:i0+b, p0:p0+b] @ B[p0:p0+b, j0:j0+b]

内側ループの1ステップが触るのは $b \times b$ タイル3枚 — $3b^2 \times 8$ バイト — で、その上で $2b^3$ FLOP を実行する:

q(b) \;=\; \frac{2b^3}{24\, b^2}\ \tfrac{\text{FLOP}}{\text{B}} \;=\; \frac{b}{12}\ \tfrac{\text{FLOP}}{\text{B}},

強度がタイルサイズとともに育つ。このコアをメモリ律速から計算律速へ持ち上げるには $q \gtrsim 67 / 12.2 \approx 5.5$ 、つまり $b$ はおよそ70 — そして float64 の 70×70 タイル3枚は 118 KB、256 KB の L2 に余裕で収まる。トリックはこれだけであり、この分野の文献でキャッシュ容量 $M$ がいつも平方根の下に現れるのはそのためだ。一辺 $b \sim \sqrt{M/3}$ のタイルは、総トラフィックを素朴な $\mathcal{O}(n^3)$ ワードから

\mathcal{O}\!\left(\frac{n^3}{\sqrt{M}}\right)

まで削る。これが古典アルゴリズムのどんな実行順序についても漸近最適であることは、Hong と Kung が1981年に証明した。キャッシュが大きいほどトラフィックは少ない — 平方根のぶんだけ。

本物のBLASが加えるもの

OpenBLAS はこのアイデアを執念で実装したものだ。キャッシュ階層ごとに調律されたタイルサイズ、オペランドのパッキング(各タイルを連続かつアラインされたバッファへコピーし、 $B$ の下降さえライン単位の逐次ストリームに変える)、そして $C$ の小ブロックをベクトルレジスタに置いたまま FMA を流し続ける最内側のレジスタ・マイクロカーネル。成果はグラフ1枚に収まる:

OpenBLASはサイズ256から4096まで51〜54 GFLOP/sを維持 — 破線の1コアピーク67 GFLOP/sの約4分の3の高さの平坦な線 — 図2 — 1スレッドの A @ B をサイズ横断で計測。平坦な線こそがすべて

平坦だ。 3枚の行列が L3 に収まる $n = 256$ から、268 MB のオペランド(結果も数えれば 400 MB)が DRAM に置かれる $n = 4096$ まで、スループットは 51〜54 GFLOP/s — 最大サイズで理論1コアピークの79% — を保つ。ブロッキングが FLOP対バイト比を、問題サイズではなくタイルの性質にしてしまうからだ。定義どおりのループは、行列がキャッシュ階層をひとつ超えるたびに、1 FLOP あたりの時間が延びていく。ブロック化されたループは意に介さない。

同じマシンからの脚注をふたつ、どちらも実測で。第一に並列性: 物理4コアに固定すると $n = 4096$ は 179 GFLOP/s — 3.4倍のスピードアップで、4倍に届かないぶんの主因はシングルコアターボより低い全コアターボ、それに L3 と DRAM チャネルの共有だ。デフォルト設定のままだと OpenBLAS はこのマシンで8スレッド — ハイパースレッド1本につき1スレッド — を立ち上げ、118 GFLOP/s まで落ちる。2本の SMT 兄弟スレッドは同じ物理コアの FMA パイプを分け合うので、余分なスレッドはスケジューリングのコストだけ足して計算を1つも足さない。数えるべきは物理コアである。第二に、ピークの79%という数字: これは成熟した BLAS がこのマイクロアーキテクチャの float64 で到達する水準で、欠けた2割はパッキングのトラフィック、タイルの端、ループの管理に消える。100%は誰にも出せない。

同じアイデアが、いちばん上まで

「L2」を「シェアードメモリ」に置き換えれば、これは CUDA の講義になる。GPU の行列積カーネルは $A$ と $B$ のタイルを各 SM のシェアードメモリに載せ、同期し、乗算し、進む — 同じ $\sqrt{M}$ の議論を、数千のスレッドが実行しているだけだ。テンソルコアはタイルをデータパスそのものに取り込む。ワープが固定サイズのフラグメントを供給し、ハードウェアが小さな密行列積を1命令として実行する。さらに一段上では、アテンション層も MLP ブロックもこうした積のバッチであり、だからこそ演算強度 — いまでは「重み1バイトあたりの FLOP」と言い換えられることが多い — が、H100 上のトランスフォーマーが計算律速か帯域律速かを今日も決めている。

定義は3重ループのまま、ビネが1812年に「行×列」の規則を書き下して以来、本質的に変わっていない。9.3秒と5.1ミリ秒のあいだにあるものはすべて、自分のキャッシュラインがどこにあるかを知っているかどうかだ。

再現する

matmul_bench.py — 標準ライブラリと NumPy だけの約130行の Python。遅い2つのループ(1回実行)、BLAS のサイズスイープとトラバーサルのデモ(ベスト・オブ・5)、マルチスレッド計測、すべて固定シード。計測対象は上に示した関数そのものだ。絶対値はあなたのマシンでは違ってくるが、比率は違わない。信じる前に、走らせてほしい。