Babylon.js WebGPU Cloth Simulation (2) - WebGPU プログラミング

前回アップした、XPBD(extended position-based dynamics) による Cloth シミュレーションの解説です。

１ GPU StorageBuffer
　compute shader で扱うデータは、storage buffer を介して利用する。
mesh の各vertex point のデータは、Particle 構造体とし、そのメンバーとして pos(position), vel(velocity), prePos(previous position), invMass(inverse mass) を設定し、storage buffer の変数とする。

struct Particle {
    pos : vec3<f32>,
    vel : vec3<f32>,
    prePos : vec3<f32>,
    @align(16) invMass: f32,
};
@group(0) @binding(0) var<storage, read_write> particles : array<Particle>;

メンバー変数の alignment は vec3型が 16、f32型が 4 であり、構造体自身の
alignment は、メンバー変数の最大値となるので、16 である。構造体のメモリー配置が 16 の倍数となるように、invMass 変数の前に @align(16) を入れる。これで正しくメモリー配置される。

CPU 側でデータを Storage Buffer にセットする場合、以下のようにする。
this.numParticles　は Particle 数を表す変数。

this.particleData = new Float32Array(this.numParticles * 16);
for (let i = 0; i < this.numParticles; ++i) {
   let p = this.points[i];
   // pos
   this.particleData[16 * i + 0] = p.pos.x;
   this.particleData[16 * i + 1] = p.pos.y;
   this.particleData[16 * i + 2] = p.pos.z;
   this.particleData[16 * i + 3] = 1.0;
   // vel
   this.particleData[16 * i + 4] = p.vel.x;
   this.particleData[16 * i + 5] = p.vel.y;
   this.particleData[16 * i + 6] = p.vel.z;
   this.particleData[16 * i + 7] = 1.0;
   // prePos
   this.particleData[16 * i + 8]  = p.prevPos.x;
   this.particleData[16 * i + 9]  = p.prevPos.y;
   this.particleData[16 * i + 10] = p.prevPos.z;
   this.particleData[16 * i + 11] = 1.0;
   // invMass
   this.particleData[16 * i + 12] = p.invMass;
   // dummy
   this.particleData[16 * i + 13] = 0.0;
   this.particleData[16 * i + 14] = 0.0;
   this.particleData[16 * i + 15] = 0.0;
}
this.particleBuffers = new BABYLON.StorageBuffer(engine,
     this.particleData.byteLength,
     BABYLON.Constants.BUFFER_CREATIONFLAG_VERTEX |
     BABYLON.Constants.BUFFER_CREATIONFLAG_READWRITE);
this.particleBuffers.update(this.particleData);

GPU 側での alignment に合わせて、position データ等は 16 バイト（4 バイト X 4）にする。

   // pos
   this.particleData[16 * i + 0] = p.pos.x;
   this.particleData[16 * i + 1] = p.pos.y;
   this.particleData[16 * i + 2] = p.pos.z;
   this.particleData[16 * i + 3] = 1.0;

invMass データも alignment 16 にしてあるので、もとの invMass データに加えて、12 バイト分のダミーのデータを入れる。

   // invMass
   this.particleData[16 * i + 12] = p.invMass;
   // dummy
   this.particleData[16 * i + 13] = 0.0;
   this.particleData[16 * i + 14] = 0.0;
   this.particleData[16 * i + 15] = 0.0;

２ solve Constraints
　 GPU 計算で拘束を解く方法は、参照ページ “16 Simulation on the GPU” の “Graph Coloring” の方法で行なっている。
　Stretching、Shearing、Bending Constraints で計算する２点は、以下のようにする。GPU 計算では、１回の dispatch（実行）が終了するまで buffer の値は更新されない。そのため、mesh point が重ならないように、passNr や iw、ih パラメータを使って計算する mesh point をコントロールする。

Stretching Constraints
Stretch 1

Stretch 2

分割数が 6 の場合、Stretch 1 における passNr 0, 1 に対応した２点の配置は次のようになる。

Shearing Constraints
Shear 1

Shear 2

Bending Constraints
Bend 1

Bend 2

Stretch 1 を計算する compute shader

solveStretchConst1
struct Params {
    passNr: u32,
};
@group(0) @binding(0) var<uniform> params : Params;

struct Particle {
    pos : vec3<f32>,
    vel : vec3<f32>,
    prePos : vec3<f32>,
    @align(16) invMass: f32,
};
@group(0) @binding(1) var<storage, read_write> particles : array<Particle>;

@compute @workgroup_size(1)
// dispatch iw_max
fn main(@builtin(global_invocation_id) GlobalInvocationID : vec3<u32>) {
    let numSubsteps: u32 = ${numSubsteps};
    let dt: f32 = ${dt};
    let sdt: f32 = dt / f32(numSubsteps);
    let d0: f32 = ${this.defaultSpacingY};
    let compliance: f32 = ${stretchCompliance};

    let iw_max: u32 = ${this.iw_max};
    let ih_max: u32 = ${this.ih_max};
    let passNr: u32 = params.passNr;

    var index : u32 = GlobalInvocationID.x;

    for (var ih: u32 = passNr; ih <= (ih_max - 1u) - 2u + passNr; ih = ih + 2u) {
        var id0: u32 = index + ih * iw_max;
        var id1: u32 = index + (ih + 1u) * iw_max;
        
        var w0: f32 = particles[id0].invMass;
        var w1: f32 = particles[id1].invMass;
        var w: f32 = w0 + w1;

        if (w == 0.0) {return;}

        var p0 : vec3<f32> = particles[id0].pos;
        var p1 : vec3<f32> = particles[id1].pos;
        var grad: vec3<f32> = p1 - p0;
        var d: f32 = length(grad);

        let alpha: f32 = compliance / sdt / sdt;
        var lambda: f32 = (d - d0) / (w + alpha);

        grad = normalize(grad);
        
        var delta_p0: vec3<f32> = grad * (w0 * lambda);
        particles[id0].pos = p0 + delta_p0;

        var delta_p1: vec3<f32> = grad * (-w1 * lambda);
        particles[id1].pos = p1 + delta_p1;
    }
}

拘束を解く式は、
J. Bender, M. Müller, M. Macklin
Position-Based Simulation Methods in Computer Graphics 2017
EUROGRAPHICS 2017 Tutorial (2017)
https://matthias-research.github.io/pages/publications/PBDTutorial2017-CourseNotes.pdf
Course notes の 4.2.6 XPBD と 5.1 Stretching の distance constraints の式を使う。

2点 $\boldsymbol{x}_1$ , $\boldsymbol{x}_2$ の位置補正の式は次のようになる。
　 $\Delta {\boldsymbol{x}}_1 = w_1 \lambda {\boldsymbol{n}}$
　 $\Delta {\boldsymbol{x}}_2 = -w_2 \lambda {\boldsymbol{n}}$

　 $\lambda = \frac{1}{w_1 + w_2 + \tilde{\alpha} \ \ } (|{\boldsymbol{x}}_{2, 1}| - d)$
　 ${\boldsymbol{x}}_{2, 1} = {\boldsymbol{x}}_2 - {\boldsymbol{x}}_1$
　 ${\boldsymbol{n}} = \frac{{\boldsymbol{x}}_{2,1}}{\ |{\boldsymbol{x}}_{2,1}| \ }$
　 $\tilde{\alpha} = \frac{\alpha}{\Delta t^2 \ \ }$

　 $d$ : rest length , $\alpha$ : compliance
　 $w_1$ , $w_2$ : inverse mass of particles 1, 2

Shearing Bending にも同じ distance constraints の式を使う。

３ Uniform Buffer を用いた繰返し処理
　拘束を解く際に用いる passNr や iw、ih の値は uniform buffer を通して shader に渡し、それらの値を変えて繰返し処理を行う。
uniform のある値に対し shader を計算し、次に新しい値に対して同じ shader を使って計算する場合は、別々の ComputeShader object として定義する。そして、それぞれの ComputeShader を使って、uniform の値を変えて計算する。
ComputeShader の定義

// cs solveStretchConst1
   this.csSolveStretchConst1 = [];
   for (let passNr = 0; passNr < 2; passNr++) {
      const csSolve1 = new BABYLON.ComputeShader("csSolve1", engine,
         { computeSource: computeShader.solveStretchConst1 },
         { bindingsMapping: {
             "params": { group: 0, binding: 0 },
             "particle": { group: 0, binding: 1 },
         }
      });
      csSolve1.setUniformBuffer("params", this.uBuffer);
      csSolve1.setStorageBuffer("particle", this.particleBuffers);

      this.csSolveStretchConst1.push(csSolve1);
   }

ComputeShader の実行

for (let passNr = 0; passNr < 2; passNr++) {
   this.updateParams(passNr);
   this.csSolveStretchConst1[passNr].dispatch(this.iw_max);
}

UniformBuffer の設定とupdate メソッド

// Uniform buffer
this.uBuffer = new BABYLON.UniformBuffer(engine, undefined, undefined,
    "uBuffer");
this.uBuffer.addUniform("passNr", 1);

updateParams(passNr) {
   this.uBuffer.updateInt("passNr", passNr);
   this.uBuffer.update();
}

４ワークグループ
　compute shader の integrate と updateVel では GPU のワークグループサイズを 64 に設定している。

@workgroup_size(64)

　integrate では、重力による速度の変化とそれによる位置の変化を計算している。updateVel では、拘束を解いた後の速度を計算している。この部分は、各 mesh point で同じ処理を行なっているので、並列計算が可能である。ワークグループサイズを 64 に設定し、同時処理する mesh point 毎の処理（ワークアイテム）の数を 64 にしている。
（ワークグループサイズの上限はより大きな値であるが、ここでは 64 にしている。）
　shader を実行する際は、dispatch のワークグループ数を numParticles / 64 にする。

// Integrate
   this.csIntegrate.dispatch(Math.ceil(this.numParticles / 64));