IwakenLabShader勉強会でBabylon.jsにおけるComputeShaderでGPUパーティクルを行うデモの準備

にー兄さん

もともとこちらの方で調査を行ってきたが
デモを作るにあたり作業ログをこっちに分離した

にー兄さん

今回作るものの要件などをまとめていく

にー兄さん

Babylon.js WebGPU Engine
Compute Shader
PointsCloundSystsem

上記を用いて、パーティクルの位置をGPGPUで計算するような簡単なデモを作成する
テキトーにパーティクルの位置をもとにsinでアニメーションさせるとかでいいような気がしている
（本当はCurlノイズでうにゃうにゃさせたいけど）

にー兄さん

3d sin waveやるか
下記のイメージをGPUパーティクルで表現

にー兄さん

式としては下記らしい

z = \sin\left(\sqrt{x2+y2}\right)

にー兄さん

まずはPointCloudかParticleSystemで高さ0の地点に
$\left(x,y\right)=\left(\left[-\pi,\pi\right],\left[-\pi,\pi\right]\right)$ の範囲にランダムなパーティクルを敷き詰める
（数式の表記的には誤りだと思うけど雰囲気伝われ）

にー兄さん

敷き詰める、の際に普通にパーティクルに値を入れていく感じではなく
パーティクルとComputeShaderの間で値の参照などをする必要があるため
位置情報は配列に入れておきたい
Vector3[]をComputeShaderに渡すのめんどいかな？
そしたらfloat32[]でもいい

にー兄さん

一番やりたいことに近いのが、配列に計算された値を書き込むBabylon.jsのサンプル

そのComputeShaderはこんな感じだった

matrixMulComputeShader

struct Matrix {
    size : vec2<f32>,
    numbers: array<f32>,
};

@group(0) @binding(0) var<storage,read_write> firstMatrix : Matrix;
@group(0) @binding(1) var<storage,read_write> secondMatrix : Matrix;
@group(0) @binding(2) var<storage,read_write> resultMatrix : Matrix;

@compute @workgroup_size(1, 1, 1)
fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
    resultMatrix.size = vec2<f32>(firstMatrix.size.x, secondMatrix.size.y);

    let resultCell : vec2<u32> = vec2<u32>(global_id.x, global_id.y);
    var result : f32 = 0.0;
    for (var i : u32 = 0u; i < u32(firstMatrix.size.y); i = i + 1u) {
    let a : u32 = i + resultCell.x * u32(firstMatrix.size.y);
    let b : u32 = resultCell.y + i * u32(secondMatrix.size.y);
    result = result + firstMatrix.numbers[a] * secondMatrix.numbers[b];
    }

    let index : u32 = resultCell.y + resultCell.x * u32(secondMatrix.size.y);
    resultMatrix.numbers[index] = result;
}

にー兄さん

そんでもって、このComputeShaderに値をセットする部分のコードはCore

const firstMatrix = new Float32Array([
    2 /* rows */, 4 /* columns */,
    1, 2, 3, 4,
    5, 6, 7, 8
]);

const bufferFirstMatrix = new BABYLON.StorageBuffer(engine, firstMatrix.byteLength);
bufferFirstMatrix.update(firstMatrix);

// ...

cs3.setStorageBuffer("firstMatrix", bufferFirstMatrix);

にー兄さん

つまり、Matrixの定義で
最初二つのfloat32をとり
その後の配列はnumberになる
って感じのsequentialなメモリ配置になっているため

const firstMatrix = new Float32Array([
    2 /* rows */, 4 /* columns */,
    1, 2, 3, 4,
    5, 6, 7, 8
]);

これをしたときにrawとcolumnsはsizeに格納されて
そのほかはnumbersに行くらしい、すごいな

にー兄さん

なるべくパーティクル数は可変にしたい感じはあるけど
いったん100x100に敷き詰めるのがいいのかなって思った

パーティクルの実態は100x100x3の長さのarrayにpositionが記録されている感じで
ComputeShaderのカーネルは100x100x1で起動

パーティクルIDをparticleId = idx.x*100+idx.yで取得できるので

const particleId = idx.x*100+idx.y

arr[particleId*3] = x
arr[particleId*3+1] = y
arr[particleId*3+2] = z

によって位置を設定できる感じだ

にー兄さん

上記はTypeScriptで書いてるけど、本番はWGSLで書く必要がある

にー兄さん

ここまでできれば、あとはパーティクルのupdate関数の中で
位置を配列から参照して適用すればよいということかな

にー兄さん

位置を配列から参照して適用

個々の部分はCPUの処理だからGPUパーティクルとは言えないかな～

にー兄さん

computeShaderのdispatchWhenReadyは非同期メソッドらしい

cs3.dispatchWhenReady(firstMatrix[0], secondMatrix[1]).then(() => {
    bufferResultMatrix.read().then((res) => {
        // we know the result buffer contains floats
        const resFloats = new Float32Array(res.buffer);
        console.log(resFloats);
    });
});

これは以下のように書き直すことも可能だな

await cs3.dispatchWhenReady(firstMatrix[0], secondMatrix[1]);

const res = await bufferResultMatrix.read();

// we know the result buffer contains floats
const resFloats = new Float32Array(res.buffer);
console.log(resFloats);