コグノスケ

SIMDを使ったお手軽最適化 - その1

目次: ベンチマーク

お手軽最適化のメモです。行列の掛け算を題材にします。

最初にお断りしておくとGEMMのような汎用処理の場合は、自分で最適化せずにOpenBLASを使ってください（素人が最適化しても勝てません）。しかしOpenBLASのような限界まで最適化されたコードは誰でも簡単には書ける、とは言えません。

スカラー処理だと遅いけれど、お手軽に最適化（数倍程度）がしたいときの参考になれば幸いです。

GEMMってなんですか？

GEMMはGeneral Matrix Multiplyの略で、高校数学辺りでやった（はず）の行列の掛け算のことです。floatの場合はSGEMMと呼ばれ、doubleの場合はDGEMMと呼ばれます。最適化の題材はどちらでも良いんですけど、今回はSGEMMを使います。

忘れている方のために2行3列の行列Aと3行2列Bの掛け算A x B = Cだとこんな感じです。

行列の掛け算の例

Aの列数とBの行数は一致していなければなりません。行数と列数の関係を表すと、行列A（M行K列）x行列B（K行N列）= 行列C（M行N列）となります。

行列の掛け算の行数と列数

Cの1要素を計算するには、Aの1行とBの1列が必要です。式、および、視覚的に示すと下記のようになります。

行列の掛け算（式）

行列の掛け算（Aの行とBの列の関係）

説明はこれくらいにしてコードを見ましょう。

基本コース - 素朴に演算

SGEMMを素直にコードにするとこんな感じです。行方向にデータを格納（Row-major orderといいます）しているので、N列の行列Cのi行j列（以降Ci,jと書く）にアクセスする際はc[i * N + j] とします。

void sgemm_naive(const float *a, const float *b, float *c, int mm, int nn, int kk)
{
	for (int i = 0; i < mm; i++) {
		for (int j = 0; j < nn; j++) {
			c[i * nn + j] = 0.0f;
			for (int k = 0; k < kk; k++) {
				c[i * nn + j] += a[i * kk + k] * b[k * nn + j];
			}
		}
	}
}

行列のサイズを適当に設定（M = 1519, N = 1517, K = 1523）して、実行時間を測ります。CPUはRyzen 7 5700Xです。

$ gcc -Wall -g -O2 -static sgemm.c

$ ./a.out
matrix size: M:1519, N:1517, K:1523
time: 2.277758

実行時間は実行するたびに変わりますが、大体2.27秒くらいでしょうか。OpenBLASのシングルスレッド（環境変数OPENBLAS_NUM_THREADS=1にするとシングルスレッド動作になります）で計算した時間を見ると、

	c_ex = malloc(m * n * sizeof(float) * 2);

	// C = alpha AB + beta C
	float alpha = 1.0f, beta = 0.0f;
	int lda = k, ldb = n, ldc = n;

	printf("----- use CBLAS\n");

	gettimeofday(&st, NULL);
	cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
		m, n, k, alpha, a, lda, b, ldb, beta, c_ex, ldc);
	gettimeofday(&ed, NULL);
	timersub(&ed, &st, &ela);

	printf("verify: %d.%06d\n", (int)ela.tv_sec, (int)ela.tv_usec);

$ export OPENBLAS_NUM_THREADS=1

$ gcc -Wall -g -O2 -static -L path/to/openblas/ sgemm.c -lopenblas

$ ./a.out
matrix size: M:1519, N:1517, K:1523
----- use CBLAS
verify: 0.052149

わずか0.05秒、実に43倍という驚異のスピードです。すごいですね……。

行列の掛け算の説明でほぼ終わってしまいました。お手軽最適化はまた次に。