式テンプレートを使った高速な線形代数ライブラリを作った話

研究で作っているD言語用のライブラリの一部として式テンプレート（Expression Template）を使った高速な線形代数ライブラリを作りましたのでその紹介をします．

単にBLASのライブラリを呼ぶためのラッパーですが，小さな行列でも大きな行列でも効率よく計算できるライブラリを作成できたと思います．

ちなみに，このライブラリはまだ非公開です．．．

設計指針その１：パフォーマンスについて

行や列のサイズが100を超えるような大きな行列では単にBLASを呼び出すことが高速化への近道です． というか，このサイズの行列では素人が頑張ってチューニングしてもOpenBLASやMKLに勝てません． そして，大きな行列においてはラッパーライブラリによるオーバヘッドが相対的に小さくなるため，BLASさえ使っておけばライブラリ作成者がパフォーマンスについて考える必要はありません． たとえば，Pythonのnumpyのようなライブラリで，大きな行列やベクトルの演算が非常に高速に実行できる理由はこれです．

一方で，行や列のサイズが8以下のような小さなサイズの行列は話が変わってきます． このサイズでは行列演算の計算時間とラッパーのコストが相対的に近くなってきます． さらに，今日のコンパイラは賢いので非常に効率のよいコードを吐き出してくれます． 結果的に，小さな行列ではBLASを自作のラッパーライブラリ経由で呼び出すよりも，愚直にforループを書いたほうが性能が出てしまいます．

つまり，小さな行列においてどれだけforループに近い性能が出せるかが勝負です．

設計指針その２：コンパイル時の式変形

たとえば，一般の$N \times N$行列と$N$要素のベクトルの積は$\mathcal{O}(N^2)$の計算量となります． もし，この行列がDFT行列であって$N$が2のべき乗であればこれは$\mathcal{O}(N \log N)$まで削減できます．

このようにベクトルにDFT行列を作用させた場合には行列とベクトルの積を計算するのではなく，FFTを実行するようにライブラリ側で操作できないでしょうか？

もしくは，ゼロ行列をベクトルや行列に作用させると，結果は必ずゼロ行列になります． 単位行列とある行列の積はかならず元の行列になります． これらのような型レベルの演算の削減をライブラリ側で行うことはできないでしょうか？

また，計算量削減のためのコンパイル時の式変形として，面白い例がもう一つあります． 二つの行列$A,B$とベクトル$v$の積$ABv$を考えたとき，$AB$を先に計算するのか$Bv$を先に計算するのかで計算量が変わります． 単にプログラムのソースコード上でA * B * vと与えられたときにA * (B * v)として最適化してくれるライブラリは作成できないでしょうか？

式テンプレートによる式木の構築

これらを実現するための手法として式テンプレートというプログラミング技法が有効です． 式テンプレートは，プログラムから静的な型情報を保存した式木を構築する技法です． 式テンプレートでは，式木を構築するときや，その式木を実行するときに，コンパイル時情報を用いて式木を操作することができます． また，コンパイラが十分に賢ければ複雑な行列の式をオーバーヘッドなしにforループと同等な処理に置き換えることも可能です． 式木の中に$aAB + bC$というようなBLASのGEMMを呼び出して積和を一度に計算できる部分を静的に見つけて，GEMMの呼び出しgemm(a, A, B, b, C)に変換することもできます．

ライブラリの例

それではまずは簡単な例として行列積をどのように書くのかを以下に例示します．

import dffdd.math.matrix;

void main()
{
    // 1000x1000の行列を3つ定義（0初期化）
    auto A = matrix!float(1000, 1000, 0);
    auto B = matrix!float(1000, 1000, 0);
    auto C = matrix!float(1000, 1000, 0)l

    // A * B + Cの結果を格納する行列（0初期化）
    auto C = matrix!float(1000, 1000, 0);

    // BLASのGEMMを実行して格納
    D[] = 2 * A * B + C * 3;
}

左辺のDは右辺に含まれていないことをライブラリ側に示すことでより最適化したコードが生成される可能性があり，特に行列のサイズが小さいときには有効です． 行列のサイズが4x4程度以上あれば.noalias無しの場合にも十分に性能が出るようになっていますが，.noaliasがないとメモリ領域を2倍消費します．

D.noalias = 2 * A * B + C * 3;

なお，以上のように呼び出した場合にはライブラリの内部では次のようにGEMMが呼び出されます． これは人間が書くBLASのGEMMを用いた行列積と同じです．

// D.noalias = 2 * A * B + C * 3;の疑似コード
D[] = C;                // Dの領域にCをコピー
gemm(2, A, B, 3, D);    // BLASのGEMMでD <- 2 * A * B + 3 * Dを実行

もちろん + C * 3の部分がなくても，ライブラリは次のようにGEMMを呼び出します． これも人間が書くBLASのGEMMを用いた行列積と同じです．

// D.noalias = 2 * A * B;の疑似コード
D[] = 0;                // Dをゼロ初期化
gemm(2, A, B, 0, D);    // BLASのGEMMでD <- 2 * A * B + 0 * Dを実行

式テンプレートと遅延評価

このライブラリでは式テンプレートによる遅延評価（Lazy Evaluation）を採用しているため，numpyの線形代数ライブラリなどのような一般的な正格評価（Eager Evaluation）とは異なった動きをすることがあります．

numpyのような正格評価のライブラリでは行列Aとベクトルvの積A @ vの結果はその計算結果のベクトルになります．

print(A @ v)    # 計算結果を表示

一方で，私のライブラリで同様な演算A * vを単に行うと，積の評価は行わずに式木を返します．

writeln(typeof(A * v).stringof);
// MatrixVectorMulGEMV!(float, MatrixedSlice!(float*, mir_slice_kind.contiguous), VectoredSlice!(float*, mir_slice_kind.contiguous), ConstAll!(float, 0.0F, 1LU))
// これはGEMV(a, A, v, b, Z)を表す式の型，ただしa, bはfloat型で，Zはゼロ行列

以下のような複雑な式でも，その式を表す式木になります．

writeln(typeof((A * A + 2 * A) * v + v).stringof);
// VectorAxpby!(float, MatrixVectorMulGEMV!(float, MatrixedSlice!(float*, mir_slice_kind.contiguous), MatrixVectorMulGEMV!(float, MatrixedSlice!(float*, mir_slice_kind.contiguous), VectoredSlice!(float*, mir_slice_kind.contiguous), ConstAll!(float, 0.0F, 1LU)), MatrixVectorMulGEMV!(float, MatrixedSlice!(float*, mir_slice_kind.contiguous), VectoredSlice!(float*, mir_slice_kind.contiguous), ConstAll!(float, 0.0F, 1LU))), VectoredSlice!(float*, mir_slice_kind.contiguous))

以下の場合には式木がexprに保存されるだけで演算は評価されません．

auto expr = A * v;

式木を評価して代入するためにはv[] = ...;という演算子を利用します． この演算子によって，一時オブジェクトの生成を最小限にすることができます．

// 式木の結果を評価してvに格納
v[] = A * v;

また，ある式木exprの評価結果を一時オブジェクトとして返す関数としてforceEvaluateがあります．

auto newVec = forceEvaluate(expr);

遅延評価その1（行列 x 行列 x ベクトルの例）

それでは具体的に式テンプレートによる遅延評価によってどのようなことが実現できるか例を示します．

たとえば，以下のようにnumpyで2つの行列A, Bとベクトルvの積を愚直に書いてしまうと， 演算子の結合規則から行列同士の積を計算した後にベクトルとの積を計算するので非効率です．

# 以下は (A @ B) @ vと解釈されて非効率
v = A @ B @ v

# 以下のように書くべき
v = A @ (B @ v)

本ライブラリでは式の構造をコンパイル時に解析してライブラリ側が自動的により良い計算順序に組み替えることができます．

// 式テンプレートによりA * (B * v)に式変形される
v[] = (A * B) * v;

一方で，このような式変形はあまり好みではなく，プログラマが指定した順番で実行すべしという人もいるかもしれません． その場合は単に一つずつ評価していけばいいだけです．

C[] = A * B;    // まず A * Bを評価する
v[] = C * v;    // その後に C * vを評価する

遅延評価その2（一時オブジェクトの削減）

たとえば以下のような三つの行列A, B, Cとスカラa, bからなる演算は，numpyのような正格評価のライブラリではA @ Bの部分でBLASのGEMMが呼び出され，その後に行列の和が計算されます．

C = a * A @ B + b * C
# （Pythonインタプリタが賢くない限り）これは次のように計算されるため
# 途中結果の行列T1〜T4が生成される
# T1 = a * A
# T2 = T1 @ B   (GEMM呼び出し)
# T3 = b * C
# T4 = T2 + T3

一方で私のライブラリでは次のようにBLASのGEMMを直接呼び出すため， 途中結果の行列が生成されません．

// コンパイラがどれだけ頭が悪くてもGEMM(a, A, B, b, C)の呼び出しへ変換
C[] = a * A * B + b * C;

// もちろん以下のような式もGEMM(a*a, A, B, a*b, C)へ変換
C[] = (b * C + A * (a * B)) * a;

このような一時オブジェクトの削減は行列のサイズが小さく，各演算の計算時間が短いときに非常に大きな意味を持ちます． 逆にいえば，巨大行列の積においてはほとんど改善されませんが．

静的式変形

次のように単位行列と一般行列の積の式木を構築したとき，その型は元の行列の型に一致します． つまり，単位行列と行列の積は，元の行列をそのまま返しています．

auto I = identity!float(3);
auto A = matrix!float(3, 3, 0);


auto B1 = I * A;
pragma(msg, typeof(B1).stringof);   // MatrixedSlice!(float*, mir_slice_kind.contiguous)

auto B2 = A * I;
pragma(msg, typeof(B2).stringof);   // MatrixedSlice!(float*, mir_slice_kind.contiguous)

原理的にはDFT行列を作業させたときにFFTを実行することも可能ですが，現状はDFT行列を実装していません．

パフォーマンス評価

ここでは本ライブラリの行列積$aAB + bC$のパフォーマンスを評価します． ソースコードは記事の末尾に示しています．

パフォーマンスを比較するのは以下の五つです．

• BLAS: D言語からCBLASの行列積cblas_sgemmの直接呼び出し
• For: D言語のforeachを用いて愚直に行列積
• dffdd(Lazy): 本ライブラリの行列積
• dffdd(Eager): 本ライブラリを用いて擬似的に正格評価をした行列積
• numpy: Python3のnumpyの行列積

行列のサイズを変えたときの結果は以下の表の通りで，単位はGFLOPSです．　

（グラフを作るのがめんどくさかったので許してください）

行数や列数の値（M=N=K）が5以下であればBLASよりもforループを書いたほうが速いです． そのため，私のライブラリでは$MNK$の積が100以下であればforループを利用するようにしています． ちなみに，この領域ですとnumpyや正格評価するライブラリでは全く性能が出ません． サイズが2x2行列のとき，正格評価版は遅延評価版の17分の1のパフォーマンスしか出ていません． さらに，遅延評価版の本ライブラリはnumpyの100倍も速く計算できます．

驚いてほしいことは，2x2程度の行列積であっても私のライブラリとForループがほとんど全く同じ性能を出していることです． これはコンパイラの最適化によって式テンプレートの式木を構築するオーバーヘッドが完全に除去され，Forループと全く同じようなコードが生成されていることを意味しています．

行列のサイズが256x256以上になると，numpyとBLASの性能はほとんど同等になり，Forループを圧倒します． これはnumpyも裏ではBLASを呼び出しているだけだからです． 私のライブラリもBLASを呼び出すため，BLASやnumpyと同等の性能が出ています．

numpyがBLASに追いつきだしてくるのは行列のサイズが128x128以上付近ですが，これよりも小さな行列ですとnumpyはほとんど性能が出ていません． 正格評価版も128x128のときには遅延評価版の63%しか速度がでません．

あと，よくわからないのが1024x1024で正格評価版のものが90GFLOPS程度しかでていないのに，numpyで130GFLOPSを超えているのはよくわかりません． numpyはなにかしら一時オブジェクトの再利用などをしているのでしょうか？

まとめ

サイズが小さな行列についてはforループ直書きと同程度に低コストで，サイズが大きな行列についてはBLASの呼び出しに変換してくれる式テンプレートを用いたライブラリを作りました． 式テンプレートを用いて一時オブジェクトを削減することで，オーバーヘッドを小さくすることに成功しました．

ちなみに，このライブラリはまだ非公開です．．．

ベンチマークソースコード