D言語で始めるOpenCL(4) プライベートメモリー編

これまでのあらすじ

• D言語で始めるOpenCL(1) 導入編
• D言語で始めるOpenCL(2) ローカルメモリー編
• D言語で始めるOpenCL(3) 下準備＆行列パディング編
• D言語で始めるOpenCL(4) プライベートメモリー編 (今回)

行列をパディングしてワークグループサイズに合わせ、条件分岐を無くすことで高速化に成功しました。
今回は、最後のメモリー階層であるプライベートメモリーを積極的に使って更なる高速化を目指します。

どうでもいいですが「プライベートメモリー」ってなんだか響きがエモいですね。

参考サイト

今回の高速化について、下記サイトがとても参考になりました。ソースコード等も公開されているので是非読みましょう。

Tutorial: OpenCL SGEMM tuning for Kepler

こちらもボリュームがあって凄かったです。

プログラムを高速化する話Ⅱ 〜GPGPU編〜

何はともあれ、GPUのリファレンスは読むべきです。

https://docs.nvidia.com/cuda/index.html
https://docs.nvidia.com/cuda/kepler-tuning-guide/index.html (今回使っているK80向け)

プライベートメモリー活用の検討

だいぶ高速化できた前回ですが、下記の部分をじっくり見ると無駄が見えてきます。

examples/product.cl

        for(size_t lk = 0; lk < localCols; ++lk) {
            value += localLhs[localI * localCols + lk] * localRhs[lk * localCols + localJ];
        }

1. 各ワークアイテムは、結果の行列の1要素を計算するために、上記のループを実行します。
2. 1要素を計算するためには、lhsの1行と、rhsの1列をローカルメモリーから読み込む必要があります。
3. 結果行列の同じ行の要素（ワークアイテム）は、すべて同じlhsの要素を読んでいます。
4. 同様に、結果行列の同じ列の要素（ワークアイテム）は、すべて同じrhsの要素を読んでいます。

ローカルメモリーからの読み出しも、GPUコアの計算スピードに比べたらまだまだ遅いです。
そこで、別々のワークアイテムでそれぞれ読み出しているローカルメモリーの内容をプライベートメモリーにキャッシュできれば、それだけ高速化ができそうです。

では、ローカルメモリーの読み込み結果をキャッシュして複数の要素で活用するためにはどうすれば良いか……。
参考サイトなどで取られている方法は、ワークアイテムの担当要素を1つではなく複数に広げるというものです。

複数の要素のブロックを1ワークアイテムで処理できれば、以下の効果が狙えます。

• ローカルメモリーの読み込み結果をプライベートメモリーにキャッシュし、メモリーアクセスを省略できる。
• 実際に動くワークアイテムの数が少なくなり、オーバーヘッドが減少する。

割と良いことずくめです。ただし以下の制約があります。

• 小さいサイズの行列の場合、動くワークアイテムの数が減るので、動かないコアが出てきて逆効果になるかもしれない。
• プライベートメモリーのサイズに制約があるので、1ワークアイテム当たりの担当要素はそれほど多くできない。
  • 多くできても、すごく遅くなる場合がある。（OpenCLの実装によっては、ローカルメモリー・グローバルメモリーで無理やり代替する場合がある？）

とはいえ、大きいサイズの行列の計算では速くできる可能性は高いので（そもそも小さいサイズであれば多少遅くても問題ない）、ワークアイテムの担当する要素の数を増やす方向で頑張ってみます。

ローカルメモリーの整理

プライベートメモリーを色々頑張る前に、ローカルメモリーへのキャッシュの部分を整理します。

一度に処理するkの数を指定できるようにしたい

examples/product.cl

        barrier(CLK_LOCAL_MEM_FENCE);
        localLhs[localI * localCols + localJ] = lhs[globalI * cols + (k + localJ)];
        localRhs[localI * localCols + localJ] = rhs[(k + localI) * resultCols + globalJ];
        barrier(CLK_LOCAL_MEM_FENCE);

ここでは、lhs・rhsのそれぞれについて、ワークグループのサイズと同じ要素数のローカルメモリーを用意してコピーを行っています。

これからワークアイテムの担当要素を増やすと、ワークグループのサイズと必要なローカルメモリーのサイズが一致しなくなります。1ワークアイテムに担当要素分のデータが必要になるので、必要なローカルメモリーも大きくなります。

また、そもそもワークグループサイズ分ずつ計算させる必要は必ずしもありませんでした。
極端な話、lhs・rhsからそれぞれ1要素ずつ読んで計算を進めても処理としては普通に終わります。

examples/product.cl

    for(size_t k = 0; k < cols; k += 1) { // lhs・rhsそれぞれ1要素ずつ読んで計算を進める
        barrier(CLK_LOCAL_MEM_FENCE);
        // 無駄なコピー
        localLhs[localI * localCols + localJ] = lhs[globalI * cols + (k + localJ)];
        localRhs[localI * localCols + localJ] = rhs[(k + localI) * resultCols + globalJ];
        barrier(CLK_LOCAL_MEM_FENCE);

        for(size_t lk = 0; lk < 1; ++lk) { // 1要素分の計算で終了
            value += localLhs[localI * localCols + lk] * localRhs[lk * localCols + localJ];
        }
    }

この計算処理を進める単位については、ワークグループサイズに依存しないパラメーターとしてちゃんと管理したいところです。

メモリーアクセスの高速化

基本的にNVIDIAのGPU依存の問題ですが、グローバルメモリーへのアクセスにはトランザクションの結合・ローカルメモリーへのアクセスにはバンクのコンフリクトという問題があります。

せっかくOpenCLだし、私もまだ理解できていない部分があるので、あまり詳細に突っ込まないで言うと、

• 連続したワークアイテムで
• 連続したメモリーに
• ワークアイテムの順序通りアクセスする

上記を全て満たすと速いということです。

逆に言えば、

• 不連続なワークアイテムを一度に使用するか
• 連続していない（ストライドのある）メモリーか
• ワークアイテムの順序通りでないアクセス

は遅いということになります。

なお、この連続したメモリーアクセスがCPUのコードと少し感覚が異なります。
コード上では連続アクセスに見えないが、同時に動く個別のワークアイテムが連続した番地にアクセスしているというのが最も速いです。

具体的にはこういうアクセスが速いです。

const size_t groupSize = get_local_size(0) * get_local_size(1);
const size_t localId = get_local_id(0) + get_local_size(0) * get_local_id(1);
for(size_t i = 0; i < N; ++i) {
    localMemory[(groupSize * i) + localId] = globalMemory[(groupSize * i) + localId];
}

えっ、これgroupSizeごとの不連続アクセスじゃん！と一瞬思うのですが、ワークアイテムを広げて考えるとちゃんと連続アクセスになっているのです。

ちなみに下記のようだとかなり遅くなるので注意が必要です。

// NG
const size_t groupSize = get_local_size(0) * get_local_size(1);
const size_t localId = get_local_id(1) + get_local_size(1) * get_local_id(0); // ここが違う
for(size_t i = 0; i < N; ++i) {
    localMemory[(groupSize * i) + localId] = globalMemory[(groupSize * i) + localId];
}

ローカルIDの次元の0と1を変えただけです。トータルのグループサイズは変わらないのでコピーは行えていますが、アクセスの順序が異なるため、速度が大きく変わります。

私はこれでだいぶ悩みました……。

ローカルメモリーの修正に合わせて、こういった点もなるべく改善してみます。

kの増加分を指定できるようにコードを修正する

とりあえずkの増加分を指定できるようにします。
kの増加分は行列サイズとは無関係のため、今回はパラメーターにせず定数にします。

product.d

    // 行列サイズ
    enum {
        ROWS = 1000,
        COLS = 2000,
        RESULT_COLS = 3000,
        BATCH_ROWS = 32,
        BATCH_COLS = 32,
        BATCH_SIZE_K = 32 // 追加。現状ではBATCH_ROWS・BATCH_COLSの倍数でないといけない……。
    }

// 中略

    // ソースコードをformatすることにして、その引数に与える。
    auto program = cl.createProgramFromSource(
        context, import("product.cl").format(BATCH_ROWS, BATCH_COLS, BATCH_SIZE_K));
    scope(exit) cl.releaseProgram(program);
    cl.buildProgram(program, deviceIds);

// 中略

    // カーネル引数
    cl.setKernelArg(kernel, 0, lhsBuffer);
    cl.setKernelArg(kernel, 1, rhsBuffer);
    cl.setKernelArg(kernel, 2, resultBuffer);
    cl.setKernelArg(kernel, 3, bufferRows);
    cl.setKernelArg(kernel, 4, bufferCols);
    cl.setKernelArg(kernel, 5, bufferResultCols);
    // ローカルメモリーのサイズはソースコードに与えた定数から求まるので、カーネル引数からは削除

product.cl

__kernel void product(
        __global const float *lhs,
        __global const float *rhs,
        __global float *result,
        uint rows,
        uint cols,
        uint resultCols) {
    // 追加
    enum {
        BATCH_ROWS = %d,
        BATCH_COLS = %d,
        BATCH_SIZE_K = %d
    };

    // ローカルメモリーをハードコードするようにした。
    // こちらの方がループを簡単に書けるし、unrollもされやすいもよう。
    __local float localLhs[BATCH_ROWS][BATCH_SIZE_K];
    __local float localRhs[BATCH_SIZE_K][BATCH_COLS];

    const size_t localJ = get_local_id(0);
    const size_t localI = get_local_id(1);
    const size_t globalJ = get_global_id(0);
    const size_t globalI = get_global_id(1);

    const size_t groupJ = get_group_id(0) * BATCH_COLS;
    const size_t groupI = get_group_id(1) * BATCH_ROWS;

    // 追加。ローカルメモリーへのコピー時に使用する値。
    const size_t localSize = get_local_size(0) * get_local_size(1);
    const size_t localId = get_local_id(0) + get_local_size(0) * get_local_id(1);
    const size_t localCopyCountLhs = BATCH_ROWS * BATCH_SIZE_K / localSize;
    const size_t localCopyCountRhs = BATCH_SIZE_K * BATCH_COLS / localSize;

    float value = 0.0f;

    for(size_t k = 0; k < cols; k += BATCH_SIZE_K) {
        barrier(CLK_LOCAL_MEM_FENCE);

        // lhsをローカルメモリーにコピー
        for(size_t i = 0; i < localCopyCountLhs; ++i) {
            // 連続するIDのワークアイテムが、グローバルメモリーの連続領域にアクセスするよう考慮
            const size_t id = (i * localSize) + localId;
            const size_t row = id / BATCH_SIZE_K;
            const size_t col = id %% BATCH_SIZE_K; // formatの都合でパーセントをエスケープ
            localLhs[row][col] = lhs[(groupI + row) * cols + (k + col)];
        }

        // rhsも同様にローカルメモリーにコピーする
        for(size_t i = 0; i < localCopyCountRhs; ++i) {
            const size_t id = (i * localSize) + localId;
            const size_t row = id / BATCH_COLS;
            const size_t col = id %% BATCH_COLS; // formatの都合でパーセントをエスケープ
            localRhs[row][col] = rhs[(k + row) * resultCols + (groupJ + col)];
        }
        barrier(CLK_LOCAL_MEM_FENCE);

        for(size_t lk = 0; lk < BATCH_SIZE_K; ++lk) {
            value += localLhs[localI][lk] * localRhs[lk][localJ];
        }
    }
    result[globalI * resultCols + globalJ] = value;
}

結果

cpu: 45636 msecs, gpu: 55 msecs (218.2 GFLOPS, faster 829.7 times)

微妙に遅くなりましたね……。ただ、今後のプライベートメモリーの活用でBATCH_SIZE_Kの値を調整する必要があるので、いったんこのままにしておきます。

ワークアイテムの担当要素を増やす

さて、次にワークアイテムの担当要素を増やしていきます。

product.d

    // 行列サイズ
    enum {
        ROWS = 1000,
        COLS = 2000,
        RESULT_COLS = 3000,
        BATCH_ROWS = 64,
        BATCH_COLS = 64,
        BATCH_SIZE_K = 16,
        PRIVATE_ROWS = 8, // ワークアイテム当たりの担当行数
        PRIVATE_COLS = 8  // ワークアイテム当たりの担当列数
    }

// 中略

    // OpenCL Cソースに定数の値を埋め込む。
    auto program = cl.createProgramFromSource(
        context, import("product.cl").format(BATCH_ROWS, BATCH_COLS, BATCH_SIZE_K, PRIVATE_ROWS, PRIVATE_COLS));
    scope(exit) cl.releaseProgram(program);
    cl.buildProgram(program, deviceIds);

// 中略

    // 1ワークアイテム当たりの担当要素数が増えた分、ワークアイテム数は減らす。
    immutable(size_t)[] globalWorkSizes = [
        bufferResultCols / PRIVATE_COLS,
        bufferRows / PRIVATE_ROWS
    ];
    immutable(size_t)[] localWorkSizes = [BATCH_COLS / PRIVATE_COLS, BATCH_ROWS / PRIVATE_ROWS];

D言語側の修正は上記で済みますが、OpenCL C側で変化が激しいです。
面倒くさいので全部載せてしまいます。

product.cl

__kernel void product(
        __global const float *lhs,
        __global const float *rhs,
        __global float *result,
        uint rows,
        uint cols,
        uint resultCols) {
    enum {
        BATCH_ROWS = %d,
        BATCH_COLS = %d,
        BATCH_SIZE_K = %d,
        // ここ以下を追加
        PRIVATE_ROWS = %d,
        PRIVATE_COLS = %d,

        // ワークグループ内のワークアイテムの数(行方向・列方向別および合計)
        LOCAL_WORK_COUNT_ROW = BATCH_ROWS / PRIVATE_ROWS,
        LOCAL_WORK_COUNT_COL = BATCH_COLS / PRIVATE_COLS,
        LOCAL_SIZE = LOCAL_WORK_COUNT_ROW * LOCAL_WORK_COUNT_COL,

        // ワークアイテムのローカルメモリーへのコピー回数
        // ワークアイテムの担当要素が増えた分、
        // 1ワークアイテムでローカルメモリーへ複数要素コピーする必要が出てきた。
        LOCAL_COPY_COUNT_LHS = (BATCH_ROWS * BATCH_SIZE_K) / LOCAL_SIZE,
        LOCAL_COPY_COUNT_RHS = (BATCH_SIZE_K * BATCH_COLS) / LOCAL_SIZE
    };

    __local float localLhs[BATCH_ROWS][BATCH_SIZE_K];
    __local float localRhs[BATCH_SIZE_K][BATCH_COLS];

    const size_t localJ = get_local_id(0);
    const size_t localI = get_local_id(1);

    const size_t groupJ = get_group_id(0) * BATCH_COLS;
    const size_t groupI = get_group_id(1) * BATCH_ROWS;

    const size_t localId = get_local_id(0) + get_local_size(0) * get_local_id(1);

    // 担当要素分だけ結果値の変数を持つ。
    float values[PRIVATE_ROWS][PRIVATE_COLS];
    for(size_t i = 0; i < PRIVATE_ROWS; ++i) {
        for(size_t j = 0; j < PRIVATE_COLS; ++j) {
            values[i][j] = 0.0f;
        }
    }

    for(size_t k = 0; k < cols; k += BATCH_SIZE_K) {
        barrier(CLK_LOCAL_MEM_FENCE);
        // ローカルメモリーへの値のコピー。
        // 各ワークアイテムをIDをもとにコピー元の要素へ割り当て、コピーする。
        // ワークアイテムがIDの順序通りグローバルメモリーの連続領域へアクセスするよう気を付ける。
        for(size_t i = 0; i < LOCAL_COPY_COUNT_LHS; ++i) {
            const size_t id = (i * LOCAL_SIZE) + localId;
            const size_t row = id / BATCH_SIZE_K;
            const size_t col = id %% BATCH_SIZE_K;
            localLhs[row][col] = lhs[(groupI + row) * cols + (k + col)];
        }
        // rhs分も同様
        for(size_t i = 0; i < LOCAL_COPY_COUNT_RHS; ++i) {
            const size_t id = (i * LOCAL_SIZE) + localId;
            const size_t row = id / BATCH_COLS;
            const size_t col = id %% BATCH_COLS;
            localRhs[row][col] = rhs[(k + row) * resultCols + (groupJ + col)];
        }
        barrier(CLK_LOCAL_MEM_FENCE);

        for(size_t lk = 0; lk < BATCH_SIZE_K; ++lk) {
            // 担当要素の分だけループし、計算を行う。
            for(size_t i = 0; i < PRIVATE_ROWS; ++i) {
                for(size_t j = 0; j < PRIVATE_COLS; ++j) {
                    // 1つのワークアイテムは、連続ではなく飛び飛びの要素(LOCAL_WORK_COUNT_ROW・COL間隔)を担当することに注意。
                    values[i][j] += localLhs[localI + i * LOCAL_WORK_COUNT_ROW][lk] * localRhs[lk][localJ + j * LOCAL_WORK_COUNT_COL];
                }
            }
        }
    }

    // 計算結果の格納。計算時と逆に代入先アドレスの展開を行う。
    for(size_t i = 0; i < PRIVATE_ROWS; ++i) {
        const size_t rowOffset = (groupI + localI + i * LOCAL_WORK_COUNT_ROW) * resultCols;
        for(size_t j = 0; j < PRIVATE_COLS; ++j) {
            result[rowOffset + (groupJ + localJ + j * LOCAL_WORK_COUNT_COL)] = values[i][j];
        }
    }
}

計算結果

cpu: 48790 msecs, gpu: 13 msecs (923.1 GFLOPS, faster 3753.1 times)

やった、ついにTFLOPS級が出ました！

おそらく、ワークアイテム切り替えのオーバーヘッドが軽減され、より効率的にコアが動けるようになったのでしょう。
ローカルメモリーのコピーもかなり効率化できています。

しかしながら、まだ我々は最後の変身を残しています……。フフフ……。

プライベートメモリーに値をキャッシュする

以下の部分に更なる最適化を施すことができます。

product.cl

        for(size_t lk = 0; lk < BATCH_SIZE_K; ++lk) {
            for(size_t i = 0; i < PRIVATE_ROWS; ++i) {
                for(size_t j = 0; j < PRIVATE_COLS; ++j) {
                    // ループに応じて、localLhs・localRhsの同じ要素を複数回読み込んでいる。
                    values[i][j] += localLhs[localI + i * LOCAL_WORK_COUNT_ROW][lk] * localRhs[lk][localJ + j * LOCAL_WORK_COUNT_COL];
                }
            }
        }

これを下記のようにキャッシュできます。

product.cl

        for(size_t lk = 0; lk < BATCH_SIZE_K; ++lk) {
            // localRhsの値をキャッシュする。
            float cols[PRIVATE_COLS];
            for(size_t j = 0; j < PRIVATE_COLS; ++j) {
                cols[j] = localRhs[lk][localJ + j * LOCAL_WORK_COUNT_COL];
            }
            for(size_t i = 0; i < PRIVATE_ROWS; ++i) {
                // localLhsの値をキャッシュする
                const float row = localLhs[localI + i * LOCAL_WORK_COUNT_ROW][lk];

                // キャッシュを使って計算する
                for(size_t j = 0; j < PRIVATE_COLS; ++j) {
                    values[i][j] += row * cols[j];
                }
            }
        }

計算結果

cpu: 49112 msecs, gpu: 14 msecs (857.1 GFLOPS, faster 3508.0 times)

ぐお、遅くなった……。

しかしさらにもうひとつ最適化を重ねることができます……。

ローカルメモリーを転置する

product.cl

            for(size_t i = 0; i < PRIVATE_ROWS; ++i) {
                // localLhsの値をキャッシュする
                const float row = localLhs[localI + i * LOCAL_WORK_COUNT_ROW][lk];

                // キャッシュを使って計算する
                for(size_t j = 0; j < PRIVATE_COLS; ++j) {
                    values[i][j] += row * cols[j];
                }
            }

実はこのlocalLhsへのアクセスの仕方はあまりよくありません。ワークアイテムが列方向に並んでアクセスしていることになっています。
これはバンクコンフリクトの原因になっている可能性が高いです。
これを解消するためには、localRhsと同様のアクセスができるようlocalLhsを転置させます。

product.cl

    __local float localLhs[BATCH_SIZE_K][BATCH_ROWS + 2]; // 転置させるので行数・列数を逆にする。ついでにバンクコンフリクト除けにパディングを入れる。
    __local float localRhs[BATCH_SIZE_K][BATCH_COLS];

// 中略
        for(size_t i = 0; i < LOCAL_COPY_COUNT_LHS; ++i) {
            const size_t id = (i * LOCAL_SIZE) + localId;
            const size_t row = id / BATCH_SIZE_K;
            const size_t col = id %% BATCH_SIZE_K;
            localLhs[col][row] = lhs[(groupI + row) * cols + (k + col)]; // ここでcol・rowを逆にする
        }

// 中略

            for(size_t i = 0; i < PRIVATE_ROWS; ++i) {
                // 晴れて行への連続アクセスが可能に
                const float row = localLhs[lk][localI + i * LOCAL_WORK_COUNT_ROW];

                // キャッシュを使って計算する
                for(size_t j = 0; j < PRIVATE_COLS; ++j) {
                    values[i][j] += row * cols[j];
                }
            }

ついでに、各種定数もK80向けに調整してしまいます。

product.d

    // matrix size.
    enum {
        // 理由あって行列サイズを大きくした。
        ROWS = 4096,
        COLS = 4096,
        RESULT_COLS = 4096,
        BATCH_ROWS = 128,
        BATCH_COLS = 128,
        //BATCH_SIZE_K = 8, // NVIDIA GeForce GTX 960M ではこれがよかった
        BATCH_SIZE_K = 32, // Tesla K80ではこれがよかった。WARPサイズ関係?
        PRIVATE_ROWS = 8,
        PRIVATE_COLS = 8
    }

計算結果

cpu: 2099435 msecs, gpu: 91 msecs (1510.3 GFLOPS, faster 23070.7 times)

1.5テラ！しかもCPUの23000倍高速！！
もうCPUで計算する気が無くなる速度ですね。色々やってようやくここまで来ました……。

理論性能の1/3くらいまで来ていて、参考サイトのコードとほぼ同じ速度が出ているので、いったんはここで満足しようと思います。

～～おわり～～