前回までのあらすじ

ようやくCPU(のむちゃくちゃ遅い実装)に対して20倍速い行列積の計算ができるようになりました。

これはまだまだ全然遅いので、今回はそれをどんどん改善していきます。

なお、記事内で使用しているD言語のOpenCLラッパーライブラリは下記です。(自作)

現状の速度について

1000 * 2000の行列と2000 * 3000の行列の積を計算した場合、FLOPSはたぶん下記の通りになります。

結果の行列は1000 * 3000で、要素数は1e3 * 3e3 = 3e6 になる。
この3e6の要素それぞれについて、乗算が2000回・加算が2000回発生する。つまり 3e6 * (2e3 + 2e3) = 12e9
計算回数 1e9 = 1G 回、12e9 = 12G 回
よって、1000 * 2000の行列と2000 * 3000の行列の積を仮に1秒で計算できた場合、12GFLOPS

現状ではK80を使って2秒ほど掛かっていたので、12GFLOPSの半分の6GFLOPSが計算速度になります。

さて、私がGoogleにお金を払って借りているK80の性能はどの程度でしょうか。

Up to 8.73 teraflops single-precision performance with NVIDIA GPU Boost

出展
https://www.nvidia.com/en-us/data-center/tesla-k80/

借りているのはGPU1コア(1/2ボード)なので、半分と考えて4.36TFLOPS程度でしょうか。

つまり、いま出せている速度は、カタログスペックの1/725です。

CPUと比較して喜んでいる場合ではなかった……。
少しでもこのカタログスペックにたどり着くよう、最適化バトルが始まります。

ローカルメモリを使う

OpenCLで扱えるメモリには、下記のような階層があります。

ホストのメモリー
====== CommandQueueなどの壁(PCIeなどのバス) ======
GPUのグローバルメモリー
====== GPU内部バスの壁 ======
GPUのローカルメモリー。ワークグループ間で共有可能。
====== 共有メモリ同期の壁 ======
プライベートメモリー。いわゆるレジスタ？

下に行くほど計算ユニットに近く、高速で、小容量です。
また、速度の壁は上に行くほど厚い(はず)です。
つまり、上の方にある壁ほど、そこを超えてデータを読み書きするのに時間がかかります。

さて、現状のコードはどうなっているでしょうか。

__kernel void product(
        __global const float *lhs, // グローバルメモリー
        __global const float *rhs, // グローバルメモリー
        __global float *result,    // グローバルメモリー
        uint rows,
        uint cols,
        uint resultCols) {
    // ワークアイテムのグローバルIDを取得する。
    const size_t groupI = get_global_id(0);
    const size_t groupRows = get_global_size(0);
    const size_t groupJ = get_global_id(1);
    const size_t groupCols = get_global_size(1);

    for(size_t i = groupI; i < rows; i += groupRows) {
        for(size_t j = groupJ; j < resultCols; j += groupCols) {
            float value = 0.0f;
            for(size_t k = 0; k < cols; ++k) {
                // グローバルメモリーから読み込み * 2
                value += lhs[i * rows + k] * rhs[k * resultCols + j];
            }
            // グローバルメモリーに書き込み
            result[i * resultCols + j] = value;
        }
    }
}

ループのもっとも内側で、2番目に遅いグローバルメモリーから読み込みまくっています。
しかも、全ワークアイテムでその読み込みが発生しています。
さらに、よく考えると、グローバルメモリーの同じ箇所を何度も計算用にコピーしています。

これを改善する方法を考えます。

計算に使用するデータを、より速いローカルメモリーに置く。
ローカルメモリーでデータを共有することで、ワークグループの範囲で読み込み回数を削減する。
ローカルメモリーにデータを置くことで、同じ値を何度もグローバルメモリーから読み込まないようにする。

ワークアイテムの現状

現在、ワークアイテムは行列積の結果の各要素について下記のように割り当てられています。
(A〜Iが個々のワークアイテムを示します)

	0	1	2	3	4	...
0	A	B	C	A	B	C
1	D	E	F	D	E	F
2	G	H	I	G	H	I
3	A	B	C	A	B	C
4	D	E	F	D	E	F
...	G	H	I	G	H	I

このA〜Iの塊の1つ1つが、一度に計算されるようになっています。

	0	1	2
0	A	B	C
1	D	E	F
2	G	H	I

(これが行列の端まで敷き詰められている)

このひと塊を計算するのに必要なデータをローカルメモリーに共有し、グローバルメモリーへのアクセスを省略することで、高速化を図ります。

ローカルメモリーの確保

先ほどからローカルメモリーとばかり言っていますが、それはどうやって確保するのでしょうか。

私の知る限り2つ方法があります。

まず、GPUコードの中でローカル変数に__local修飾子をつける方法です。

void f(void) {
    // これ
    __local float localMemory[1024];
}

これでも良いのですが、いかんせん固定長になってしまいます。

もう1つ、動的に確保する方法があります。それは、ホスト側コードでclSetKernelArgを使用する方法です。

clSetKernelArg(kernel, 0, sizeof(float) * 1024, null);

GPU側コードは下記のようにします。

__kernel void kernel(__local float *localMemory) {
    // localMemory使ってガンガン計算
}

こうすると、localMemory引数がローカルメモリーということになり、指定したサイズ確保されている状態になります。

……正直これはAPIとして微妙な感じなので、自作ラッパーライブラリではallocateLocalMemoryという名前で関数を追加しました。

cl.allocateLocalMemory(kernel, 0, 1024 * float.sizeof);

ワークアイテムの同期

ローカルメモリーを使う場合、同一ワークグループ内のワークアイテムで同期を行う必要があります。
同期を行わないと、使いたい値を別ワークアイテムに上書きされてしまったり、消されたりと言った問題が発生します。

そういった処理の競合を防ぐために、barrier(CLK_LOCAL_MEM_FENCE)を使用します。
これを使うと、ワークグループ内のワークアイテムがbarrierのある行で一旦停止し、タイミングを合わせた上で一斉に再開することになります。

つまり、ローカルメモリーへの読み書きのタイミングを揃えられます。

気をつけないといけない点として、ワークグループの全ワークアイテムがbarrierに到達しないと正しく同期してくれません。
ifやforで一部のワークアイテムが実行しないbarrierがある場合、バグの元になります。

出典 https://www.khronos.org/registry/OpenCL/sdk/1.2/docs/man/xhtml/barrier.html

If barrier is inside a conditional statement, then all work-items must enter the conditional if any work-item enters the conditional statement and executes the barrier.

実装

上記を踏まえてようやく実装です。

ワークグループ単位で処理を行うようループを整理する

まず、全ワークアイテムがbarrierを踏む必要があるので、ループの回り方を変えます。
範囲チェックを外側のforではなく内側で行うようにします。
CPU脳で考えると超無駄に思えますが、ローカルメモリーを使うためには仕方ありません。

const size_t groupI = get_global_id(0);
const size_t groupRows = get_global_size(0);
const size_t groupJ = get_global_id(1);
const size_t groupCols = get_global_size(1);

// ここでは行列の端を超えていてもループを終了しないようにする。
for(size_t i = 0; i < rows; i += groupRows) {
    for(size_t j = 0; j < resultCols; j += groupCols) {
        // 行列範囲内かどうか
        const bool inRange = (i + groupI) < rows && (j + groupJ) < resultCols;
        float value = 0.0f;
        
        for(size_t k = 0; k < cols; ++k) {
            
            // TODO: ここでワークグループが使用する値をローカルメモリーにコピー

            // 範囲内の場合のみ計算
            if(inRange) {
                // TODO: ローカルメモリーを使用する
                value += lhs[(i + groupI) * rows + k] * rhs[k * resultCols + (j + groupJ)];
            }
        }

        // 範囲内の場合のみ値を設定
        if(inRange) {
            result[(i + groupI) * resultCols + (j + groupJ)] = value;
        }
    }   
}

さて、まだグローバルなワークアイテムIDを基準にした処理しか行なっていません。
ローカルメモリーはワークグループ内でしか使用できないので、ワークグループ内のIDであるローカルIDを基準にループするよう修正します。

// ワークグループ内部でのIDを取得
const size_t localI = get_local_id(0);
const size_t localRows = get_local_size(0);
const size_t localJ = get_local_id(1);
const size_t localCols = get_local_size(1);

        // ワークグループのサイズ(列数)ごとの処理に修正
        for(size_t k = 0; k < cols; k += localCols) {
            
            // TODO: ここでワークグループが使用する値をローカルメモリーにコピー

            // 範囲内の場合のみ計算
            if(inRange) {
                // ワークグループが処理している範囲で計算を行う。
                // 現在のワークアイテムが行列の端から出ている場合は何もしない。
                for(size_t lk = 0; lk < localCols && (k + lk) < cols; ++lk) {
                    // TODO: ローカルメモリーを使用する
                    value += lhs[(i + groupI) * rows + (k + lk)] * rhs[(k + lk) * resultCols + (j + groupJ)];
                }
            }
        }

この一番内側のlkのループでグローバルメモリーアクセスがローカルメモリーアクセスになることで高速化されます。

ローカルメモリーへのコピー

いよいよローカルメモリーへグローバルメモリーから値をコピーしますが、ここからが面倒臭いです。

今回の処理は、

直前の計算(ローカルメモリーアクセス)が終わるまで待機
ワークグループ内部の全ワークアイテムで、ローカルメモリーにコピー
ローカルメモリーのコピーが終わるまで待機
ローカルメモリーを使用した計算を開始

という順序で行います。というわけで、とりあえずbarrierを貼ります。

// TODO: ここでワークグループが使用する値をローカルメモリーにコピー
barrier(CLK_LOCAL_MEM_FENCE); // 直前の計算完了まで待つ
barrier(CLK_LOCAL_MEM_FENCE); // コピー完了まで待つ

次に、グローバルメモリーからローカルメモリーへのコピーを行います。
このコピーをまた単一ワークアイテムでシコシコ実行すると遅いです。
そこで、ワークアイテムを一斉に動かして並列処理でコピーが終わるようにします。
コードとしては、localI・localJを元にグローバルメモリーlhs・rhsからローカルメモリーへ、1要素分のコピーを行うことになります。
コード上ではただの1要素のコピーに見えますが、ワークアイテムで並列実行されるので必要な要素分コピーが行われることになります。

段階を追って実装します。まず、左辺の行列の行の要素をコピーしてみます。

barrier(CLK_LOCAL_MEM_FENCE); // 直前の計算完了まで待つ

// ワークアイテムが行列からはみ出していなければコピー実行
if((i + groupI) < rows && (k + localJ) < cols) {
    localRow[localI * localCols + localJ] = lhs[(i + groupI) * cols + (k + localJ)];
}

barrier(CLK_LOCAL_MEM_FENCE); // コピー完了まで待つ

localRowはどこから出てきたのでしょうか？これはカーネル関数の引数に追加し、先述のallocateLocalMemoryで確保します。

__kernel void product(
        __global const float *lhs,
        __global const float *rhs,
        __global float *result,
        uint rows,
        uint cols,
        uint resultCols,
        __local float *localRow) // これ

localRowは、ワークグループで一度に処理する要素の分、つまりワークグループ内のワークアイテムの数だけ必要になります。

// ローカルメモリーを確保
cl.allocateLocalMemory(kernel, 6, 32 * 32 * float.sizeof);

右辺の行列の列の要素もコピーします。

barrier(CLK_LOCAL_MEM_FENCE); // 直前の計算完了まで待つ

// ワークアイテムが行列からはみ出していなければコピー実行
if((i + groupI) < rows && (k + localJ) < cols) {
    localRow[localI * localCols + localJ] = lhs[(i + groupI) * cols + (k + localJ)];
}
if((j + groupJ) < resultCols && (k + localI) < cols) {
    localCol[localI * localCols + localJ] = rhs[(k + localI) * resultCols + (j + groupJ)];
}

barrier(CLK_LOCAL_MEM_FENCE); // コピー完了まで待つ

localColについても、ローカルメモリーを先ほどと同様に確保します。(省略)

（以下、細かい言い訳。なぜrhsからのコピーでlocalIを使用するのかについて）
さて、ここで実はワークグループサイズに暗黙の仮定を置いてしまっています。現状では上のコードはワークグループサイズの行数・列数が同じでないと動きません。
なぜなら、localColへのコピーで、コピー元のインデックス生成にlocalI(行番号)を使用してしまっているからです……。
ここは本来はlocalJ(列番号)が対応するべきですが、localJが増えるとgroupJも増える関係にあるため、うまくいきません。(対角線の要素のみコピーするような動作になってしまいます)
localIを使えば、groupJが同じで行番号が違うワークアイテムが動くことになるので、コピー元とコピー先が1:1対応してうまくコピーできます……。

ローカルメモリーで計算する

色々苦労しましたがようやくローカルメモリーで計算できます。
localRow・localColにワークグループが参照する両辺の要素はコピーされているので、これを参照して計算するようにします。

        // ワークグループのサイズ(列数)ごとの処理に修正
        for(size_t k = 0; k < cols; k += localCols) {
            
            barrier(CLK_LOCAL_MEM_FENCE); // 直前の計算完了まで待つ

            // ワークアイテムが行列からはみ出していなければコピー実行
            if((i + groupI) < rows && (k + localJ) < cols) {
                localRow[localI * localCols + localJ] = lhs[(i + groupI) * cols + (k + localJ)];
            }
            if((j + groupJ) < resultCols && (k + localI) < cols) {
                localCol[localI * localCols + localJ] = rhs[(k + localI) * resultCols + (j + groupJ)];
            }

            barrier(CLK_LOCAL_MEM_FENCE); // コピー完了まで待つ

            // 範囲内の場合のみ計算
            if(inRange) {
                // ワークグループが処理している範囲で計算を行う。
                // 現在のワークアイテムが行列の端から出ている場合は何もしない。
                for(size_t lk = 0; lk < localCols && (k + lk) < cols; ++lk) {
                    // ローカルメモリー参照
                    value += localRow[localI * localCols + lk] * localCol[lk * localCols + localJ];
                }
            }
        }

本当にこれで早くなったのか?

計算結果

ローカルメモリー使用前

cpu: 48344 msecs, gpu: 2100 msecs

ローカルメモリー使用後

cpu: 42206 msecs, gpu: 791 msecs

なんということでしょう。ほぼ2.5倍に！

あんな遅そうなコードなのに……。

FLOPSで言えば、15GFLOPSまで到達しました。とは言え、理論性能の1/290です……。
明らかにまだまだ遅そうなので、さらなる最適化を行います。