NumPyのベクトル化はなぜ速いのか — 性能の理由を測って理解する

同じ「2つの配列の要素ごとの積」を計算するのに、for ループで書くと 33.5 ミリ秒、NumPy で a * b と書くと 0.21 ミリ秒。手元の環境で測ると、その差は 161 倍 ありました。やっている演算はまったく同じ 100 万回の掛け算です。CPU が急に速くなったわけでも、掛け算の回数が減ったわけでもありません。それなのに、なぜここまで違うのでしょうか。

この「161 倍」の正体を知らないまま高速化に取り組むと、だいたい間違った場所を直します。よくあるのが「とりあえず np.einsum に書き換える」「とりあえず out= を付ける」というやり方です。実際に測ってみると、einsum を素直に使ったら BLAS 版より 1.5 倍遅くなることもあれば、out=+= を付けたほうが遅くなることもあります(後で見るように、小さい配列では += が 0.93 倍に落ちました)。速度は「NumPy っぽいコードか」では決まりません。1 要素あたり何命令実行しているかと、メモリを何バイト往復させたかで決まります。

この視点は、たとえば衛星画像の前処理で数億画素を扱うとき、あるいは機械学習の学習ループで毎イテレーション距離行列を作り直すときに、そのまま効いてきます。前者では「1 画素あたりのバイト数」を減らす設計が支配的になり、後者では「巨大な中間配列を materialize しない」書き換えが 3 倍の差を生みます。どちらも、コードの見た目ではなくコストモデルを持っているかどうかで結果が変わります。

本記事では、この差を測定で分解します。単に「NumPy は速い」で終わらせず、どのコストがどの領域で支配的になるかを、要素数 $10^2$ から $10^7$ まで振った実測で切り分けていきます。

Pythonのforループでは1要素ごとにインタプリタが型を調べて新しいオブジェクトを作るため40.1ナノ秒かかるのに対し、NumPyのベクトル化では型が確定したCのループに1回だけ丸投げするので0.198ナノ秒で済むことを示した概念図

この図が本記事の全体像です。左右で計算している中身は同じ「要素ごとの掛け算」ですが、ループを Python がまわすか C がまわすかだけが違います。左では 1 要素ごとに「型を調べる → 掛ける → 新しいオブジェクトを作る」という往復が発生し、実測で 1 要素 40.1 ナノ秒。右では配列全体を C 側に 1 回渡すだけなので 0.198 ナノ秒、つまり 203 倍の差になります。以降の節では、この差がどんなコストの積み重ねでできているかを 1 つずつ剥がしていきます。

本記事の内容

  • ベクトル化の速度差を「インタプリタ/ボックス化/メモリ帯域/SIMD」の4コストに分解する
  • ルーフラインモデルで「演算律速かメモリ律速か」を事前に見積もる
  • for / リスト内包 / map / NumPy / einsum を要素数を振って実測し、ログログで比較する
  • 一時配列の生成コストを測り、out=+= が効く場面と効かない場面を切り分ける
  • C 連続 / F 連続とアクセス順序が食い違ったときのキャッシュミスを実測する

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

ベクトル化とは — 「1個ずつ」を「まとめて」に置き換える

ベクトル化を説明する前に、身近な比喩から入りましょう。段ボール 100 箱を倉庫から運び出す場面を想像してください。1 箱ずつ「倉庫のドアを開けて、箱を探して、抱えて、ドアを閉めて、運ぶ」を 100 回繰り返すのが Python のループです。一方、台車を持ち込んで「ドアを開けて、100 箱まとめて積んで、一度に運ぶ」のがベクトル化です。運ぶ箱の数は同じでも、ドアの開け閉め(オーバーヘッド)が 100 回から 1 回に減る。これがベクトル化の本質です。

もう少し正確に言い換えます。ベクトル化とは、「要素ごとの制御をPythonから追い出し、配列全体に対する1回の関数呼び出しに置き換える」 操作です。

$$ \underbrace{\text{for } i: \; c_i = a_i \times b_i}_{\text{Python が } n \text{ 回まわす}} \quad\Longrightarrow\quad \underbrace{\bm{c} = \bm{a} \odot \bm{b}}_{\text{C が } n \text{ 回まわす}} $$

ここで $\odot$ はアダマール積(要素ごとの積)です。式の上では同じですが、ループを誰がまわすかが違います。左辺は CPython のバイトコードインタプリタが、右辺は NumPy のコンパイル済み C ループがまわします。この「誰がまわすか」の違いが、後で見る 100 倍超の差を生みます。

大事なのは、ベクトル化は「計算量を減らす最適化ではない」という点です。$O(n)$ の計算はベクトル化しても $O(n)$ のままで、定数倍だけが改善します。しかしその定数倍が 3 桁あるので、実務では計算量の改善に匹敵するインパクトを持ちます。逆に言えば、定数倍の出どころを理解していないと、改善の上限が見えません。

では、その定数倍はどこから来ているのでしょうか。1 つの原因ではなく、独立した 4 つのコストが積み重なっています。次節でそれを 1 つずつ剥がしていきます。

なぜ速いのか — 4つのコストに分解する

コスト1: インタプリタのディスパッチ

out[i] = a[i] * b[i] という 1 行を CPython が実行するとき、内部では概ね次のことが起きます。

  1. バイトコードを 1 命令ずつフェッチし、巨大な switch で分岐する
  2. ai をスタックに積み、BINARY_SUBSCR を実行する
  3. a の型を調べ、__getitem__ に相当するスロットを引く
  4. 取り出した要素は PyObject* なので、参照カウントを増やす
  5. 掛け算では両オペランドの型を調べ、nb_multiply スロットを引く
  6. 結果として 新しい float オブジェクトをヒープに確保する
  7. リストに格納し、古い値の参照カウントを減らす(必要なら解放する)

掛け算そのもの(浮動小数点乗算 1 命令)は、この一連の流れのごく一部でしかありません。型チェックと動的ディスパッチが、要素ごとに毎回発生します。

NumPy の ufunc は、この型チェックを配列に対して 1 回だけ行います。afloat64 の連続配列、bfloat64 の連続配列だと分かった瞬間に、「float64 同士の乗算専用の C ループ」を 1 本選んで、そこに丸投げします。ループの中身は while (n--) *out++ = *in1++ * *in2++; のような、型が確定した単純な機械語です。要素あたりのオーバーヘッドはほぼゼロになります。

コスト2: ボックス化とポインタ追跡

2 つ目は、データがメモリ上でどう並んでいるかの違いです。CPython の floatボックス化されたオブジェクトで、値だけでなく型ポインタと参照カウントを持ちます。実測すると 1 個 24 バイトです。

Python のリストは、この 24 バイトのオブジェクトへのポインタ(8 バイト)の配列です。つまり [0.1, 0.2, ...] という 100 万要素のリストは、

  • ポインタの配列: $10^6 \times 8 = 8$ MB
  • float オブジェクト本体: $10^6 \times 24 = 24$ MB(ヒープ上に散らばる)

の合計 32 MB を占め、しかも本体はヒープのあちこちに散在しています。値を 1 つ読むたびにポインタをたどる必要があり、CPU のプリフェッチャは次にどこを読むか予測できません。

対して NumPy の ndarray は、$10^6 \times 8 = 8$ MB の1 本の連続したバイト列です。メモリ量が 4 分の 1 で済むだけでなく、アドレスが等間隔に並ぶのでプリフェッチャが完璧に働きます。この「連続していること」が、次の 2 つのコストの前提になります。

Pythonのリストはポインタ配列とヒープに散らばる24バイトのfloatオブジェクトで合計32MBを占めるのに対し、NumPyのndarrayは8MBの連続したバイト列で128バイトのキャッシュラインに16要素が丸ごと収まることを示した図

左右を見比べると、同じ 100 万個の数値でも「アクセス 1 回で何をしなければならないか」がまったく違うことが分かります。リストでは値に届くまでに必ずポインタを 1 段たどる必要があり、飛び先が不規則なので CPU は次に読む場所を予測できません。ndarray では 8 バイトごとに次の値が並ぶので、128 バイトのキャッシュライン 1 本を引くだけで 16 要素が手に入ります。この「1 回の転送で何要素まかなえるか」は、記事の後半でストライドを振って実測するときにもう一度出てきます。

コスト3: メモリ帯域

要素ごとの演算では、CPU は演算よりもデータの搬入出で忙しいというのが実情です。c = a * b を $n$ 要素について行うとき、

  • 読み込み: $a$ と $b$ で $2 \times 8n$ バイト
  • 書き込み: $c$ で $8n$ バイト
  • 演算: 乗算 $n$ 回

つまり 1 回の浮動小数点乗算あたり 24 バイトを動かしています。現代の CPU は 1 コアあたり毎秒数十 GFLOPS の演算能力を持ちますが、メモリ帯域は数十 GB/s しかありません。24 バイト/1 演算という比率では、演算器は圧倒的に暇です。要素ごとの演算はメモリ律速(memory-bound) である、と覚えておいてください。

この事実は重要な帰結を持ちます。「メモリ律速なら、演算を速くしても意味がない。減らすべきはメモリの往復回数だ」ということです。後で out= や一時配列の議論をするとき、この基準がそのまま判断材料になります。

コスト4: SIMD

4 つ目は SIMD(Single Instruction, Multiple Data)です。1 つの命令で複数のデータを同時に処理する CPU 機能で、x86 なら AVX2(256 ビット = double 4 個)や AVX-512(8 個)、Arm なら NEON/ASIMD(128 ビット = double 2 個)が該当します。NumPy 1.17 以降は「universal SIMD」という仕組みで、実行時に CPU の機能を検出して最適な実装を選びます。

自分の環境で何が有効かは、次のように確認できます。

import numpy as np

# NumPy が検出した CPU の SIMD 機能
feats = np.core._multiarray_umath.__cpu_features__
print("有効な機能:", [k for k, v in feats.items() if v])
print("NumPy バージョン:", np.__version__)

この記事の測定環境(Apple Silicon, arm64)では NEON, ASIMD, ASIMDDP などが有効で、AVX2 系はすべて False でした。NEON は 128 ビット幅なので double なら一度に 2 要素です。x86 の AVX-512 環境なら 8 要素になり、演算律速な処理では差が出ます。ただし前節で見たとおり要素ごとの演算はメモリ律速なので、SIMD の幅が広くても要素ごと演算はそれほど速くなりません。SIMD が効くのは、行列積のように演算密度が高い処理です。

ここまでで 4 つのコストが揃いました。ではこの 4 つのうち、どれが自分の計算で支配的になるのか。それを事前に見積もる道具が次のルーフラインモデルです。

性能の見取り図 — ルーフラインで「どこで詰まるか」を予測する

高速化で最初にやるべきことは、プロファイラを起動することではなく、理論上の上限を見積もることです。上限が分かれば「あと何倍縮められるのか」「もう限界なのか」が判断できます。

ある処理が実行する浮動小数点演算の総数を $W$[flop]、メインメモリとやり取りするバイト数を $Q$[byte]とします。CPU のピーク演算性能を $F$[flop/s]、メモリ帯域を $B$[byte/s]とすると、実行時間 $T$ の下限は次の 2 つの制約のうち厳しいほうで決まります。

$$ T \ \geq\ \max\!\left( \frac{W}{F},\ \frac{Q}{B} \right) $$

この式を「1 秒あたり何 flop 出せるか」に直します。両辺の逆数を取って $W$ を掛けると、

$$ \frac{W}{T} \ \leq\ \min\!\left( F,\ \ \frac{W}{Q} \cdot B \right) $$

となります。ここで現れた $I = W/Q$[flop/byte]を 演算強度(arithmetic intensity) と呼びます。演算強度は「1 バイト運ぶごとに何回計算するか」という、アルゴリズム固有の量です。上の式は「達成できる性能は、ピーク性能 $F$ と、演算強度に帯域を掛けた $I \cdot B$ の小さいほう」と読めます。これがルーフラインモデルです。

実際に 2 つのケースで $I$ を計算してみましょう。まず要素ごとの積 c = a * b です。$W = n$、$Q = 24n$ なので、

$$ I_{\text{elementwise}} = \frac{n}{24n} = \frac{1}{24} \approx 0.042 \ \text{[flop/byte]} $$

きわめて小さい値です。この環境で実測した飽和帯域は $B = 100$ GB/s だったので、達成しうる性能は $0.042 \times 100 \times 10^9 \approx 4.2$ GFLOPS。同じマシンで行列積を測ると 169 GFLOPS 出るので、その 2.5% しか使えない計算になります。要素ごと演算は、原理的にメモリ帯域でしか語れないわけです。

次に $n \times n$ の行列積 $\bm{C} = \bm{A}\bm{B}$ です。演算は $W = 2n^3$(積と和)、データは $\bm{A}, \bm{B}, \bm{C}$ の 3 枚で $Q = 3 \times 8n^2 = 24n^2$(キャッシュが完璧に効いた理想の場合)。したがって、

$$ I_{\text{matmul}} = \frac{2n^3}{24n^2} = \frac{n}{12} \ \text{[flop/byte]} $$

$n$ に比例して増えます。$n = 4000$ なら $I \approx 333$ で、要素ごと演算の 8000 倍です。だから行列積は演算律速(compute-bound) になり、SIMD や FMA、マルチスレッドが素直に効きます。BLAS ライブラリが何十年もチューニングされ続けているのは、ここに投資する価値があるからです。

この 2 つのケースを、実測値で 1 枚の図に重ねてみます。

実測した帯域100GB/sと行列積の実測性能169GFLOPSで描いたルーフライン図。要素ごとの積は演算強度1/24で4.2GFLOPS、3000角の行列積は演算強度250で169GFLOPSに位置し、分かれ目は演算強度1.68であることを示す

2 つの点はどちらも屋根の線にぴったり載っています。つまり NumPy はどちらの計算でも「そのアルゴリズムの演算強度で許される限界」まで出し切っており、残っているのは書き方の問題ではなく $I$ そのものを上げる余地だけだということです。分かれ目は $I = F/B = 169/100 \approx 1.68$[flop/byte]で、これより左は何をしてもメモリ帯域が天井、右に来て初めて演算器の性能が意味を持ちます。要素ごとの積は分かれ目の 40 分の 1 の位置にあり、絶望的にメモリ側です。

この $I$ の見積もりが、高速化の方針を決めます。

演算強度 $I$ 支配要因 有効な手
$I \lesssim 1$(要素ごと、リダクション) メモリ帯域 配列の往復回数を減らす、dtype を小さくする、融合する
$I \gtrsim 10$(行列積、畳み込み) 演算性能 BLAS に落とす、スレッド数を確保する、SIMD 幅を活かす

ここまでは理論です。では実際に測ってみて、この見立てが合っているかを確かめましょう。まずは測定の道具を用意します。

計測の作法 — ベンチマーク関数を作る

測定は雑にやると簡単に嘘をつきます。特に NumPy では、次の 4 つが典型的な落とし穴です。

  1. ウォームアップ不足 — 初回呼び出しには BLAS のスレッドプール起動や、np.empty が返した仮想メモリのページフォルトが混ざります
  2. 平均値を使う — OS のスケジューリングで時々大きな外れ値が出ます。最小値のほうが「邪魔が入らなかったときの真の実行時間」に近づきます
  3. キャッシュの状態が揃っていない — 直前に大きな配列を触ったかどうかで 2 倍変わります
  4. 測定対象が消える — 結果を使わないと、遅延評価やコンパイラ最適化で消える場合があります(NumPy では起きにくいですが、返り値は受け取る習慣を)

これらを踏まえた最小限のベンチマーク関数を作ります。

import timeit
import numpy as np


def bench(f, target=0.05, repeat=7, warmup=2):
    """f() の実行時間[秒]を返す。
    target: 1セットあたりに費やしたい時間[秒](自動で試行回数を決める)
    repeat: セットの繰り返し数。最小値を採用する
    warmup: 計測前に捨てる実行回数(ページフォルト・BLAS起動を除去)
    """
    for _ in range(warmup):
        f()
    # 1回の実行時間をざっくり測って、number を決める
    t1 = timeit.timeit(f, number=1)
    number = max(1, int(target / max(t1, 1e-9)))
    times = timeit.repeat(f, repeat=repeat, number=number)
    return min(times) / number   # 平均でなく最小値


def fmt(t):
    """秒を読みやすい単位に整形"""
    if t < 1e-6:
        return f"{t*1e9:8.2f} ns"
    if t < 1e-3:
        return f"{t*1e6:8.2f} us"
    return f"{t*1e3:8.3f} ms"

この bench はミリ秒級の処理なら数百回、秒級なら 1 回に自動で試行回数を合わせます。min を採るのは、ノイズが「常に時間を増やす方向にしか働かない」からです。平均を採ると、たまたま別プロセスが走った回に引きずられます。

以降のすべての測定は、この bench を使い、Python 3.10.17 / NumPy 1.26.4 / arm64(Apple Silicon, 18コア, キャッシュライン 128 バイト, L1d 128 KB, L2 16 MB) で行いました。掲載する絶対値は環境に強く依存します。読者の環境では違う数字が出るはずなので、比の傾向に注目してください。同じスクリプトを手元で走らせて、自分のマシンの数字を持つのが理想です。

道具が揃ったので、いよいよ本題の測定に入ります。

実測1: for / 内包表記 / map / NumPy をログログで比べる

最初の実験は、もっとも単純な要素ごとの積です。同じ計算を 4 通りに書き、要素数 $n$ を $10^2$ から $10^7$ まで 10 倍刻みで振ります。

import numpy as np

rng = np.random.default_rng(0)
results = {}

for n in [10**2, 10**3, 10**4, 10**5, 10**6, 10**7]:
    a = rng.random(n)
    b = rng.random(n)
    la, lb = a.tolist(), b.tolist()

    def by_for():                      # 素朴な for ループ
        out = [0.0] * n
        for i in range(n):
            out[i] = la[i] * lb[i]
        return out

    def by_comp():                     # リスト内包表記
        return [x * y for x, y in zip(la, lb)]

    def by_map():                      # map + 演算子
        return list(map(float.__mul__, la, lb))

    def by_numpy():                    # ベクトル化
        return a * b

    row = {}
    if n <= 10**6:                     # 10^7 の Python ループは現実的でないので除外
        row["for"] = bench(by_for, target=0.2)
        row["内包"] = bench(by_comp, target=0.2)
        row["map"] = bench(by_map, target=0.2)
    row["NumPy"] = bench(by_numpy)
    results[n] = row
    print(n, {k: fmt(v) for k, v in row.items()})

手元での結果は次のとおりでした(単位はミリ秒、最小値)。

$n$ for ループ リスト内包 map NumPy for/NumPy
$10^2$ 0.0026 0.0019 0.0026 0.00016 16
$10^3$ 0.0298 0.0180 0.0253 0.00037 80
$10^4$ 0.403 0.177 0.257 0.00158 255
$10^5$ 3.051 1.689 2.433 0.0145 210
$10^6$ 33.49 20.07 28.50 0.208 161
$10^7$ 2.160

この表を両対数でプロットすると、傾きと切片から性質が読めます。

4通りの書き方の実行時間を両対数でプロットした図と、forループがNumPyの何倍遅いかを要素数に対してプロットした図。4本とも傾き1で計算量は同じだが縦位置が違うこと、速度比は10の4乗で255倍と最大になりその後161倍まで下がることを示す

左の両対数プロットでは、4 本の線がどれも傾き 1 の直線に乗ります。つまりどの書き方も計算量は $O(n)$ で同じであり、違うのは縦方向の位置(定数倍)だけです。これはベクトル化が「計算量の改善ではなく定数倍の改善」であることの、目に見える証拠です。

右の速度比のグラフはもっと示唆的です。$n = 100$ では 16 倍しかないのに、$n = 10^4$ で 255 倍まで上がり、そこからは逆に下がって $n = 10^6$ で 161 倍になります。山が立つのです。

左肩が低いのは、NumPy 側にも固定費があるからです。a * b は Python レベルで見ると 1 回の関数呼び出しですが、その中で引数の型チェック、ブロードキャスト形状の解決、出力配列の確保、ループ関数の選択が走ります。実際に $n$ を 1 から振って測ると、$n \le 100$ ではどの $n$ でも一定の 165 ナノ秒前後、$n = 300$ でようやく 213 ns と増え始めます。つまり 1 回の a * b には約 0.17 マイクロ秒の固定費があり、$n = 100$ ではそれが実行時間のほぼ全部です。要素数が数百以下の配列では、ベクトル化のうまみはあまり出ません。これは意外と実務で刺さるポイントで、「小さい配列に対して NumPy 関数を何万回も呼ぶ」パターンは、むしろ純 Python と大差なくなります。

では右肩が下がるのはなぜでしょうか。次項で見るように、NumPy 側が大きい $n$ でメモリ帯域に張り付くからです。Python ループは 1 要素あたりの仕事が重すぎてメモリ帯域には届かないので、$n$ を増やすと NumPy だけが先に頭打ちになり、比が縮みます。

大きい配列で NumPy が頭打ちになる理由

NumPy 単体の 1 要素あたりコストを、$n$ を細かく振って測ってみます。

NumPyのa*bの1要素あたり時間を要素数に対してプロットした図と、同じ測定を実効メモリ帯域に換算した図。3万要素で0.147ナノ秒・163GB/sと最も速く、大きくなると0.239ナノ秒・100GB/sに漸近することを示す

左のグラフは U 字になります。$n = 1000$ では固定費が効いて 0.374 ns/要素と割高、$n = 3 \times 10^4$(作業領域 0.7 MiB)で 0.147 ns/要素と最速になり、そこから $n$ を増やすと 0.239 ns/要素まで悪化してそこで完全に横ばいになります。悪化幅はわずか 1.6 倍で、「キャッシュから落ちた瞬間に何十倍も遅くなる」ようなことは起きていません。

右のグラフに換算すると理由が一目で分かります。$c = a \times b$ は 1 要素あたり 24 バイト動かすので、実効帯域は $24n/T$ です。これを描くと、小さい $n$ では 163 GB/s(キャッシュの帯域)まで出て、$n$ が大きくなるにつれて 100 GB/s に漸近して止まります。この 100 GB/s こそ、前節のルーフラインで屋根として使った値です。つまり大きい配列での a * b は、もう NumPy の実装の問題ではなく、メモリの物理的な速度そのものを測っていることになります。ここから先を速くしたいなら、書き方を変えるのではなく「動かすバイト数を減らす」しかありません。

ループのコストを分解する

33.5 ミリ秒のうち、何にいくら使われているのでしょうか。ループの中身を段階的に削って測ると分解できます。

import numpy as np

n = 10**6
rng = np.random.default_rng(0)
a = rng.random(n); b = rng.random(n)
la, lb = a.tolist(), b.tolist()

def empty_loop():                       # ループの骨格だけ
    for i in range(n):
        pass

def index_only():                       # + 添字アクセス1回
    s = 0.0
    for i in range(n):
        s = la[i]
    return s

def full_loop():                        # + 乗算と代入
    out = [0.0] * n
    for i in range(n):
        out[i] = la[i] * lb[i]
    return out

for name, f in [("空ループ", empty_loop),
                ("+ la[i] 読み出し", index_only),
                ("+ 乗算と out[i] 代入", full_loop)]:
    t = bench(f, target=0.3, repeat=3)
    print(f"{name:<22} {fmt(t)}  ({t/n*1e9:6.1f} ns/要素)")

t_np = bench(lambda: a * b)
print(f"{'NumPy a*b':<22} {fmt(t_np)}  ({t_np/n*1e9:6.3f} ns/要素)")

出力はこうなりました。

空ループ                  5.564 ms  (   5.6 ns/要素)
+ la[i] 読み出し         14.633 ms  (  14.6 ns/要素)
+ 乗算と out[i] 代入     40.087 ms  (  40.1 ns/要素)
NumPy a*b                 0.198 ms  (  0.198 ns/要素)

ループの中身を段階的に足したときの1要素あたり時間を対数軸の棒グラフで示した図と、Pythonループの内訳をループの骨格5.6ナノ秒・添字アクセス9.1ナノ秒・乗算とオブジェクト生成25.5ナノ秒に積み上げて示した図

この分解は雄弁です。掛け算を1回もしていない空のループだけで、すでに 1 要素あたり 5.6 ナノ秒かかっています。NumPy が乗算まで含めて 0.198 ns なので、「何もしないループ」の時点で NumPy の 28 倍遅いわけです。さらに添字アクセスを 1 つ足すと 14.6 ns(+9.1 ns)、掛け算と代入を足すと 40.1 ns(+25.5 ns)。右のグラフはこの増分を積み上げたもので、増分の大半は浮動小数点乗算そのものではなく、PyObject の取り出し・型ディスパッチ・新しい float オブジェクトの確保と破棄だと読めます。

つまり、「ループ内の演算を工夫して減らす」という素朴な最適化は、Python ループでは効果が薄いということです。1 要素 40.1 ns のうち浮動小数点乗算は 1 ns にも満たないので、演算を半分にしても全体は 1% も変わりません。削るべきはループそのものです。この認識が、ベクトル化を「小手先の技」ではなく「構造の変更」として捉えるための出発点になります。

なお、リスト内包が for より一貫して速い($n=10^6$ で 20.07 ms 対 33.49 ms、1.7 倍)のは、内包表記が専用のバイトコードを使い、リストの伸長を C レベルで行うためです。map は両者の中間(28.50 ms)で、これは float.__mul__ の呼び出しが 1 要素ごとに発生するぶん、内包表記の BINARY_MULTIPLY より重いからです。とはいえ、この 3 者の差はせいぜい 1.7 倍。NumPy との 161 倍に比べれば誤差のようなもので、「for か内包表記か」を悩むより「ループを消せるか」を考えるほうが桁で効くということです。

ループを消すところまでは分かりました。しかし NumPy に書き換えたあとも、まだ 3 倍以上縮む余地が残っていることがあります。次はその原因である一時配列を見ていきます。

実測2: 一時配列とメモリ帯域

d = (a + b) * c - a のような式を書くと、NumPy は二項演算ごとに新しい配列を作ります。概念的には次のように分解されます。

$$ \bm{t}_1 = \bm{a} + \bm{b}, \qquad \bm{t}_2 = \bm{t}_1 \odot \bm{c}, \qquad \bm{d} = \bm{t}_2 – \bm{a} $$

$\bm{t}_1, \bm{t}_2$ が 一時配列(temporary) です。$n = 10^7$ なら 1 本 80 MB。式の途中で 80 MB の配列が生まれては捨てられます。メモリ律速の世界では、これは無視できないコストに見えます。

ここで多くの入門記事は「だから out=+= でインプレース化しましょう」と書きます。しかしそれは半分しか正しくありません。実際に測って確かめます。

割り当てはタダに近い、ゼロ埋めは高い

まず、配列の確保そのもののコストを測ります。

import numpy as np

for n in [10**5, 10**6, 10**7]:
    t_empty = bench(lambda: np.empty(n), repeat=9)
    t_zeros = bench(lambda: np.zeros(n), repeat=9)
    t_fill = bench(lambda: np.empty(n).fill(1.0), repeat=9)
    print(f"n={n:>9}  np.empty {fmt(t_empty)} / np.zeros {fmt(t_zeros)}"
          f" / empty+fill {fmt(t_fill)}")
n=   100000  np.empty     0.16 us / np.zeros     4.01 us / empty+fill     6.39 us
n=  1000000  np.empty     0.51 us / np.zeros    32.94 us / empty+fill    62.19 us
n= 10000000  np.empty     0.51 us / np.zeros   323.51 us / empty+fill   613.10 us

np.emptyの時間が要素数によらず0.5マイクロ秒未満で一定なのに対し、np.zerosとempty+fillは要素数に比例して伸びることを示した対数グラフと、1000万要素で温かいバッファへのfillが656.9マイクロ秒・冷たいバッファへのfillが681.7マイクロ秒とわずか4%しか違わないことを示す棒グラフ

ここに決定的な事実が現れています。np.empty は要素数に依存せず 0.5 マイクロ秒程度で終わるのです。80 MB の確保でも 0.51 us。これは OS に「この仮想アドレス空間を使う」と宣言するだけで、物理メモリの割り当てが実際に触るまで遅延されるからです。一方 np.zeros は $n = 10^7$ で 324 us と、$n$ に比例して伸びます。empty + fill はさらに重い 613 us で、これは 80 MB を実際に書き込むコスト($80\ \text{MB} / 613\ \mu\text{s} \approx 130$ GB/s、ほぼメモリ帯域そのもの)です。

右のグラフは「では冷たいページに初めて書き込むのはどれだけ高くつくのか」を切り分けたものです。すでに使い回している温かい配列への fill が 656.9 us、確保したての冷たい配列への fill が 681.7 us で、差はわずか 4%。初回のページフォルト代は、80 MB を書き込むメモリトラフィックの前ではほとんど誤差だということです。

この結果から導かれる結論は明快です。「一時配列が遅い」の正体は、確保そのものでもページフォルトでもなく、そのバイトを実際に読み書きするメモリトラフィックである。 そして out= を付けても、書き込むバイト数は 1 バイトも減りません。この時点で、インプレース化の効果には低い天井があると予想できます。

out=+= はどこまで効くか

では実際に測りましょう。

import numpy as np

rng = np.random.default_rng(0)
for n in [10**5, 10**6, 10**7, 3*10**7]:
    x = rng.random(n); s = rng.random(n)
    o = np.empty(n); o.fill(0.0)          # 出力先のページを事前に触っておく
    t_new = bench(lambda: x + s, repeat=9)                 # 新しい配列を作る
    t_inp = bench(lambda: np.add(x, s, out=x), repeat=9)   # x を書き換える
    t_out = bench(lambda: np.add(x, s, out=o), repeat=9)   # 既存の o に書く
    print(f"n={n:>9}  x+s {fmt(t_new)} / x+=s {fmt(t_inp)} / out=o {fmt(t_out)}")
n=   100000  x+s     0.015 ms / x+=s     0.016 ms / out=o     0.015 ms
n=  1000000  x+s     0.203 ms / x+=s     0.173 ms / out=o     0.202 ms
n= 10000000  x+s     2.141 ms / x+=s     1.850 ms / out=o     2.140 ms
n= 30000000  x+s     6.778 ms / x+=s     5.822 ms / out=o     6.532 ms

新規確保・xへの上書き・別配列への出力の3通りを4つの要素数で比べた対数棒グラフと、新規確保を1としたときの速さの比を示す棒グラフ。上書きは最大1.17倍にとどまり、10万要素では0.93倍と逆に遅くなることを示す

期待していたほど劇的ではありません。x += s にあたる out=x は新規確保より 1〜2 割速い($n = 10^6$ 以上で 1.16〜1.17 倍)だけで、out=o に至ってはほぼ同じ(1.00〜1.04 倍)です。しかも $n = 10^5$ では +=0.93 倍と逆に遅くなっています。

理由は 2 つあります。第一に、前項で見たとおり確保はほぼ無料です。第二に、NumPy は解放した一時配列のバッファを内部でキャッシュして再利用するので、同じ式を繰り返し評価すると「新規確保」も実質は温かいバッファの使い回しになります。out=x だけがわずかに勝つのは、読み込む x と書き込む先が同じ領域なので、触るキャッシュラインの本数が 3 本ぶんから 2 本ぶんに減るからです。逆に out=ox, s, o の 3 領域を触るので、新規確保と同じだけのトラフィックになります。

「冷たいバッファに out= を渡すと遅くなる」という話も検証しました。$n = 10^7$ で (x + y) * z - x を測ると、素直に書いた場合が 6.34 ms、out= で温かいバッファを使い回した場合が 6.37 ms、確保したてのバッファを out= に渡した場合が 6.41 ms。差は 1% 程度しかありません。前項で見たとおりページフォルト代がメモリトラフィックに埋もれるうえ、そもそも NumPy がバッファを再利用するので「本当に冷たいバッファ」はループの中ではめったに生じないのです。

まとめると、インプレース化で得られるのはせいぜい 1〜2 割であり、しかも小さい配列では逆効果になりうる。これがこの節の実測から言えることです。

NumPyの一時配列消去(temporary elision)

もうひとつ、知らないと損をする挙動があります。NumPy 1.13 以降、参照カウントが 1 の一時配列に対する演算は、暗黙にインプレース化されます(temporary elision)。(a + b) * c(a + b) は誰も名前を持っていないので、* c はその一時配列を上書きして使い回します。

これは tracemalloc で確認できます。

import numpy as np, tracemalloc

rng = np.random.default_rng(0)
M, D = 2000, 3
P = rng.random((M, D)); Q = rng.random((M, D))

def one_expression():                      # 一気に書く
    return np.sqrt(((P[:, None, :] - Q[None, :, :])**2).sum(-1))

tracemalloc.start(); one_expression()
_, peak1 = tracemalloc.get_traced_memory(); tracemalloc.stop()

tracemalloc.start()                        # 途中結果を変数に束縛する
diff = P[:, None, :] - Q[None, :, :]
sq = diff**2
res = sq.sum(-1)
_, peak2 = tracemalloc.get_traced_memory(); tracemalloc.stop()

print(f"1つの式にまとめた場合      : ピーク {peak1/2**20:6.1f} MiB")
print(f"途中結果を変数に入れた場合 : ピーク {peak2/2**20:6.1f} MiB")
print(f"理論値: 差分配列 {M*M*D*8/2**20:.1f} MiB, 結果 {M*M*8/2**20:.1f} MiB")
1つの式にまとめた場合      : ピーク  122.1 MiB
途中結果を変数に入れた場合 : ピーク  213.7 MiB
理論値: 差分配列 91.6 MiB, 結果 30.5 MiB

1つの式にまとめた場合のピークメモリ122.1MiBと途中結果に名前を付けた場合の213.7MiBを比べた棒グラフと、参照カウントが1なら二乗が差分配列を上書きできる仕組みを段階図で示した説明

同じ計算なのに、ピークメモリが 122.1 MiB と 213.7 MiB で 1.75 倍違います。しかも差はぴったり $213.7 – 122.1 = 91.6$ MiB、つまり差分配列 1 枚ぶんです。1 つの式にまとめた場合は $91.6 + 30.5 = 122.1$ MiB になっており、二乗の一時配列が消えているdiff を上書きした)ことが分かります。変数に束縛した瞬間、diff の参照カウントが 1 を超えるので elision が無効になり、91.6 MiB がもう 1 枚必要になりました。右の図は、この上書きが起きるか起きないかを段階ごとに追ったものです。

読みやすさのために途中結果へ名前を付けるのは良い習慣ですが、巨大配列に対しては、名前を付けること自体がメモリを増やす。これは覚えておく価値があります。逆に言えば、名前を付けたい場合は del diff で明示的に手放すか、最初から out= で書くかを選ぶことになります。

本当に効くのは「巨大な中間を作らないこと」

ここまでの測定をまとめると、一時配列に関する処方箋は次のように整理できます。

やること 実測した効果 理由
+=out=x)で確保を避ける 1〜2 割(小さい配列では 0.93 倍と悪化) 確保自体がほぼ無料なので上限が低い
別配列への out=o ほぼ変わらない(1.00〜1.04 倍) 触る領域が 3 つのままでトラフィックが減らない
冷たいバッファに out= ほぼ変わらない(1% 程度) ページフォルト代がメモリ帯域に埋もれる
途中結果に名前を付けない ピークメモリ 1.75 倍改善 temporary elision が効く
中間配列そのものを消す 最大 7.9 倍(次節以降で実測) 動かすバイト数が桁で減る

最後の行が本命です。out= で 1 割を削るより、そもそも $M^2 D$ 要素の中間配列を作らない書き方に変えるほうが桁違いに効きます。その具体例を、次の 2 つの節(einsum とペア距離)で見ていきます。

実測3: einsum — 添字で縮約を書くと何が変わるか

np.einsum は、アインシュタインの縮約記法を文字列で書くと、その通りの計算をしてくれる関数です。たとえば行列積は $C_{ik} = \sum_j A_{ij} B_{jk}$ なので 'ij,jk->ik' と書きます。出力側の添字に現れない添字は、自動的に和を取って消えるというのが唯一のルールです。

$$ \underbrace{\texttt{‘ij,jk->ik’}}_{j \text{ は出力にないので } \sum_j} \qquad \underbrace{\texttt{‘ij,ji->’}}_{i,j \text{ ともに消えるのでスカラー}} $$

なぜこれが性能の話に関係するかというと、einsum「何を計算したいか」だけを宣言し、「どういう中間配列を経由するか」を指定しないからです。うまく使えば中間を作らずに済みます。

trace(AB) の罠

典型例が $\mathrm{tr}(\bm{A}\bm{B})$ です。定義どおりに書けば、まず $\bm{A}\bm{B}$ を計算して対角和を取ります。しかしトレースの定義を展開すると、

$$ \mathrm{tr}(\bm{A}\bm{B}) = \sum_i (\bm{A}\bm{B})_{ii} = \sum_i \sum_j A_{ij} B_{ji} $$

となり、対角成分しか使わないのに $n \times n$ の行列積をまるごと計算するのは無駄だと分かります。右辺は $A_{ij} B_{ji}$ の総和なので $O(n^2)$ です。$\bm{A}\bm{B}$ の計算は $O(n^3)$ なので、$n$ 倍の差があります。

import numpy as np

rng = np.random.default_rng(0)
B = rng.random((3000, 3000))

t1 = bench(lambda: np.trace(B @ B), target=1.0, repeat=3)      # O(n^3)
t2 = bench(lambda: np.einsum('ij,ji->', B, B), target=1.0, repeat=3)  # O(n^2)
t3 = bench(lambda: (B * B.T).sum(), target=1.0, repeat=3)      # O(n^2) だが一時配列あり

print(f"np.trace(B @ B)          {fmt(t1)}")
print(f"np.einsum('ij,ji->', B, B) {fmt(t2)}   ({t1/t2:.0f} 倍速い)")
print(f"(B * B.T).sum()          {fmt(t3)}")
print("一致:", np.allclose(np.trace(B @ B), np.einsum('ij,ji->', B, B)))
np.trace(B @ B)            319.252 ms
np.einsum('ij,ji->', B, B)   8.645 ms   (37 倍速い)
(B * B.T).sum()             12.360 ms
一致: True

37 倍の差がつきました。これは定数倍ではなく、$O(n^3)$ を $O(n^2)$ に落としたことによる、計算量そのものの改善です。einsum の真価はここにあります。「出力に現れない添字は消える」という記法のおかげで、中間の $n \times n$ 行列を一度も作らずに済むのです。

比較対象の (B * B.T).sum() も $O(n^2)$ で 12.4 ms と同じオーダーですが、einsum に 1.4 倍負けています。B * B.T という 72 MB の一時配列を作るぶんだけ余計にメモリを往復するからです。しかも B.T はストライドが逆なので、要素ごとの積の段階でキャッシュミスが起きます。einsum はこの中間を作らずに直接和を積み上げるので、その差が出ます。

einsum が BLAS に負ける場面

一方で、einsum を万能薬のように使うと痛い目を見ます。二次形式 $y_i = \bm{x}_i^\top \bm{A} \bm{x}_i$($N$ 個のサンプルそれぞれについてマハラノビス距離的な量を計算する)を 3 通りで書いてみます。

import numpy as np

rng = np.random.default_rng(0)
N, D = 20000, 50
X = rng.random((N, D)); A = rng.random((D, D))

def f_einsum():            # 3項の縮約を einsum に丸投げ
    return np.einsum('ij,jk,ik->i', X, A, X)

def f_einsum_opt():        # 最適な縮約順序を探索させる
    return np.einsum('ij,jk,ik->i', X, A, X, optimize=True)

def f_blas():              # 行列積に落としてから要素ごとに畳む
    return ((X @ A) * X).sum(axis=1)

print("一致:", np.allclose(f_einsum(), f_blas()))
for name, f in [("einsum(既定)", f_einsum),
                ("einsum(optimize=True)", f_einsum_opt),
                ("((X@A)*X).sum(1)", f_blas)]:
    print(f"{name:<24} {fmt(bench(f, target=0.5, repeat=5))}")
print("縮約の順序:", np.einsum_path('ij,jk,ik->i', X, A, X, optimize=True)[0])
一致: True
einsum(既定)             52.229 ms
einsum(optimize=True)      16.442 ms
((X@A)*X).sum(1)           34.154 ms
縮約の順序: ['einsum_path', (0, 1), (0, 1)]

トレースを3通りで計算した対数横棒グラフと、二次形式を3通りで計算した棒グラフ。トレースはeinsumが319.3ミリ秒に対し8.6ミリ秒で37倍速く、二次形式は既定のeinsumが52.2ミリ秒で最も遅くoptimize指定が16.4ミリ秒で最速であることを示す

3 つとも同じ結果ですが、速度は 3 倍以上ばらけました。ポイントは 2 つです。

第一に、np.einsum は既定(optimize=False)では BLAS を呼びません。添字の入れ子ループを NumPy 自身の汎用 C ループで回します。行列積のように演算律速な処理では、何十年もチューニングされた BLAS の dgemm に敵うはずがありません。実際、素朴に書いた既定の einsum が 52.2 ms で 3 つの中で最も遅くX @ A として BLAS に落とす書き方(34.2 ms)に 1.5 倍負けています。

第二に、optimize=True はこのケースでは劇的に効きました。既定の 52.2 ms に対して 16.4 ms、3.2 倍です。このオプションは「どの順序で縮約するか」を探索し、見つけた経路を tensordot(つまり BLAS)と縮約の組み合わせに落とします。出力された経路 [(0, 1), (0, 1)] は「まず XA を縮約し(=行列積)、その結果と X を縮約する」という意味です。手書きの ((X @ A) * X).sum(1) と同じ順序ですが、最後の 'ik,ik->i' を一時配列なしの 1 パスで畳むぶんだけ速くなっています(手書き版は $N \times D$ の一時配列を作ってから sum します)。

ただし optimize=True が常に得とは限りません。経路探索そのものに時間がかかるので、小さいテンソルや 2 項の縮約では、探索のオーバーヘッドが利得を食い潰します。必ず既定と比較して測る、というのが唯一の正しい運用です。

使い分けの指針

以上をまとめると、einsum の判断基準はこうなります。

  • 使うべき — 出力に現れない添字が多く、素直に書くと巨大な中間を作る場合(トレース、部分的な縮約、バッチ化された内積)。$O(n^3) \to O(n^2)$ のようなオーダーの改善が狙えるとき
  • 既定のまま行列積を書かない — 実質が行列積なら、既定の einsum は BLAS のマルチスレッドと SIMD を捨てることになる。@ に落とすか、optimize=True を付けて BLAS へ回してもらう
  • optimize=True は測ってから — 3 テンソル以上の縮約では大きく効くことがある(ここでは 3.2 倍)が、小さい縮約では探索コストが勝つ。必ず既定と比較する

einsum は「中間を作らない」ための道具であり、それ自体が「速い演算カーネル」なのではない、と覚えておくとよいでしょう。次の節では、同じ「中間を作らない」問題を、より実務的なペア距離の計算で扱います。

実測4: ペア距離 — ブロードキャストを行列積に落とす

$M$ 個の $D$ 次元点の集合 $\{\bm{p}_i\}$ と $\{\bm{q}_j\}$ の全ペア距離行列を作る、という計算を考えます。$k$ 近傍法、クラスタリング、カーネル法など、あらゆる場面で出てくる基本操作です。

素直に書くと、ブロードキャストを使ってこうなります。

D_ij = np.sqrt(((P[:, None, :] - Q[None, :, :])**2).sum(-1))

P[:, None, :] は形状 $(M, 1, D)$、Q[None, :, :] は $(1, M, D)$ なので、引き算の結果は $(M, M, D)$ です。ここに問題があります。中間配列のサイズが $M^2 D$ に膨れるのです。float64 なら $8 M^2 D$ バイト。$M = 2000$, $D = 50$ なら、

$$ 8 \times 2000^2 \times 50 = 1.6 \times 10^9\ \text{バイト} = 1.6\ \text{GB} $$

最終的に欲しいのは $M^2 = 4 \times 10^6$ 要素(32 MB)だけなのに、その 50 倍のメモリを経由しています。メモリ律速の世界で 50 倍のバイトを動かせば、当然 50 倍近く遅くなる方向に働きます。

これを回避するのが、ノルムの展開です。

$$ \|\bm{p}_i – \bm{q}_j\|^2 = (\bm{p}_i – \bm{q}_j)^\top (\bm{p}_i – \bm{q}_j) $$

右辺を展開します。内積は分配法則が使えるので、

$$ = \bm{p}_i^\top \bm{p}_i – \bm{p}_i^\top \bm{q}_j – \bm{q}_j^\top \bm{p}_i + \bm{q}_j^\top \bm{q}_j $$

内積は可換($\bm{p}^\top \bm{q} = \bm{q}^\top \bm{p}$)なので中央の 2 項をまとめると、

$$ \|\bm{p}_i – \bm{q}_j\|^2 = \|\bm{p}_i\|^2 – 2\,\bm{p}_i^\top \bm{q}_j + \|\bm{q}_j\|^2 $$

これを行列でまとめて書きます。$\bm{P} \in \mathbb{R}^{M \times D}$, $\bm{Q} \in \mathbb{R}^{M \times D}$ とし、行ごとのノルム二乗を並べたベクトルを $\bm{u}, \bm{v} \in \mathbb{R}^M$($u_i = \|\bm{p}_i\|^2$)とすると、距離二乗行列 $\bm{S}$ は

$$ \bm{S} = \bm{u}\bm{1}^\top – 2\,\bm{P}\bm{Q}^\top + \bm{1}\bm{v}^\top $$

となります。ここで $\bm{1}$ は全要素 1 のベクトルで、$\bm{u}\bm{1}^\top$ と $\bm{1}\bm{v}^\top$ はブロードキャストで表現できます。中心にあるのは $\bm{P}\bm{Q}^\top$ という単なる行列積です。中間配列は $(M, M)$ の 1 枚だけで、$D$ の次元は BLAS の内部で畳まれてしまいます。

数値誤差への注意

この変形にはひとつ落とし穴があります。$\bm{p}_i \approx \bm{q}_j$ のとき、$\|\bm{p}_i\|^2$ と $2\bm{p}_i^\top\bm{q}_j$ と $\|\bm{q}_j\|^2$ はどれも同じくらい大きな値で、その差がほぼゼロになります。これは桁落ち(catastrophic cancellation) の典型で、大きな数どうしの引き算で有効数字が失われます。結果として、本来 0 になるべき距離が $-10^{-12}$ のようなわずかな負の値になり、np.sqrtnan を返すことがあります。

対策として np.maximum(S, 0) でクリップします。scikit-learn の euclidean_distances も同じ手法と同じ対策を使っています。距離が極端に小さい領域での精度が要求される用途では、この展開は使わないほうが安全です。

import numpy as np, tracemalloc

rng = np.random.default_rng(0)

def dist_broadcast(P, Q):
    """素直なブロードキャスト。中間配列は (M, M, D)"""
    return np.sqrt(((P[:, None, :] - Q[None, :, :])**2).sum(-1))

def dist_gram(P, Q):
    """ノルム展開+行列積。中間配列は (M, M) だけ"""
    u = (P**2).sum(1)[:, None]      # (M, 1)
    v = (Q**2).sum(1)[None, :]      # (1, M)
    S = u - 2.0 * (P @ Q.T) + v     # 距離の二乗
    return np.sqrt(np.maximum(S, 0.0))   # 桁落ちで負になった分をクリップ

for M, D in [(1000, 3), (2000, 3), (2000, 50)]:
    P = rng.random((M, D)); Q = rng.random((M, D))
    ok = np.allclose(dist_broadcast(P, Q), dist_gram(P, Q), atol=1e-6)
    t1 = bench(lambda: dist_broadcast(P, Q), target=0.5, repeat=5)
    t2 = bench(lambda: dist_gram(P, Q), target=0.5, repeat=5)
    print(f"M={M:>5} D={D:>3}: broadcast {fmt(t1)} / gram {fmt(t2)}"
          f"  比 {t1/t2:5.2f}  中間配列 {M*M*D*8/2**20:7.1f} MiB  一致={ok}")
M= 1000 D=  3: broadcast   12.910 ms / gram   15.050 ms   比  0.86  中間配列    22.9 MiB  一致=True
M= 2000 D=  3: broadcast   53.485 ms / gram   19.798 ms   比  2.70  中間配列    91.6 MiB  一致=True
M= 2000 D= 50: broadcast  191.740 ms / gram   24.157 ms   比  7.94  中間配列  1525.9 MiB  一致=True

ブロードキャスト版とノルム展開版のペア距離計算時間を3つの条件で比べた棒グラフと、中間配列サイズに対する速度比をプロットした図。中間配列が22.9MiBでは0.86倍とブロードキャストが速く、1525.9MiBでは7.94倍まで差が開くことを示す

読み取りどころは、比が中間配列のサイズとともに素直に伸びることです。$M = 1000$, $D = 3$ の小さいケースでは、中間配列が 23 MiB 程度なのでブロードキャスト版のほうがむしろ 1.16 倍速い(比 0.86)。BLAS 呼び出しの固定費や、np.maximumsqrt の追加パスのぶんで gram 版が損をするからです。$M$ を倍にして中間配列が 91.6 MiB になると比は 2.70 に、$D = 50$ で 1.5 GiB になると 7.94 倍まで開きます。「行列積に落とすのが常に正解」ではなく、$M$ と $D$ が大きいほど有利になるという条件付きの結論です。

ピークメモリも測っておきます。$M = 2000, D = 3$ で tracemalloc を仕掛けると、ブロードキャスト版が 122.1 MiB、gram 版が 91.6 MiB でした。$D = 3$ ではまだ 1.33 倍の差ですが、$D = 50$ ならブロードキャスト版だけが 1.5 GiB を要求し、gram 版は $D$ に依存しないのでほぼ同じ 91.6 MiB のままです。$D$ を大きくしていくと、速度の問題である以前に動くか動かないかの問題になります。

この節で見た「中間配列のサイズを式から見積もる」という習慣は、NumPy コードを読むときの最重要スキルのひとつです。[:, None, :] のようなブロードキャスト構文を見たら、反射的に掛け算した結果の形状と、それが何バイトかを暗算する。それだけで、多くの性能問題は書く前に避けられます。

さて、ここまで「どれだけのバイトを動かすか」を見てきました。最後に、同じバイト数でも並び方が違うと速度が変わるという、より深い層の話に進みます。

実測5: メモリレイアウト — C連続とF連続、そしてキャッシュライン

strides を読む

NumPy の配列は「1 本の連続したバイト列」+「そこをどう歩くかの規則」でできています。この規則が strides です。

import numpy as np

A = np.arange(12, dtype=np.float64).reshape(3, 4)
print("C 連続:", A.strides, A.flags['C_CONTIGUOUS'], A.flags['F_CONTIGUOUS'])
F = np.asfortranarray(A)
print("F 連続:", F.strides, F.flags['C_CONTIGUOUS'], F.flags['F_CONTIGUOUS'])
print("転置ビュー:", A.T.strides, "データは共有:", A.T.base is A)
C 連続: (32, 8) True False
F 連続: (8, 24) False True
転置ビュー: (8, 32) データは共有: True

C 連続(行優先)の A では、strides = (32, 8) です。「行を 1 つ進むには 32 バイト、列を 1 つ進むには 8 バイト」という意味で、同じ行の隣り合う要素がメモリ上で隣り合っていることを表します。F 連続(列優先、Fortran 由来)では (8, 24) で、逆に同じ列の要素が隣接します。

転置 A.Tstrides を入れ替えるだけの ビュー で、データのコピーは起きません(A.T.base is ATrue)。これは NumPy の優れた設計ですが、同時に落とし穴でもあります。転置はタダだが、転置された配列を走査するのはタダではないからです。

行走査と列走査で 4.9 倍

同じ行列の全要素を足すのに、行方向に走査するか列方向に走査するかで何が変わるか測ります。

import numpy as np

rng = np.random.default_rng(0)
K = 3000
C = np.asarray(rng.random((K, K)), order='C')   # 行優先, 72 MB
F = np.asfortranarray(C)                        # 列優先, 同じ中身

def by_row(A):
    s = np.zeros(K)
    for i in range(K):
        s += A[i, :]     # 行を1本ずつ足す
    return s

def by_col(A):
    s = np.zeros(K)
    for j in range(K):
        s += A[:, j]     # 列を1本ずつ足す
    return s

print("C連続  行走査:", fmt(bench(lambda: by_row(C), target=0.3, repeat=5)))
print("C連続  列走査:", fmt(bench(lambda: by_col(C), target=0.3, repeat=5)))
print("F連続  行走査:", fmt(bench(lambda: by_row(F), target=0.3, repeat=5)))
print("F連続  列走査:", fmt(bench(lambda: by_col(F), target=0.3, repeat=5)))
C連続  行走査:    2.09 ms
C連続  列走査:   10.27 ms
F連続  行走査:   10.13 ms
F連続  列走査:    1.97 ms

C連続とF連続の行列に対して行方向と列方向に走査したときの実行時間を比べた棒グラフと、キャッシュライン128バイトのうち何バイトを実際に使うかを図解した説明。並びと走査が合っていれば使用率100%、ずれていると6%しか使わないことを示す

きれいな鏡像になりました。同じ 900 万要素を足すのに、走査の向きが並びと合っているかどうかで 4.9 倍違うのです。C 連続の配列を列方向に走査すると、1 要素読むごとに 24000 バイト($3000 \times 8$)飛びます。CPU はメモリを 1 バイト単位ではなくキャッシュライン単位で読むので、右の図のように128 バイトのラインを引っ張ってきて 8 バイトだけ使い、残り 120 バイトを捨てるという動作を繰り返すことになります。使用率は 6% です。

ここで大事なのは、「C 連続が速い / F 連続が遅い」ではないことです。合っているかどうかだけが問題であり、F 連続の配列を列方向に走査すれば 1.97 ms と、C 連続の行走査(2.09 ms)とほぼ同じ、むしろわずかに速い性能が出ます。scikit-learn や scipy の一部関数が F 連続を要求するのは、内部のアルゴリズムが列方向にアクセスするからです。

ストライドを振ってキャッシュラインを見る

キャッシュラインの効果は、ストライドを段階的に変えると数値として見えます。

import numpy as np

rng = np.random.default_rng(0)
a = rng.random(16_000_000)
print("ステップ  要素数     時間      1要素あたり")
for step in [1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 128]:
    v = a[::step]                      # ストライド step のビュー
    t = bench(lambda: v.copy(), target=0.3, repeat=15)
    print(f"  {step:>3}   {v.size:>8}  {fmt(t)}  {t/v.size*1e9:6.2f} ns/要素")
ステップ  要素数     時間      1要素あたり
    1   16000000     2.240 ms    0.14 ns/要素
    2    8000000     2.351 ms    0.29 ns/要素
    4    4000000     1.919 ms    0.48 ns/要素
    8    2000000     1.673 ms    0.84 ns/要素
   12    1333334     1.298 ms    0.97 ns/要素
   16    1000000     6.247 ms    6.25 ns/要素
   24     666667     1.705 ms    2.56 ns/要素
   32     500000     3.037 ms    6.07 ns/要素
   48     333334     1.876 ms    5.63 ns/要素
   64     250000     1.336 ms    5.34 ns/要素
  128     125000     0.500 ms    4.00 ns/要素

ストライドを1から128まで振ったときの1要素あたり時間をプロットした図。16要素すなわち128バイトを境に0.14ナノ秒から6.25ナノ秒へ45倍跳ね上がり、その後は4から6ナノ秒で頭打ちになることを示す

1 要素あたりのコストは、ステップ 12 までは 0.14 ns から 0.97 ns へゆるやかに増えるだけですが、ステップ 16 で 6.25 ns へ一気に 45 倍跳ね上がり、そこから先は 4〜6 ns 台で頭打ちになります。この「段差の位置」に意味があります。この環境のキャッシュラインは 128 バイト、double は 8 バイトなので、$128/8 = 16$、つまり 16 要素飛ばすとラインが完全に別物になるのです。ステップ 16 以降は「1 要素読むために毎回 1 本のラインを丸ごと転送する」状態で飽和し、それ以上悪くなりようがありません。x86 の多くの環境ではキャッシュラインが 64 バイトなので、段差はステップ 8 で現れるはずです。この測定は、自分のマシンのキャッシュラインを実験的に求める方法にもなっています。

なお、ステップ 24 だけ 2.56 ns と不自然に低いのが目につきます。これは 2 の冪でないストライドの効果です。ステップが 2 の冪だとアドレスの下位ビットが揃うため、アクセスがキャッシュの一部のセットに集中してセット競合ミスが起きます。24(192 バイト)のような半端なストライドではアドレスがセット全体に散るので、競合が緩和されるのです。細かい話ですが、「性能が 2 の冪のサイズで落ち込む」現象(行列サイズを 1024 から 1025 に変えると速くなる、など)の正体はこれです。

転置コピーはなぜ非線形に悪化するのか

もう一段深い現象を見ます。転置された配列を連続化(実体化)するコストは、行列サイズに対して非線形に増えます。

import numpy as np

rng = np.random.default_rng(0)
for K in [1000, 2000, 4000]:
    C = np.asarray(rng.random((K, K)), order='C')
    t_copy = bench(lambda: C.copy(), target=0.3, repeat=5)
    t_trans = bench(lambda: np.ascontiguousarray(C.T), target=0.3, repeat=5)
    print(f"K={K:>5} ({C.nbytes/2**20:6.1f} MiB):"
          f" 素直なコピー {fmt(t_copy)} / 転置の連続化 {fmt(t_trans)}"
          f"  比 {t_trans/t_copy:5.1f}")
K= 1000 (   7.6 MiB): 素直なコピー    0.130 ms / 転置の連続化    0.367 ms  比   2.8
K= 2000 (  30.5 MiB): 素直なコピー    0.578 ms / 転置の連続化    3.109 ms  比   5.4
K= 4000 ( 122.1 MiB): 素直なコピー    2.346 ms / 転置の連続化   26.922 ms  比  11.5

バイト数は 4 倍ずつ増えているのに、転置コピーの時間は $0.367 \to 3.11 \to 26.9$ ms と 8.5 倍、8.7 倍のペースで伸びています。素直なコピーが $0.130 \to 0.578 \to 2.346$ とほぼバイト数どおり(4.4 倍、4.1 倍)なのと対照的で、比は 2.8 → 5.4 → 11.5 と一貫して開いていきます。

原因は 2 段階あります。$K = 1000$ 程度なら、1 列($K \times 8 = 8$ KB 相当のライン群)がまだキャッシュに収まるので、列を読み進めても前に引いたラインが生き残ります。$K$ が大きくなると 1 列を読み終わる前にキャッシュが一巡してしまい、すべてのアクセスがミスになります。さらに $K = 4000$ では TLB ミスが加わります。列方向のアクセスは $4000 \times 8 = 32000$ バイトごとに次の要素を読むので、この環境のページサイズ 16 KB では 1 要素ごとに別ページです。TLB のエントリ数を超えると、アドレス変換のたびにページテーブルを歩くことになります。

タイル化で緩和する

この問題は、小さなブロックに区切って転置することで緩和できます。ブロックが L1 に収まるサイズなら、ブロック内では読みも書きも局所的になります。

import numpy as np

rng = np.random.default_rng(0)
K = 4000
C = np.asarray(rng.random((K, K)), order='C')

def blocked_transpose(A, bs=64):
    """bs x bs のタイル単位で転置する"""
    n = A.shape[0]
    out = np.empty((n, n))
    for i in range(0, n, bs):
        for j in range(0, n, bs):
            out[j:j+bs, i:i+bs] = A[i:i+bs, j:j+bs].T
    return out

t_naive = bench(lambda: np.ascontiguousarray(C.T), target=1.0, repeat=3)
print("素直な連続化:", fmt(t_naive))
for bs in [16, 32, 64, 128, 256, 512]:
    t = bench(lambda: blocked_transpose(C, bs), target=1.0, repeat=3)
    print(f"{bs:>4}タイル分割: {fmt(t)}  ({t_naive/t:.2f} 倍速い)")
print("一致:", np.array_equal(np.ascontiguousarray(C.T), blocked_transpose(C, 64)))
素直な連続化:   30.240 ms
  16タイル分割:   44.750 ms  (0.68 倍速い)
  32タイル分割:   28.960 ms  (1.04 倍速い)
  64タイル分割:   24.600 ms  (1.23 倍速い)
 128タイル分割:   23.360 ms  (1.29 倍速い)
 256タイル分割:   19.750 ms  (1.53 倍速い)
 512タイル分割:   13.610 ms  (2.22 倍速い)
一致: True

転置の連続化コストが行列サイズに対して非線形に悪化する様子を両対数で示した図と、タイルの一辺を16から512まで振ったときの実行時間をプロットした図。512角タイルで素直な連続化の2.22倍速くなることを示す

Python レベルで二重ループを回しているにもかかわらず、タイルの一辺が 32 以上なら素直な連続化に勝ち、512 角タイルでは 2.22 倍速くなりました。これは「Python ループは常に悪」という素朴な信念への反例です。

一方でタイルが小さすぎると負けます。一辺 16 では $250^2 = 62500$ 回の反復が必要になり、1 回あたりの仕事は 256 要素しかありません。1 反復あたりの実測は $44.75\ \text{ms} / 62500 \approx 0.72$ マイクロ秒で、そのうち相当部分がスライス 2 回と代入 1 回にかかる Python 側のオーバーヘッドです。これが支配的になり、0.68 倍と逆に遅くなります。一辺 512 なら反復は $8^2 = 64$ 回で、1 回あたり 262144 要素。ループのコストは完全に無視できます。

ループ回数が少なく、1 回あたりの仕事が大きいなら、Python ループを足しても構わない。ベクトル化の目的は「Python ループを消すこと」ではなく「1 回の C 呼び出しあたりの仕事を増やすこと」だと理解していれば、この結果は自然に受け止められるはずです。

行列積とオーダーの関係

最後に、行列積 A @ B でオーダーがどう効くかです。直感的には、$\bm{A}$ は行方向、$\bm{B}$ は列方向にアクセスするので、$\bm{A}$ は C 連続、$\bm{B}$ は F 連続が理想に思えます。実際、素朴な三重ループ実装ならそのとおりです。

しかし NumPy の @ は BLAS の dgemm を呼びます。dgemm は内部で両方のオペランドをキャッシュに収まるパネルに再パック(コピー)してから計算するので、入力のオーダーにあまり敏感ではありません。再パックのコストは $O(n^2)$、計算は $O(n^3)$ なので、$n$ が大きいほど再パックの相対コストは無視できます。演算強度が高い($I = n/12$)処理は、そもそもメモリレイアウトの影響を受けにくいのです。

つまり、メモリレイアウトを気にすべきなのは BLAS の外側です。具体的には、

  • リダクション(sum, mean, maxaxis 指定)
  • 要素ごとの演算で、片方だけが転置ビューの場合(A * B.T など)
  • スライスの取り出しと代入
  • reshape — 連続でない配列に対する reshape は暗黙にコピーを発生させる

こうした場所で「並びと走査の向きが合っているか」を確認し、合っていなければ np.ascontiguousarray / np.asfortranarray一度だけ変換して以後使い回す、というのが定石になります。変換自体は上で見たとおり高くつくので、ループの中で毎回やってはいけません。

ここまでで、4 つのコストすべてを実測で確認できました。最後に、これらを実際の高速化作業の手順に落とし込みます。

高速化の意思決定フロー

測定結果を踏まえると、NumPy コードの高速化は次の順序で進めるのが合理的です。上から順に効果が大きく、労力が小さい順に並んでいます。

ステップ1: Python ループが残っていないか確認する 要素数が数千以上のループが 1 つでも残っていれば、そこが 100 倍以上のボトルネックです。他の最適化を考える前にここを潰します。ただし、ループ 1 回あたりの仕事が十分大きい(数千要素以上を処理する)なら、残しても構いません。

ステップ2: 演算強度を見積もる $I = W/Q$ を暗算します。$I \lesssim 1$ ならメモリ律速なので、動かすバイト数を減らす方向に手を打ちます。$I \gtrsim 10$ なら演算律速なので、BLAS に落として並列化する方向です。

ステップ3: 中間配列のサイズを式から数える [:, None, :] のようなブロードキャスト構文を見たら、結果の形状を掛け算して 8 を掛けます。最終出力より 1 桁以上大きい中間が出るなら、数学的に書き換えられないか検討します(ノルム展開、einsum の縮約、np.dot への帰着)。

ステップ4: dtype を見直す float64 を float32 にすれば、動かすバイト数が半分になります。メモリ律速の処理なら、それだけで 2 倍近く速くなります。実測すると、$n = 2 \times 10^7$ の a * b は float64 で 5.02 ms、float32 で 2.10 ms と 2.4 倍でした(バイト数が半分になるうえ、SIMD で一度に扱える要素数も倍になるため、2 倍をやや上回ります)。精度要件が許すかどうかだけが判断基準です。

ステップ5: メモリレイアウトを揃える リダクションの axis や要素ごと演算で、走査の向きと strides が食い違っていないか確認します。食い違っていれば、ループの外で 1 回だけ変換します。

ステップ6: 最後に out=+= ここまでやってまだ足りないなら、インプレース化を検討します。効果は 1〜2 割程度で、しかも小さい配列では逆効果になりうる(実測 0.93 倍)と分かったうえで、測定しながら入れます。

この順序が重要です。多くの人がステップ 6 から始めて、「NumPy っぽく書いたのに速くならない」と悩みます。効果の大きさは上から順に、桁で違います

それでも足りないとき

NumPy のベクトル化には原理的な限界があります。「配列全体に対する演算」に分解できない処理は、どうしてもループが残ります。典型は、前の要素の結果が次の計算に必要な逐次依存(IIR フィルタ、累積的な状態更新、条件分岐の多い処理)です。

そうした場合の選択肢を、性質の違いとともに挙げておきます。

  • Numba@njit を付けるだけで Python の関数を LLVM でネイティブコンパイルします。ループがそのまま高速化されるので、逐次依存のある処理に最適です。parallel=True でループの自動並列化もできます
  • Cython — 型注釈を付けて C に変換します。既存の C ライブラリとつなぐ場合に強みがあります
  • BLAS のスレッド数OMP_NUM_THREADS などの環境変数で制御します。他のプロセスと競合している環境では、スレッド数を減らしたほうが速いことがあります
  • np.memmap — メモリに載らない巨大配列を、ファイルとして扱いながら NumPy の演算を適用します
  • Dask / CuPy — 前者は分割してコアを跨ぐ並列化、後者は GPU への移行です

ただし、これらに手を出す前に、本記事で見た測定を必ずやってください。メモリ律速の処理を GPU に移しても、GPU の帯域とホストへの転送コストで期待ほど速くならないのはよくある失敗です。演算強度を見積もっていれば、その失敗は事前に避けられます。

よくある誤解 Q&A

Q. 「NumPy を使えば速い」は正しい? 配列が十分大きい場合に限って正しい、というのが答えです。実測では $n = 100$ で 16 倍、$n = 10^4$ で 255 倍、$n = 10^6$ で 161 倍でした。逆に、100 要素の配列に対して NumPy 関数を 10 万回呼ぶようなコードは、1 回あたり約 165 ナノ秒の固定費が支配的になり、純 Python と大差なくなります。配列を大きくして呼び出し回数を減らすのが正しい方向です。

Q. なぜ速度比は $n = 10^4$ をピークに下がるの? NumPy 側が先にメモリ帯域へ張り付くからです。NumPy の a * b は $n = 3\times10^4$ で 0.147 ns/要素(163 GB/s)が最速で、大きくすると 0.239 ns/要素(100 GB/s)で頭打ちになります。一方 Python ループは 1 要素 40 ns も使っているので帯域には遠く届かず、$n$ を増やしても大して悪化しません。結果として比が縮みます。

Q. einsum は常に速い? いいえ。既定では BLAS を呼ばないので、実質が行列積なら @ のほうが速いです。実測では二次形式で既定の einsum 52.2 ms に対し ((X@A)*X).sum(1) が 34.2 ms でした。einsum が輝くのは、素直に書くと巨大な中間ができる縮約(トレースで 37 倍)です。

Q. optimize=True を付ければ安心? 測ってから決めてください。実測の二次形式では 52.2 ms → 16.4 ms と 3.2 倍速くなりました(縮約が BLAS に落ち、最後の畳み込みも一時配列なしで済むため)。ただし経路探索そのものにコストがかかるので、小さい縮約では逆効果になります。必ず既定と比較するのが唯一の運用方法です。

Q. out= は必ず速くなる? いいえ。np.empty による確保はほぼ無料($10^7$ 要素で 0.51 us)で、冷たいページへの初回書き込みも温かいバッファ比 4% しか高くありません。実測では out=x+= 相当)が 1〜2 割速い程度、別配列への out=o はほぼ同じ、$n = 10^5$ では += が 0.93 倍と逆に遅くなりました。

Q. 途中結果を変数に入れると何が起きる? NumPy の一時配列消去(temporary elision)が無効になり、ピークメモリが増えます。実測では 122.1 MiB が 213.7 MiB(1.75 倍)になりました。差はちょうど中間配列 1 枚ぶんです。巨大配列を扱うときは、式を分けずに書くか del で明示的に手放してください。

Q. C 連続と F 連続、どちらを使うべき? どちらが速いということはなく、アクセスの向きと合っているかだけが問題です。実測では C 連続の行走査 2.09 ms に対し列走査 10.27 ms(4.9 倍)で、F 連続では完全に逆転しました(行走査 10.13 ms、列走査 1.97 ms)。使う関数がどちらを前提にしているかを調べ、ループの外で 1 回だけ変換するのが定石です。

Q. Python のループは絶対に消すべき? いいえ。ループ 1 回あたりの仕事が十分大きければ、インタプリタのオーバーヘッドは相対的に消えます。実測のタイル転置では、Python の二重ループを追加したのにキャッシュ局所性の改善で 2.22 倍速くなりました(一辺 512 のタイル、反復は $8^2 = 64$ 回)。ただし一辺 16(反復 62500 回)まで細かくすると 0.68 倍と逆転します。ループを消すこと自体が目的ではなく、1 回の C 呼び出しあたりの仕事を増やすことが目的です。

まとめ

本記事では、NumPy のベクトル化がなぜ速いのかを、実測で分解しながら解説しました。

  • 速度差は 4 つのコストの積み重ね — インタプリタのディスパッチ(要素あたり 5.6 ns の固定費)、ボックス化された float(24 バイト/個で 32 MB 対 8 MB)、メモリ帯域、SIMD。空のループだけで NumPy の 28 倍遅いという事実が、ボトルネックの所在を示しています
  • ルーフラインで方針が決まる — 演算強度 $I = W/Q$ を暗算し、$I \lesssim 1$ ならバイト数を減らす、$I \gtrsim 10$ なら BLAS に落とす。要素ごとの積は $I = 1/24$ で 4.2 GFLOPS しか出ず完全にメモリ律速、$n \times n$ の行列積は $I = n/12$ で 169 GFLOPS まで届きます
  • 速度比は $n$ に依存し、山になる — $n = 100$ で 16 倍、$n = 10^4$ で 255 倍と最大、$n = 10^6$ では 161 倍に戻ります。小さい側は NumPy の固定費(約 165 ns/呼び出し)、大きい側は NumPy がメモリ帯域(100 GB/s)に張り付くためです
  • out= より「中間を作らない」 — 確保はほぼ無料($10^7$ 要素で 0.51 us)なのでインプレース化の効果は 1〜2 割。それより $M^2 D$ の中間配列を $M^2$ に落とす書き換えのほうが桁で効きます($M=2000, D=50$ で 7.94 倍)
  • einsum は縮約の道具であって速いカーネルではない — $\mathrm{tr}(\bm{A}\bm{B})$ を $O(n^3)$ から $O(n^2)$ に落として 37 倍。一方、既定の einsum で行列積を書くと BLAS 版に 1.5 倍負けます(optimize=True で BLAS に回せば 3.2 倍取り返せます)
  • レイアウトは「合っているか」がすべて — C 連続の行走査 2.09 ms 対列走査 10.27 ms(4.9 倍)。ストライドを振ると、キャッシュラインの幅(この環境では 128 バイト= 16 要素)でコストが 45 倍に跳ねる様子が数値で見えます

ここで身につけた「バイト数と演算数を数える」習慣は、NumPy に限らず、PyTorch のテンソル演算でも、C++ の数値計算でも、そのまま使えます。フレームワークが変わっても、メモリ階層と演算器の関係は変わらないからです。

次のステップとして、以下の記事も参考にしてください。