SIMD と AVX を EC2 上の Python + NumPy で試してみた
はじめに
AWS ParallelCluster や AWS Batch で HPC ワークロードを実行する際、SIMD や AVX といった用語を目にすることがあります。これらの技術が実際にどの程度パフォーマンスに影響するのか、Python と NumPy を使って EC2 上で試してみました。
本記事では用語の簡単な解説と、AVX2 の有効・無効による行列計算を試してみて、約 3.9 倍の早くなったことを紹介します。
SIMD と AVX について
主に科学技術計算や機械学習などの数値計算が多いワークロードで登場するキーワードです。
SIMD (Single Instruction, Multiple Data)
SIMD とは 1 つの命令で複数のデータを同時に処理する並列処理技術です。
例えば通常の処理は命令を 1 つずつ逐次実行します。
命令1: データA + データB = 結果A
命令2: データC + データD = 結果B
命令3: データE + データF = 結果C
命令4: データG + データH = 結果D
SIMD 命令では、1 つの命令で同時処理ができます。この例では理論上 4 倍速くなります。
命令1: [データA, データC, データE, データG] + [データB, データD, データF, データH]
= [結果A, 結果B, 結果C, 結果D]
AVX (Advanced Vector Extensions)
AVX は CPU の演算性能を向上させるための拡張命令セットです。SSE の後継として登場しました。
AVX の世代
AVX にも世代があります。新しくなるにつれて浮動小数点の演算を同時処理できる数が増えています。SSE も複数バージョンがありましたが長くなるので割愛しました。
命令セット | 登場年 | レジスタ幅 | 同時処理数(64 ビットの場合) |
---|---|---|---|
SSE | 1999 年 | 128 ビット | 2 個 |
AVX | 2011 年 | 256 ビット | 4 個 |
AVX2 | 2013 年 | 256 ビット | 4 個 |
AVX-512 | 2016 年 | 512 ビット | 8 個 |
試してみた
SIMD と AVX についてはわかりました。実際に試してみて学んでみます。
Python で NumPy を利用して AVX あり/なしによる行列計算の違いを試してみます。
検証環境
EC2 インスタンス
項目 | 内容 |
---|---|
インスタンスタイプ | m7i.xlarge |
CPU | Intel(R) Xeon(R) Platinum 8488C |
OS | Amazon Linux 2023 |
Python | 3.13.3 |
NumPy | 2.3.4 |
OpenBLAS | 0.3.30 (Haswell ビルド) |
CPU の AVX サポート状況
AVX、AVX2、AVX-512 すべてをサポートしていました。第 7 世代のインスタンスタイプで新しいですからね。
$ cat /proc/cpuinfo | grep flags | head -1 | grep -o 'avx[^ ]*' | sort -u
avx
avx2
avx512_bf16
avx512_bitalg
avx512_fp16
avx512_vbmi2
avx512_vnni
avx512_vpopcntdq
avx512bw
avx512cd
avx512dq
avx512f
avx512ifma
avx512vbmi
avx512vl
avx_vnni
Python の実行環境
Python 3.13 と NumPy をインストールします。
# Python 3.13 のインストール
$ sudo dnf install -y python3.13-devel python3.13-pip
# バージョン確認
$ python3.13 --version
Python 3.13.3
# 仮想環境の作成と有効化
$ python3.13 -m venv venv
$ source venv/bin/activate
# pip と NumPy のインストール
$ pip install --upgrade pip
$ pip install numpy
NumPy/OpenBLAS 設定の確認
NumPy は行列計算を高速化するために BLAS(Basic Linear Algebra Subprograms)ライブラリを使用します。どの BLAS ライブラリが使用されているか、AVX のサポート状況を確認します。
NumPy バージョン
$ python3 -c "import numpy as np; print(f'NumPy: {np.__version__}')"
NumPy: 2.3.4
NumPy 詳細確認
python3 -c "import numpy as np; np.show_config()"
- BLAS: scipy-openblas(OpenBLAS 0.3.30)を使用
- OpenBLAS 設定: Haswell(Intel のマイクロアーキテクチャ名)→ AVX2 まで対応
CPU は AVX-512 をサポートしていました。ですが、実際の行列演算を行う OpenBLAS は Haswell ビルド(AVX2 まで)のため、AVX-512 は使用されないようでした。AVX-512 でビルドすることも検討しましたが、効果は限定的な可能性が示唆されていたため、そのまま利用することにしました。
一般的なアプリケーションにおける有用性の低さについて
wiki より
実行結果の折りたたみ
/home/ec2-user/venv/lib64/python3.13/site-packages/numpy/__config__.py:155: UserWarning: Install `pyyaml` for better output
warnings.warn("Install `pyyaml` for better output", stacklevel=1)
{
"Compilers": {
"c": {
"name": "gcc",
"linker": "ld.bfd",
"version": "14.2.1",
"commands": "cc"
},
"cython": {
"name": "cython",
"linker": "cython",
"version": "3.1.4",
"commands": "cython"
},
"c++": {
"name": "gcc",
"linker": "ld.bfd",
"version": "14.2.1",
"commands": "c++"
}
},
"Machine Information": {
"host": {
"cpu": "x86_64",
"family": "x86_64",
"endian": "little",
"system": "linux"
},
"build": {
"cpu": "x86_64",
"family": "x86_64",
"endian": "little",
"system": "linux"
}
},
"Build Dependencies": {
"blas": {
"name": "scipy-openblas",
"found": true,
"version": "0.3.30",
"detection method": "pkgconfig",
"include directory": "/opt/_internal/cpython-3.13.8/lib/python3.13/site-packages/scipy_openblas64/include",
"lib directory": "/opt/_internal/cpython-3.13.8/lib/python3.13/site-packages/scipy_openblas64/lib",
"openblas configuration": "OpenBLAS 0.3.30 USE64BITINT DYNAMIC_ARCH NO_AFFINITY Haswell MAX_THREADS=64",
"pc file directory": "/project/.openblas"
},
"lapack": {
"name": "scipy-openblas",
"found": true,
"version": "0.3.30",
"detection method": "pkgconfig",
"include directory": "/opt/_internal/cpython-3.13.8/lib/python3.13/site-packages/scipy_openblas64/include",
"lib directory": "/opt/_internal/cpython-3.13.8/lib/python3.13/site-packages/scipy_openblas64/lib",
"openblas configuration": "OpenBLAS 0.3.30 USE64BITINT DYNAMIC_ARCH NO_AFFINITY Haswell MAX_THREADS=64",
"pc file directory": "/project/.openblas"
}
},
"Python Information": {
"path": "/tmp/build-env-xnu9j23j/bin/python",
"version": "3.13"
},
"SIMD Extensions": {
"baseline": [
"SSE",
"SSE2",
"SSE3"
],
"found": [
"SSSE3",
"SSE41",
"POPCNT",
"SSE42",
"AVX",
"F16C",
"FMA3",
"AVX2",
"AVX512F",
"AVX512CD",
"AVX512_SKX",
"AVX512_CLX",
"AVX512_CNL",
"AVX512_ICL",
"AVX512_SPR"
],
"not found": [
"AVX512_KNL",
"AVX512_KNM"
]
}
}
実行プログラム
Claude Code を利用してシンプルな行列計算スクリプトを作成しました。
#!/usr/bin/env python3
"""
行列計算だけのシンプルなベンチマーク
2つの行列を掛け算するだけのプログラムです。
例:
行列A (100x100) × 行列B (100x100) = 行列C (100x100)
"""
import numpy as np
import time
def main():
print("=" * 60)
print("行列計算ベンチマーク")
print("=" * 60)
# 設定
matrix_size = 1000 # 1000x1000の行列
iterations = 5 # 5回繰り返し
print(f"\n【設定】")
print(f" 行列サイズ: {matrix_size} x {matrix_size}")
print(f" 繰り返し回数: {iterations} 回")
print(f" 計算内容: A × B = C (行列の掛け算)")
# ステップ1: ランダムな行列を2つ作成
print(f"\n【ステップ1】ランダムな行列を2つ作成中...")
A = np.random.rand(matrix_size, matrix_size)
B = np.random.rand(matrix_size, matrix_size)
print(f" 行列A: {matrix_size}x{matrix_size} の行列")
print(f" 行列B: {matrix_size}x{matrix_size} の行列")
print(f"\n 行列Aの一部:")
print(f" {A[0, :5]}")
print(f" 行列Bの一部:")
print(f" {B[0, :5]}")
# ステップ2: 行列の掛け算を繰り返す
print(f"\n【ステップ2】行列の掛け算を実行中...")
times = []
for i in range(iterations):
# 時間計測開始
start = time.perf_counter()
# 行列の掛け算(ここでAVX2が使われる)
C = A @ B
# 時間計測終了
end = time.perf_counter()
elapsed = end - start
times.append(elapsed)
print(f" {i+1}回目: {elapsed*1000:.2f} ミリ秒")
# ステップ3: 結果を計算
print(f"\n【ステップ3】結果")
avg_time = sum(times) / len(times)
min_time = min(times)
max_time = max(times)
print(f" 平均時間: {avg_time*1000:.2f} ミリ秒")
print(f" 最速: {min_time*1000:.2f} ミリ秒")
print(f" 最遅: {max_time*1000:.2f} ミリ秒")
# 計算量の情報
total_operations = matrix_size * matrix_size * matrix_size * 2 # 乗算と加算
gflops = (total_operations / avg_time) / 1e9
print(f"\n 計算量: {total_operations/1e9:.2f} 億回の演算")
print(f" 性能: {gflops:.2f} GFLOPS (1秒間に何億回計算できるか)")
# 結果の確認
print(f"\n【確認】計算結果のサンプル")
print(f" 行列C (A×Bの結果) の一部:")
print(f" {C[0, :5]}")
print("\n" + "=" * 60)
print("完了!")
print("=" * 60)
if __name__ == "__main__":
main()
実行方法
AVX2 有効の場合
デフォルト状態で AVX2 が利用されます。何度か検証していたので環境変数をクリアしてから実行しています。
$ unset OPENBLAS_CORETYPE # 環境変数をクリア(AVX2有効)
$ python3 benchmark_matrix_simple.py
AVX 無効の場合
拡張命令セットの AVX を使わず、基本的な命令セットのみで計算してもらいます。
$ export OPENBLAS_CORETYPE=KATMAI # 基本的な命令セット指定
$ python3 benchmark_matrix_simple.py
実行結果
AVX2 有効
============================================================
行列計算ベンチマーク(シンプル版)
============================================================
【設定】
行列サイズ: 1000 x 1000
繰り返し回数: 5 回
計算内容: A × B = C (行列の掛け算)
【ステップ1】ランダムな行列を2つ作成中...
行列A: 1000x1000 の行列
行列B: 1000x1000 の行列
行列Aの一部:
[0.97522309 0.25590702 0.82768389 0.10400997 0.41126751]
行列Bの一部:
[0.22073047 0.55123485 0.76811404 0.53638863 0.86763605]
【ステップ2】行列の掛け算を実行中...
1回目: 13.71 ミリ秒
2回目: 13.13 ミリ秒
3回目: 12.83 ミリ秒
4回目: 12.96 ミリ秒
5回目: 12.05 ミリ秒
【ステップ3】結果
平均時間: 12.94 ミリ秒
最速: 12.05 ミリ秒
最遅: 13.71 ミリ秒
計算量: 2.00 億回の演算
性能: 154.60 GFLOPS (1秒間に何億回計算できるか)
【確認】計算結果のサンプル
行列C (A×Bの結果) の一部:
[254.37613647 257.99070529 252.56361471 264.80468734 261.48124516]
============================================================
完了!
============================================================
AVX 無効(基本的な命令セット指定)
============================================================
行列計算ベンチマーク(シンプル版)
============================================================
【設定】
行列サイズ: 1000 x 1000
繰り返し回数: 5 回
計算内容: A × B = C (行列の掛け算)
【ステップ1】ランダムな行列を2つ作成中...
行列A: 1000x1000 の行列
行列B: 1000x1000 の行列
行列Aの一部:
[0.36582787 0.18759122 0.72934862 0.72053441 0.29003918]
行列Bの一部:
[0.4128825 0.01193105 0.15556012 0.02282569 0.86116398]
【ステップ2】行列の掛け算を実行中...
1回目: 52.15 ミリ秒
2回目: 50.06 ミリ秒
3回目: 50.05 ミリ秒
4回目: 50.24 ミリ秒
5回目: 49.41 ミリ秒
【ステップ3】結果
平均時間: 50.38 ミリ秒
最速: 49.41 ミリ秒
最遅: 52.15 ミリ秒
計算量: 2.00 億回の演算
性能: 39.70 GFLOPS (1秒間に何億回計算できるか)
【確認】計算結果のサンプル
行列C (A×Bの結果) の一部:
[250.3054255 254.9739997 238.24762122 249.58004213 255.37609278]
============================================================
完了!
============================================================
実行結果の比較
比較すると AVX2 により約 3.9 倍の高速化を確認できました。基本的な命令セットと AVX2 では演算性能に大きな違いがでました。
項目 | AVX2 有効 | AVX 無効 | 比較 |
---|---|---|---|
平均時間 | 12.94 ms | 50.38 ms | 約 3.9 倍 |
性能(GFLOPS) | 154.60 | 39.70 | 約 3.9 倍 |
まとめ
-
行列計算で AVX2 による約 3.9 倍の高速化を確認しました
- AVX2 有効: 12.94 ms(154.60 GFLOPS)
- AVX 無効: 50.38 ms(39.70 GFLOPS)
-
今回使用した NumPy の OpenBLAS は Haswell ビルド(AVX2 まで)対応でした
- CPU は AVX-512 対応だが、ライブラリが対応していない
- 実際の性能はライブラリに依存する
おわりに
AWS ParallelCluster や、AWS Batch で複数インスタンスタイプを混在させて起動確率を高める方法をよくとっています。AVX でも AVX-512 指定があった際に使用するライブラリはもちろん、混在可能なインスタンスタイプの CPU を調べる必要があることを理解できました。手を動かしてみると学びが多くていいですね。