【OpenCL】行列の積のカーネルの最適化

はじめに

Cra2yPierr0tさんが書いている記事軽率にGPUを使っていこう、OpenCL入門の最後に
「データの再利用を最適化するためにはブロック化の技術が必要になる。行列をプライベートメモリにちょうど収まるようにタイルに分解したり、タイルをローカルメモリにコピーしたり、タイル間で乗算を行う技術である。」とあります。
「タイル」というのを使えばまだ早くなるということなのかな。
また時間測定については、clEnqueueNDRangeKernel、clFinishのみを測定します。clEnqueueReadBufferは測定しません。
軽率にGPUを使っていこう、OpenCL入門を読んだ後に、本記事をお読みください。

測定時間

1024x1024の行列の掛け算の測定です。
時間は32回実行したときの平均時間を載せています。
最後のタイルを除いて軽率にGPUを使っていこう、OpenCL入門のパクリです。

名称	処理時間	備考
ベース(mat_mul_base)	0.063908 s	これを基準とします。
行ごとに計算(mat_mul_row)	0.508425 s	ベースより遅くなる。
L1L2キャッシュ(mul_mat_l1l2)	0.373094 s	ベースより遅くなる。
タイル(mat_mul_tile)	0.006084 s	これでいいんじゃない？めちゃ速い。
おまけ(CPU)	8.879012 s	直列って遅ーい。

ベース(mat_mul_base)

work-itemがNxNのカーネルです。これでもCPUと比べたらめちゃ速いです。

__kernel void mul_mat_base(__global float* A,
                           __global float* B,
                           unsigned int N,
                           __global float* C) {
  int i = get_global_id(0);
  int j = get_global_id(1);
  float tmp = 0.0f;
  for (int k = 0; k < N; k++) {
    tmp += A[i * N + k] * B[k * N + j];
  }
  C[i * N + j] = tmp;
}

行ごとに計算(mat_mul_row)

行のみをglobal_idにします。私の環境で遅くなります。

__kernel void mul_mat_row(__global float* A,
                          __global float* B,
                          unsigned int N,
                          __global float* C) {
  int i = get_global_id(0);
  float tmp;
  for (int j = 0; j < N;j++) {
    tmp = 0.0f;
    for (int k = 0; k < N; k++) {
      tmp += A[i * N + k] * B[k * N + j];
    }
    C[i * N + j] = tmp;
  }
}

L1L2キャッシュ(mul_mat_l1l2)

軽率にGPUを使っていこう、OpenCL入門では「Aの一列分だけグローバルメモリからプライベートメモリに移してやる。」とあるんですが、多分float Awrk[1024]はプライベートメモリには入りきらないのでグローバルメモリに書きこまれると思います。
それでも「行ごとに計算(mat_mul_row)」より速くなるのは、L1/L2キャッシュが効くかららしいです。

__kernel void mul_mat_l1l2(__global float* A,
                           __global float* B,
                           unsigned int N,
                           __global float* C) {
  int i = get_global_id(0);
  float tmp;

  float Awrk[1024];
  for (int j = 0; j < N; j++) {
    Awrk[j] = A[i * N + j];
  }

  for (int j = 0; j < N; j++) {
    tmp = 0.0f;
    for (int k = 0; k < N; k++) {
      tmp += Awrk[k] * B[k * N + j];
    }
    C[i * N + j] = tmp;
  }
}

タイル(mul_mat_tile)

これが目玉です、というか私の記事はここだけですね。でもこれ、めちゃ速くなります。
ブロック行列をイメージしながらソースを脳内で実行してみてください。
図を描こうと思ったんですが、かえって分かりにくいかと思って、図無しです。気が向いたら描きます。
分からない人は global を {4, 4}, Local を {2, 2} として、実際に計算してみるとよいです。
タイルのサイズが16x16なのはCL_DEVICE_MAX_WORK_GROUP_SIZEが私のGPUでは256だからです。1024の人は32x32を試してみてください。

#define TILE 16

__kernel void mul_mat_tile(__global const float* A,
                           __global const float* B,
                           unsigned int N,
                           __global float* C) {
  // C の座標
  int gx = get_global_id(0);
  int gy = get_global_id(1);
  // ローカルの座標
  int lx = get_local_id(0);
  int ly = get_local_id(1);

  // ワークグループ内で共有するタイル
  __local float Atile[TILE][TILE];
  __local float Btile[TILE][TILE];

  float sum = 0.0f;

  // タイルをスライドさせていく
  int num_tiles = (N + TILE - 1) / TILE;

  for (int t = 0; t < num_tiles; t++) {
    int ac = t * TILE + lx; // A の列
    int ar = gy;            // A の行
    int bc = gx;            // B の列
    int br = t * TILE + ly; // B の行

    // A と B のタイルを local memory にロード
    Atile[ly][lx] = A[ar * N + ac];
    Btile[ly][lx] = B[br * N + bc];

    // 全ての lx, ly が書き終わるまで待つ
    barrier(CLK_LOCAL_MEM_FENCE);

    // タイル内で積和
    for (int k = 0; k < TILE; k++) {
      sum += Atile[ly][k] * Btile[k][lx];
    }

    // タイルの計算を終えるまで待つ
    barrier(CLK_LOCAL_MEM_FENCE);
  }

  // C に書き込み
  C[gy * N + gx] = sum;  
}

最後に

行列とベクトルの積も試したんですが、速くならなかったです。
本記事は生成AI Copilot に聞いたものもあるので、間違いがあれば指摘してください。

参考リンク

軽率にGPUを使っていこう、OpenCL入門

参考文献

GPU を支える技術

機械学習基礎理論独習

誤りがあればご指摘いただけると幸いです。数式が整うまで少し時間かかります。リンクフリーです。