床井研究室

■ 2023年06月04日 - セパラブルフィルタ

ブログ書いてない

2019 年から丸 4 年ブログを更新しておりませんでした。もちろん、その間に何もしていなかったわけではなく、また書きたいと思っているネタもいくつもありました。でも、なんか毎日が自転車操業状態で、ブログを書くとか、それより大事な論文を書くとかいう作業に全然取り掛かれずにいました。コロナとかでいろいろあったってこともあるんですけど、この状況で逆に研究に集中できたという方もいらっしゃいますから、言い訳にはならんでしょうね。今もやらなきゃなんないことがいくつもあって、忙しい、時間が足りないという焦りで頭が一杯になっています。

でも、これは錯覚に過ぎないっていう話も聞きます。そういうことを ChatGPT に相談してみたりしたのですが、まあ、そのとおりやなっていう回答をもらったものの、気づきみたいなものは得られませんでした。もっと突っ込んで質問すればよかったのでしょうけど、それ以上聞きたいことを思いつきませんでした。

コンピュートシェーダのプログラミングモデル

GPU のプログラミングモデルは、大量のデータを GPU の複数のプロセッサに分配し、各プロセッサはカーネルと呼ばれる共通したプログラムにより、分配された個々のデータに対して同一の処理を同時に実行するというものです。このような処理はデータ並列性 (data parallelism) と呼ばれます。

特に GPU では、分配されたデータが互いに依存していないとして、個々のプロセッサが他のプロセッサの実行状況を考慮することなく、並列に処理を実行します。バーテックスシェーダは CPU から供給された複数の頂点の頂点属性を同時に処理しますが、一つのプロセッサは一つの頂点だけを担当し、他の頂点の頂点属性を参照することはありません。フラグメントシェーダもラスタライザにより指定されたフラグメントのみを担当し、他のフラグメントの状況を考慮することはありません。これにより、各プロセッサは他のプロセッサの状況にかかわらず、フルスピードで動作することができます。これはストリームコンピューティング (stream computing) と呼ばれます。

しかし一般には、入力された複数のデータが相互に依存しているような処理はいくらでも存在します。例えば画像のフィルタリング処理では、入力画像の複数の画素のデータから出力画像の 1 画素の値を決定します。

画像のフィルタリング

このような応用に対応するために、コンピュートシェーダには複数のプロセッサがデータを共有する機能や（共有メモリ）、共有したデータに対する複数のプロセッサからのアクセスの競合を調停する機能（アトミック変数）、およびプロセッサが互いに処理の同期をとる仕組み（バリア）などが導入されています。また、プロセッサを入力データに対してどのように配分するかという設定も、プログラマに提供されています。

コンピュートシェーダでセパラブルフィルタ

今回は 2018 年に書いたコードをネタにします。ずっと書きたかったのですが、先に書いた通り書くことがとっても億劫だったので、先延ばしにしていました。でも最近学生さんに説明する必要が出てきたので、いい加減にちゃんと書こうと思いました。元ネタは NVIDIA の（OpenGL の始祖であらせられる）Mark Kilgard 氏の 2012 年の SIGGRAPH での講演、NVIDIA OpenGL for 2012 です。

画像のフィルタリング処理では、例えば、画像中の $\left(x,y\right)$ の位置にある画素の値 $f\left(x,y\right)$ を、その位置を中心とした $\left(2N+1\right)\times\left(2N+1\right)$ 画素の近傍の画素の平均値 $h\left(x,y\right)$ で置き換えることにより、画像をぼかすことができます。

\[D=\left(2N+1\right)^2\] \[h\left(x,y\right)=\frac{1}{D}\sum_{j=-N}^{N}\sum_{i=-N}^{N}f\left(x+i,y+j\right)\]

この計算を横 $W$ 画素、縦 $H$ 画素の画像に対して画素ごとに独立に行うと、画素へのアクセス回数は（周辺部のはみ出しを無視するとして）$\left(2N+1\right)^2\times W\times H$ 回になります。しかし、これは次のように、行（横）方向の和を求めてから、それを列（縦）方向に合計しても同じです。

\[g\left(x,y\right)=\sum_{i=-N}^{N}f\left(x+i,y\right)\] \[h\left(x,y\right)=\frac{1}{D}\sum_{j=-N}^{N}g\left(x,y+j\right)\]

このとき、この $g\left(x,y\right)$ を画像全体について保持しておけば、その分のメモリが必要になりますが、画素へのアクセスが $2\left(2N+1\right)\times W\times H$ 回になるので、計算量のオーダーを下げることができます。

これはガウスフィルタのような重み付きフィルタの場合も同様です。

\[D=\sum_{j=-N}^{N}\sum_{i=-N}^{N}\exp\left(-\frac{i^2+j^2}{2\sigma^2}\right)\] \[h\left(x,y\right)=\frac{1}{D}\sum_{j=-N}^{N}\sum_{i=-N}^{N}\exp\left(-\frac{i^2+j^2}{2\sigma^2}\right)f\left(x+i,y+j\right)\]

この各画素の重みは、

\[\exp\left(-\frac{i^2+j^2}{2\sigma^2}\right)=\exp\left(-\frac{i^2}{2\sigma^2}\right)\exp\left(-\frac{j^2}{2\sigma^2}\right)\]

ですから、$j$ の項をくくり出して、

\[h\left(x,y\right)=\frac{1}{D}\sum_{j=-N}^{N}\exp\left(-\frac{j^2}{2\sigma^2}\right)\sum_{i=-N}^{N}\exp\left(-\frac{i^2}{2\sigma^2}\right)f\left(x+i,y+j\right)\]

とすると、これも次のように2つのステップに分けることができます。

\[g\left(x,y\right)=\sum_{i=-N}^{N}\exp\left(-\frac{i^2}{2\sigma^2}\right)f\left(x+i,y\right)\] \[h\left(x,y\right)=\frac{1}{D}\sum_{j=-N}^{N}\exp\left(-\frac{j^2}{2\sigma^2}\right)g\left(x,y+j\right)\]

ワークグループとスレッド

さて、このような処理をコンピュートシェーダ上に実装するわけですが、そのためにはまず、コンピュートシェーダのプロセッサを、入力データに対してどのように配分するかを決めておく必要があります。

GPU には多数のプロセッサがあり、それぞれが並列に動作します。この一つのプロセッサに割り当てられた処理をスレッド (thread) といいます。そして、互いに協力して働くことのできるスレッドの集合をワークグループ (workgroup) といいます。一つのワークグループを構成するスレッドの数は、あらかじめ決めておく必要があります。これは (local_size_x, local_size_y, local_size_z) の 3 次元の格子の数で指定します。

またコンピュートシェーダを起動する際は、このワークグループをいくつ使用するかを指定します。このワークグループの数も (num_groups_x, num_groups_y, num_groups_z) の3 次元の格子の数で指定します。ここでは対象が画像なので、local_size_z = num_groups_z = 1 として、2 次元の格子を構成するようにします。

ワークグループとスレッド

ワークグループ同士も、互いに並行して動作します。これらがどれだけ並列に動作するかどうかは、GPU のプロセッサの数などのリソース次第です。世の中、金次第ということですね。

コンピュートシェーダによる画像のフィルタリング

いま、入力画像 image の (x, y) の位置にある画素を中心とした横 filterSize.x、縦 filterSize.y の領域 filter 内の各画素の値から、入力画像と同じサイズの出力画像 filtered の (x, y) の位置にある画素の値を決定することを考えます。これを出力画像のすべての画素に対して行います。

ですが、これを一つ一つの画素に対して順番にやっていては、CPU で処理するのと変わりありません。そこで、複数の画素の値を同時に決定するように、処理を並列化します。ワークグループ内の各スレッドは並行して動作しますから、それぞれに値を決定する画素を割り当てます。ここでは一つのワークグループが処理する画素の領域を tile と呼ぶことにします。この tile の横と縦のサイズを、それぞれ tileSize.x、tileSize.y とします。

画像とタイル

コンピュートシェーダは、この tile を画像に敷き詰めていると考えて実行します。ここでは、この一つの tile に対して一つのワークグループを割り当てます。

個々のワークグループにはワークグループの番号が割り当てられ、組み込み変数 gl_WorkGroupID に格納されます。ワークグループは3 次元の格子に配列されていると見立てていますから、gl_WorkGroupID も 3 次元のベクトルです。同様にワークグループ内の個々のスレッドにもスレッドの番号が割り当てられ、組み込み変数 gl_LocalInvocationID に格納されます。これも 3 次元のベクトルです。これらの番号は、各スレッドが担当する対象（画素など）を決定するための手掛かりに用いられます。

画像のタイリング

一つの tile が参照する入力画像の領域のサイズは、tile の最外周の画素に filter の中心を置いた時の、filter の範囲を含む領域です。filter の中心の画素から最外周の画素までの画素数 $N$ は、横方向と縦方向それぞれ filterOffset.x = floor(filterSize.x / 2)、filterOffset.y = floor(filterSize.y / 2) ですから、これは横が neighborhoodSize.x = tileSize.x + filterOffset.x × 2、縦が neighborhoodSize.y = tileSize.y + filterOffset.y × 2 になります。このサイズの作業用の配列を共有メモリに作成します。フィルタリングの処理は、画像をこの共有メモリにコピーして実行します。

共有メモリのサイズ

ワークグループとタイルのサイズ

さて、ここからがコンピュートシェーダのキモです。いま、このフィルタリングの実行に使うワークグループのサイズを、横 tileSize.x、縦 neighborhoodSize.y とします。つまり、横は tile の横のサイズと同じですが、縦は tile の縦のサイズに filter の縦方向のマージン分を加えたものにします。このサイズを用いる理由は後述します。ここではワークグループのサイズを横 local_size_x = 32、縦 local_size_y = 32 とし（local_size_z は省略しているので 1 になります）、filter のサイズ filterSize を横 filterSize.x = 11、縦 filterSize.y = 11 として、これらからタイルのサイズ tileSize を決めることにします。

#version 430 core
 
// 非 0 ならガウスフィルタ、0 なら平均値（ボックス）フィルタ
#define GAUSS 1
 
// ワークグループのサイズ
layout (local_size_x = 32, local_size_y = 32) in;
 
// フィルタのサイズ
const ivec2 filterSize = ivec2(11, 11);

filter のマージン filterOffset を求めておきます。正の整数の除算なので、小数点以下は切り捨てられます。これはガウスフィルタの重みを計算するときに filter の中心位置としても使います（このため filterSize は奇数である必要があります）。

// フィルタのオフセット
const ivec2 filterOffset = filterSize / 2;

そうすると、一度に処理するタイルの領域 tile のサイズ tileSize は、ワークグループのサイズの縦方向を filter のマージン分切り詰めたものになります。ワークグループのサイズが横 local_size_x = 32、縦 local_size_y = 32、フィルタ filter の縦のサイズが filterSize.y = 11 であれば、タイルのサイズ tileSize は横 tileSize.x = 32、縦 tileSize.y = 32 - floor(11/2) × 2 = 22 になります。

// 一度に処理するタイルの領域のサイズ
const ivec2 tileSize = ivec2(gl_WorkGroupSize) - ivec2(0, filterOffset.y * 2);

また、tile のサイズに filter のマージンを加えて、tile が参照する入力画像の領域のサイズを求めます。

// 近傍を含む領域のサイズ
const ivec2 neighborhoodSize = tileSize + filterOffset * 2;

このサイズの配列 pixel を、共有メモリ上に作成します。これに加えて、横のサイズが tile と同じサイズの vec2 の配列 row も作成しておきます。平均値（ボックス）フィルタやガウスフィルタでは、事前に重みの総和 $D$ を求めることができますが、ここでは後々の都合のために row を vec2 にして、分子と分母のそれぞれの合計を別々に求められるようにしておきます。

// 処理する領域の近傍を含めたコピー
shared float pixel[neighborhoodSize.y][neighborhoodSize.x];
shared vec2 row[neighborhoodSize.y][tileSize.x];

入力画像のフォーマットは、この後の都合から r32f (32bit float, 1 channel) とします。これを image というイメージユニットから入力します。結果は同様に r32f の filtered というイメージユニットに格納します。このほか、ガウスフィルタの重みの分散を uniform 変数で取得します。

// 入力データを格納したイメージユニット
layout (r32f) readonly uniform image2D image;
 
// 出力データを格納するイメージユニット
layout (r32f) writeonly uniform image2D filtered;
 
// ガウスフィルタの重みの分散 (x: 横方向, y: 縦方向)
uniform vec2 variance;

イメージユニットから画素を取得する際、インデックスがイメージユニットの範囲から外れないようにする関数 clampLocation() を用意しておきます。

// インデックスがイメージの領域から外れないようにする
ivec2 clampLocation(ivec2 xy)
{
  return clamp(xy, ivec2(0), imageSize(image) - 1);
}

スレッドの同期をとる関数 retirePhase() を用意します。これは共有メモリへのアクセスの完了と、スレッドのここまでの処理の完了を待ちます。

// 他のスレッドの共有メモリへのアクセス完了と処理完了を待つ
void retirePhase()
{
  memoryBarrierShared();
  barrier();
}

以降、実際の処理を説明します。コンピュートシェーダはバーテックスシェーダやフラグメントシェーダと異なり、処理する対象を自分で決めることができます。しかし、適切な対象を決めるためには、何らかの手掛かりが必要になります。そのために、ワークグループやスレッドに割り当てられる番号を用います。前者は組み込み変数 gl_WorkGroupID、後者は組み込み変数 gl_LocalInvocationID に格納されています。

void main(void)
{
  // ワークグループが処理する領域の基準位置
  const ivec2 tile_xy = ivec2(gl_WorkGroupID);
  
  // スレッドが処理する画素のワークグループにおける相対位置
  const ivec2 thread_xy = ivec2(gl_LocalInvocationID);

これらより、スレッドが担当する入力画像中の画素の位置 dst_xy を求めます。ワークグループの番号に tile のサイズを乗じて、それにスレッドの位置を足します。これは結果を格納するイメージユニットの画素の位置です。

  // スレッドに割り当てられた画素位置
  const ivec2 dst_xy = tile_xy * tileSize + thread_xy;

スレッドが参照するイメージユニット image の領域の基準の画素位置は、格納する画素位置より filter の縦方向のマージン分下げたところにします。

  // スレッドが読み出すイメージ上の画素位置
  const ivec2 src_xy = ivec2(dst_xy.x, dst_xy.y - filterOffset.y);

x, y はスレッドが処理するタイル上の画素位置です。

  // スレッドが処理する画素位置
  const int x = thread_xy.x;
  const int y = thread_xy.y;

スレッドが参照するイメージユニット image の領域の各画素を、共有メモリ上の配列 pixel にコピーします。読み出す画素の位置の決定に先ほど定義した clampLocation() を使います。これにより image からはみ出した部分から読み出そうとしたときは、image の最外周の画素の値になります。これでテクスチャの GL_CLAMP_TO_EDGE に似た処理になります。

ワークグループの各スレッドは、自分がコピーする画素を 1 つ共有メモリ上の配列 pixel にコピーします。このコピーは並列に動作する多数のスレッドによって一気に実行されます。データ並列性というのは、逐次処理では繰り返しループで実行されていた複数のデータに対する処理を、このように多数のプロセッサによる並列処理に置き換えます。

ただし、ここで本当に 1 つの画素しかコピーしないでいると、filter の横のサイズ filterSize.x 分の画素がコピーされません。そこで、ワークグループの横のサイズ tileSize.x 分ずらした位置にある画素をもう一度コピーします。これを共有メモリ上の配列 pixel の横のサイズを超えない範囲で繰り返します。

  // タイルごとに処理する 1 画素を共有メモリにコピーする
  for (int i = 0; i < neighborhoodSize.x; i += tileSize.x)
  {
    if (x + i < neighborhoodSize.x)
    {
      // 読み出し位置
      const ivec2 read_at = clampLocation(ivec2(src_xy.x - filterOffset.x + i, src_xy.y));
      
      // 共有メモリに書き込む
      pixel[y][x + i] = imageLoad(image, read_at).r;
    }
  }

画素を共有メモリにコピー

そして、次の処理に移る前に、すべての画素のコピーが完了し、すべてのスレッドが次の処理に移ることができる時点に到達するのを待つ必要があります。これは、最初に pixel の左端の画素を担当したスレッドは複数の画素のコピーを行う必要があるのに対し、右端の画素を担当したスレッドは1 つの画素しかコピーする必要がないので、すべてのスレッドの処理時間が同じにはならないためです。また、スレッドが本当に並列に動作している独立したプロセッサに割り当てられているかどうかも、GPU の仕様に依存します。

  // 他のスレッドの共有メモリへのアクセス完了と処理完了を待つ
  retirePhase();

セパラブルフィルタ

セパラブルフィルタでは、処理を縦と横に分離して行います。そのため、横方向の処理に必要なスレッドの数（ワークグループのサイズ）は、横は tileSize.x ですが、縦はフィルタの縦方向のマージン分を含めた neighborhoodSize.y になります。また、この結果を使って次に縦方向の処理を行うので、結果を保存するために必要なもう一つの共有メモリ上の配列 row のサイズは、このワークグループのサイズと同じ横 tileSize.x、縦 neighborhoodSize.y になります。

変数 x と y は、ワークグループ内でのこのスレッドの「位置」に相当します。そこからフィルタの横のサイズ filterSize.x 分の画素の重み付け和を求めます。sum は vec2 型なので、第 1 要素 (sum.r) に分子となる画素値の重み付け和、第 2 要素 (sum.g) に重みの和を求めます。重みの和は事前に計算できるので、ここでこんな風に計算する必要はないのですが、今後この部分にちょっと細工するつもりなので、今はこうしておきます。

  // 対象画素の値に重みを掛けたものの合計（分子）と重みの合計（分母）
  vec2 sum = vec2(0.0);
  
  // 横方向の重み付け和を求める
  for (int i = 0; i < filterSize.x; ++i)
  {
    const float c = pixel[y][x + i];
#if GAUSS
    const float d = float(i - filterOffset.x);
    const float e = exp(-0.5 * d * d / variance.x);
    sum += vec2(c, 1.0) * e;
#else
    sum += vec2(c, 1.0);
#endif
  }

横方向の処理

横方向の画素値の重み付け和と重みの和が求まれば、これを共有メモリ上の配列 row に格納します。そのあと、他のスレッドの共有メモリへのアクセスの完了と、ここまでの処理の完了を待ちます。

  // 横方向の重み付け和を保存する
  row[y][x] = sum;
  
  // 他のスレッドの共有メモリへのアクセス完了と処理完了を待つ
  retirePhase();

次に、横方向の画素値の重み付け和と重みの和を縦方向に合計します。この処理には y が tileSize.y より小さいスレッドだけを使います。残りのスレッドの処理はここで完了します。

共有メモリ上の配列 row から横方向の画素値の重み付け和と重みの和を取り出し、その両方に縦方向の重みを乗じて、それぞれを合計します。

  // tile の高さ分の数のスレッドを使って（そのほかのスレッドは終了）
  if (y < tileSize.y)
  {
    // 対象画素の値とその重みのペアを作る
    vec2 sum = vec2(0.0);
    
    // 縦方向の重み付け和を求める
    for (int j = 0; j < filterSize.y; ++j)
    {
      const vec2 c = row[y + j][x];
#if GAUSS
      const float d = float(j - filterOffset.y);
      const float e = exp(-0.5 * d * d / variance.y);
      sum += c * e;
#else
      sum += c;
#endif
    }

縦方向の処理

横方向の画素値の重み付け和と重みの和を縦方向に合計したら、画素値の重み付け和を重みの和で割って、出力画像 filtered の対応する画素に格納します。なお、画像の端のタイルでは dst_xy が出力画像の解像度を超える場合がありますが、OpenGL の Image Load Store の仕様では境界外への imageStore() は自動的に破棄されるため、範囲外アクセスのチェックを行わなくても安全に処理されます。

    // 分子を分母で割って保存する
    imageStore(filtered, dst_xy, vec4(sum.r / sum.g, 0.0, 0.0, 1.0));
  }
}

コンピュートシェーダの起動

このコンピュートシェーダは、次のようにして起動します。タイルの横幅 tile_size_x には 32、縦幅 tile_size_y には 22 を指定します。

  ///
  /// 計算を実行する
  ///
  /// @param width 画像の横の画素数
  /// @param height 画像の縦の画素数
  /// @param tile_size_x タイルの横方向のサイズ (ここでは 32)
  /// @param tile_size_y タイルの縦方向のサイズ (ここでは 22)
  ///
  void execute(GLuint width, GLuint height, GLuint tile_size_x = 32, GLuint tile_size_y = 22) const
  {
    glDispatchCompute((width + tile_size_x - 1) / tile_size_x, (height + tile_size_y - 1) / tile_size_y, 1);
  }

（まだ追記・修正するかもしれん）