WebGL2でGPGPU近傍探索

BoidsやSPHといったパーティクル系のシミュレーションでは一定距離内の他のパーティクルとの相互作用を行うため近傍の探索を行う必要があります。単純な実装では、すべてのパーティクルについて他のすべてのパーティクルとの距離を計算するため計算量が$O(n^{2})$となります。これだと、パーティクル数が増えると著しくパフォーマンスが低下します。

この記事では近傍探索の高速化の手法として、GPUを使った近傍探索を行いたいと思います。

アルゴリズムは以下の記事を参考にします。この記事ではUnityのComputeShaderを使っていますが、WebGLにはComputeShaderがないのでテクスチャを使ったGPGPUで行います。
[\[\#Unity\] ComputeShaderで近傍探索の高速化 \- Qiita](https://qiita.com/kodai100/items/d1ddb3f86ba5dccccf9f)

## サンプル

近傍数に応じてヒートマップで色付けをおこなうサンプルを作成しました。近傍数が多いほど暖色になります。
https://aadebdeb.github.io/WebGL2_GPGPU_NeighborSearch/gpu-2d
![neighbor-search-2d.png](https://qiita-image-store.s3.amazonaws.com/0/61569/cfc5a27f-ca04-b377-3b3b-808950839dfa.png)

同じことを３次元でおこなうバージョンも作成しました。
https://aadebdeb.github.io/WebGL2_GPGPU_NeighborSearch/gpu-3d
![neighbor-search-3d.png](https://qiita-image-store.s3.amazonaws.com/0/61569/d28ba0af-18a4-9b3b-32e9-ea615cb86624.png)


Githubにプロジェクトを置いておきましたので、ソースコード全文はそちらを確認ください。
パフォーマンスの比較用にCPUで近傍探索をおこなうバージョンも置いてあります。
[aadebdeb/WebGL2\_GPGPU\_NeighborSearch: Implementation of WebGL2 GPGPU Neighbor Search](https://github.com/aadebdeb/WebGL2_GPGPU_NeighborSearch)

以下では2次元のサンプルをもとに解説していきます。

## 初期設定

### フレームバッファの作成

初期設定として以下のような計算に使用するフレームバッファとそこに紐づけるテクスチャを作成しておきます。
この記事では、シミュレーション空間を区切る一つ一つの格子のことをバケット(`bucket`)と呼びます。

* `particleFramebuffer`:
    * `positionTexture`: パーティクルの位置を保存するテクスチャ
        * x要素: パーティクルのx座標の位置
        * y要素: パーティクルのy座標の位置
    * `velocityTexture`: パーティクルの速度を保存するテクスチャ
        * x要素: パーティクルのx方向の速さ
        * y要素: パーティクルのy方向の速さ
* `bucketFramebuffer`
    * `bucketTexture`: パーティクルのインデックスとそのパーティクルが所属するバケットのインデックスを保存するテクスチャ
        * x要素: パーティクルが所属するバケットのインデックス
        * y要素: パーティクルのインデックス
* `bucketReferrerFramebuffer`
    * `bucketReferrerTexture`: 各バケットの`bucktetTexture`でのインデックスの範囲を保存するテクスチャ
        * x要素: `bucketTexture`での開始インデックス
        * y要素: `bucketTexture`での終了インデックス

`particleFramebuffer`では一つのピクセルが一つのパーティクルを表しています。今回はテクスチャのサイズが縦横同じ大きさで２の累乗になるようにしてます。以下の`particleTextureSize`がテクスチャの一辺の大きさで、`particleNum`がパーティクル数になります。(パーティクルを格納するテクスチャのサイズは2の累乗でなくても問題ないかもしれませんが、後述するバイトニックソートではパーティクル数が2の累乗になるようにする必要があります)

```js
const particleTextureSizeN = data['N'];
const particleTextureSize = 2**particleTextureSizeN;
const particleNumN = particleTextureSizeN * 2;
const particleNum = particleTextureSize * particleTextureSize;
```

`bucketTexture`も各パーティクルのidとそのバケットのidを保持するため、`particleFramebuffer`内のテクスチャと同じサイズになるようにします。

`bucketReferrerTexture`は一つのピクセルが一つのバケットを表しており、そのバケットの`bucketTexture`での開始インデックスと終了インデックスをそれぞれ保持します。バケット数は近傍の探索範囲により異なるので、きれいに2の累乗になりません。そのため、以下のようにバケット数よりもピクセル数が多くなるようにテクスチャサイズを決めておきます。

```js
const viewRadius = data['view radius'];
const bucketSize = Math.ceil(1.0 / (2.0 * viewRadius));
let bucketReferrerTextureSize;
for (let i = 0; ; i++) {
  bucketReferrerTextureSize = 2**i;
  if (bucketReferrerTextureSize > bucketSize) {
    break;
  }
}
```


### 位置と速度の初期化

Multiple Render Targetsを使って各パーティクルの初期位置と初期速度をまとめて`positionTexture`と`velocityTexture`に書き出します。

```glsl:particleFramebufferへのレンダリング
#version 300 es

precision highp float;

layout (location = 0) out vec2 o_position;
layout (location = 1) out vec2 o_velocity;

uniform vec2 u_randomSeed;

float random(vec2 x){
  return fract(sin(dot(x,vec2(12.9898, 78.233))) * 43758.5453);
}

void main(void) {
  // 位置の初期化
  o_position = vec2(
    sqrt(random(gl_FragCoord.xy * 0.013 + u_randomSeed + vec2(32.19, 27.51))),
    sqrt(random(gl_FragCoord.xy * 0.029 + u_randomSeed + vec2(19.56, 11.34)))
  );
  // 速度の初期化
  o_velocity = (vec2(
    random(gl_FragCoord.xy * 0.059 + u_randomSeed + vec2(27.31, 16.91)),
    random(gl_FragCoord.xy * 0.038 + u_randomSeed + vec2(25.95, 19.47))
  ) * 2.0 - 1.0) * 0.05;
}
```

## 位置と速度の更新

ここからの処理は、レンダリングループ内で行います。
今回はパーティクルを等速で移動させ、境界で反射させています。

```glsl
#version 300 es

precision highp float;

layout (location = 0) out vec2 o_position;
layout (location = 1) out vec2 o_velocity;

uniform sampler2D u_positionTexture;
uniform sampler2D u_velocityTexture;
uniform float u_deltaTime;

void main(void) {
  ivec2 coord = ivec2(gl_FragCoord.xy);
  vec2 position = texelFetch(u_positionTexture, coord, 0).xy;
  vec2 velocity = texelFetch(u_velocityTexture, coord, 0).xy;

  // 位置と速度の更新
  vec2 nextPosition = position + u_deltaTime * velocity;
  vec2 nextVelocity = velocity;

  // 境界を超えた場合は反射処理
  if (nextPosition.x <= 0.0) {
    nextVelocity.x *= -1.0;
    nextPosition.x += u_deltaTime * nextVelocity.x;
  }
  if (nextPosition.x >= 1.0) {
    nextVelocity.x *= -1.0;
    nextPosition.x += u_deltaTime * nextVelocity.x;
  }
  if (nextPosition.y <= 0.0) {
    nextVelocity.y *= -1.0;
    nextPosition.y += u_deltaTime * nextVelocity.y;
  }
  if (nextPosition.y >= 1.0) {
    nextVelocity.y *= -1.0;
    nextPosition.y += u_deltaTime * nextVelocity.y;
  }

  o_position = nextPosition;
  o_velocity = nextVelocity;
}
```

## バケットの構築

バケットの構築を行うJavaScriptのコードは以下のようになります。初めに各パーティクルについてどのバケットに属しているかを`bucketTexture`に書き込みます。その後、`bucketTexture`のx要素に含まれているバケットのインデックスが昇順に並ぶようにバイトニックソートで並び替えを行います。最後に`bucketReferrerTexture`に各バケットの`bucketTexture`での開始インデックスと終了インデックスを書き込みます。

```js
const constructBuckets = function() {
  // bucketTextureに各パーティクルについてバケットのインデックスを書き込み
  initializeBucket();

  // バイトニックソートでバケットのインデックスで昇順になるように並び替え
  for (let i = 0; i < particleNumN; i++) {
    for (let j = 0; j <= i; j++) {
      swapBucketIndex(i, j);
    }
  }
  
  // bucketReferrerテクスチャに書き込み
  initializeBucketRefrrer();
}
```

### `bucketTexture`の初期化

`bucketTexture`のx要素に各パーティクルが属するバケットのインデックスを、y要素にそのパーティクルのインデックスを書き込みます。`convertCoordToIndex`は2次元のテクスチャ上の座標を1次元のインデックスに変換するための関数です。位置からバケットのインデックスを求める`getBucketIndex`関数でも2次元から1次元への変換を行っています。

```glsl:bucketFramebufferへのレンダリング
#version 300 es

precision highp float;

out uvec2 o_particle;

uniform sampler2D u_positionTexture;
uniform float u_viewRadius; //  近傍の探索半径

float simulationSpace = 1.0;

// 2次元の座標を1次元のインデックスに変換する関数
uint convertCoordToIndex(uvec2 coord, uint sizeX) {
  return coord.x + sizeX * coord.y;
}

// 位置からバケットのインデックスを求める関数
uint getBucketIndex(vec2 position) {
  uvec2 bucketCoord = uvec2(position / (2.0 * u_viewRadius));
  uvec2 bucketNum = uvec2(simulationSpace / (2.0 * u_viewRadius)) + 1u;
  return bucketCoord.x + bucketCoord.y * bucketNum.x;
}

void main(void) {
  vec2 position = texelFetch(u_positionTexture, ivec2(gl_FragCoord.xy), 0).xy;
  uint positionTextureSizeX = uint(textureSize(u_positionTexture, 0).x);
  uint particleIndex = convertCoordToIndex(uvec2(gl_FragCoord.xy), positionTextureSizeX);
  uint bucketIndex = getBucketIndex(position);
  o_particle = uvec2(bucketIndex, particleIndex);
}
```

### `bucketTexture`のソート

バイトニックソートで`bucketTexture`のx要素に保持するバケットのインデックスが昇順に並ぶようにソートします。
WebGLでのバイトニックソートについては私が以前に書いた記事で解説しているため、この記事での解説は省略します。
[WebGL2でGPGPUバイトニックソート \- Qiita](https://qiita.com/aa_debdeb/items/e04aa08bcb9ff9be8e32)

### `bucketReferrerTexture`の作成

各バケットごとに`bucketTexture`のx要素でそのバケットのインデックスが開始するインデックスと終了するインデックスを求めます。今回の実装では、ソート済み配列から特定の値を見つけるということでバイナリサーチを利用します。
重複要素のあるバイナリサーチで指定した値の範囲を取得するアルゴリズムについてはJavaScriptで実装したものもあるので、必要があれば参考にしてください。
[JavaScriptで重複のあるソート済み配列から指定した値の範囲をバイナリサーチで取得する \- Qiita](https://qiita.com/aa_debdeb/items/c31cef34f1b980440fcc)

`bucketReferrerTexture`は各ピクセルが各バケットに対応しており全ピクセル数がバケット数よりも多くなるようにしているため、対応したバケットがないピクセルもあります。そのようなピクセルには適当な値`(-1, -1)`を書き込んでおきます。また対応したバケットのインデックスが`bucketTexture`にない場合、つまりそのバケットにパーティクルが一つもない場合にも`(-1, -1)`を書き込むようになっています。

```glsl
#version 300 es

precision highp float;
precision highp usampler2D;

out ivec2 o_referrer;

uniform ivec2 u_bucketReferrerTextureSize;
uniform usampler2D u_bucketTexture;
uniform float u_viewRadius;
uniform int u_particleNumN;

float simulationSpace = 1.0;

// 2次元の座標を1次元のインデックスに変換する関数 
int convertCoordToIndex(ivec2 coord, int sizeX) {
  return coord.x + sizeX * coord.y;
}

// 1次元のインデックスから2次元の座標に変換する関数
ivec2 convertIndexToCoord(int index, int sizeX) {
  return ivec2(index % sizeX, index / sizeX);
}

int getBucketIndex(int particleIndex, int particleTextureSizeX) {
  return int(texelFetch(u_bucketTexture, ivec2(convertIndexToCoord(particleIndex, particleTextureSizeX)), 0).x);
}

// バイナリサーチで指定した値を持つ最も小さいインデックスを求める関数
int binarySearchMinIndex(int target, int from, int to, int particleTextureSizeX) {
  for (int i = 0; i < u_particleNumN + 1; i++) {
    int middle = from + (to - from) / 2;
    int bucketIndex = getBucketIndex(middle, particleTextureSizeX);
    if (bucketIndex < target) {
      from = middle + 1;
    } else {
      to = middle;
    }
    if (from == to) {
      if (getBucketIndex(from, particleTextureSizeX) == target) {
        return from;
      } else {
        return -1;
      }
    }
  }
  return -1;
}

// バイナリサーチで指定した値を持つ最も大きいインデックスを求める関数
int binarySearchMaxIndex(int target, int from, int to, int particleTextureSizeX) {
  for (int i = 0; i < u_particleNumN + 1; i++) {
    int middle = from + (to - from) / 2 + 1;
    int bucketIndex = getBucketIndex(middle, particleTextureSizeX);
    if (bucketIndex > target) {
      to = middle - 1;
    } else {
      from = middle;
    }
    if (from == to) {
      if (getBucketIndex(from, particleTextureSizeX) == target) {
        return from;
      } else {
        return -1;
      }
    }
  }
  return -1;
}

// バイナリサーチで指定した値の範囲を求める関数
ivec2 binarySearchRange(int target, int from, int to) {
  int particleTextureSizeX = textureSize(u_bucketTexture, 0).x;
  from =  binarySearchMinIndex(target, from, to, particleTextureSizeX);
  to = from == -1 ? -1 : binarySearchMaxIndex(target, from, to, particleTextureSizeX);
  return ivec2(from, to);
}

void main(void) {
  int bucketIndex = convertCoordToIndex(ivec2(gl_FragCoord.xy), u_bucketReferrerTextureSize.x);
  ivec2 bucketNum = ivec2(simulationSpace / (2.0 * u_viewRadius)) + 1;
  int maxBucketIndex = bucketNum.x * bucketNum.y;

  // 対応したバケットのないピクセルには(-1, -1)を書き込む
  if (bucketIndex >= maxBucketIndex) {
    o_referrer = ivec2(-1, -1);
    return;
  }

  ivec2 particleTextureSize = textureSize(u_bucketTexture, 0);
  int particleNum = particleTextureSize.x * particleTextureSize.y;

  o_referrer = binarySearchRange(bucketIndex, 0, particleNum - 1);
}
```

以上で近傍を探索するためにバケット構造の作成が完了しました。

## 近傍の探索

構築したバケット構造をもとに近傍の探索をおこない、ヒートマップに基づき色を決定します。
パーティクルの位置情報をもとに現在のバケットとそれに隣接するバケットを求めます。各バケット内のパーティクルについて距離計算を行い、一定距離以内であればその影響を加算します。今回はパーティクル間の距離が小さいほど影響が大きくなるようにしています。

今回はバケットの大きさを近傍探索半径の2倍、つまり直径にしているのでパーティクル自身が所属するバケットとその近傍3つの計４つのバケットのみを考慮するば大丈夫です。もしバケットの大きさを近傍探索半径にした場合は、パーティクル自身が所属するバケットとその8近傍の計9バケットを探索する必要があります。

```glsl
#version 300 es

precision highp isampler2D;
precision highp usampler2D;

out vec3 o_color;

uniform sampler2D u_positionTexture;
uniform usampler2D u_bucketTexture;
uniform isampler2D u_bucketReferrerTexture;
uniform float u_viewRadius;
uniform float u_maxValue;

float simulationSpace = 1.0;

// 1次元のインデックスから2次元の座標に変換する関数
ivec2 convertIndexToCoord(int index, int sizeX) {
  return ivec2(index % sizeX, index / sizeX);
}

// 指定したバケット内の近傍パーティクルの影響を取得する関数
float findNeighbors(vec2 position, ivec2 bucketPosition, ivec2 bucketNum, int particleTextureSizeX, int bucketReferrerTextureSizeX) {
  // バケットが存在しない場合は0を返す
  if (bucketPosition.x < 0 || bucketPosition.x >= bucketNum.x || bucketPosition.y < 0 || bucketPosition.y >= bucketNum.y) {
    return 0.0;
  }
  int bucketIndex = bucketPosition.x + bucketNum.x * bucketPosition.y;
  ivec2 coord = convertIndexToCoord(bucketIndex, bucketReferrerTextureSizeX);

  ivec2 bucketReferrer = texelFetch(u_bucketReferrerTexture, coord, 0).xy;

  // バケット内にパーティクルがない場合は0を返す
  if (bucketReferrer.x == -1 || bucketReferrer.y == -1) {
    return 0.0;
  }

  float sum = 0.0;
  for (int i = bucketReferrer.x; i <= bucketReferrer.y; i++) {
    uvec2 bucket = texelFetch(u_bucketTexture, convertIndexToCoord(i, particleTextureSizeX), 0).xy;

    int particleIndex = int(bucket.y);
    // 近傍探索をおこなっているパーティクル自身のときは無視する
    if (gl_VertexID == particleIndex) {
      continue;
    }
    ivec2 particleCoord = convertIndexToCoord(particleIndex, particleTextureSizeX);
    vec2 particlePos = texelFetch(u_positionTexture, particleCoord, 0).xy;
    // 一定距離以内であれば影響を加算
    if (length(position - particlePos) < u_viewRadius) {
      // 距離が近いほど値が大きくなるようにする
      sum += (u_viewRadius - length(position - particlePos)) / u_viewRadius;
    }
  }
  return sum;
}

vec3[7] HEAT_MAP_COLORS = vec3[7](
  vec3(0.5),
  vec3(0.0, 0.0, 1.0),
  vec3(0.0, 0.5, 0.5),
  vec3(0.0, 1.0, 0.0),
  vec3(0.5, 0.5, 0.0),
  vec3(0.8, 0.4, 0.0),
  vec3(1.0, 0.0, 0.0)
);

vec3 getHeatmapColor(float value, float maxValue) {
  float step = maxValue / 7.0;
  float i = min(value, maxValue - 0.0001) / step;
  return mix(HEAT_MAP_COLORS[int(i)], HEAT_MAP_COLORS[int(i) + 1], fract(i));
}

void main(void) {
  ivec2 coord = convertIndexToCoord(gl_VertexID, textureSize(u_positionTexture, 0).x);
  vec2 position = texelFetch(u_positionTexture, coord, 0).xy;

  int particleTextureSizeX = textureSize(u_positionTexture, 0).x;
  int bucketReferrerTextureSizeX = textureSize(u_bucketReferrerTexture, 0).x;

  // 位置をもとにバケットの二次元座標を計算
  vec2 bucketPosition = position / (2.0 * u_viewRadius);
  int xOffset = fract(bucketPosition.x) < 0.5 ? -1 : 1;
  int yOffset = fract(bucketPosition.y) < 0.5 ? -1 : 1;

  // 考慮するバケットを求める
  ivec2 bucketPosition00 = ivec2(bucketPosition);
  ivec2 bucketPosition10 = bucketPosition00 + ivec2(xOffset, 0);
  ivec2 bucketPosition01 = bucketPosition00 + ivec2(0, yOffset);
  ivec2 bucketPosition11 = bucketPosition00 + ivec2(xOffset, yOffset);

  ivec2 bucketNum = ivec2(simulationSpace / (2.0 * u_viewRadius)) + 1;

  float sum = 0.0;
  // 各バケットの影響を加算
  sum += findNeighbors(position, bucketPosition00, bucketNum, particleTextureSizeX, bucketReferrerTextureSizeX);
  sum += findNeighbors(position, bucketPosition10, bucketNum, particleTextureSizeX, bucketReferrerTextureSizeX);
  sum += findNeighbors(position, bucketPosition01, bucketNum, particleTextureSizeX, bucketReferrerTextureSizeX);
  sum += findNeighbors(position, bucketPosition11, bucketNum, particleTextureSizeX, bucketReferrerTextureSizeX);

  o_color = getHeatmapColor(sum, u_maxValue);
}
```

ここで求めた色に基づきレンダリングをおこないます。レンダリングについては特筆すべきことがないので解説は省略します。

## 終わりに

WebGLでGPGPU近傍探索をおこなう方法を紹介しました。ComputeShaderのアルゴリズムをテクスチャを使ったGPGPUで実装し直しているため、シェーダーでバイナリサーチを実装するなど若干トリッキーな実装になってしまいました。より良い方法があればぜひ教えてください。

肝心のパフォーマンスですが、CPU実装と比べて動かせるパーティクル数に各段の違いがありました。特に強いGPUを積んでいるPCであれば顕著な差が出るのではないでしょうか。ただ、近傍探索範囲を大きくすると探索対象のパーティクルが増えるためかパフォーマンス的にきつそうでした。