「这是我参与2022首次更文挑战的第9天，活动详情查看：2022首次更文挑战」。

本文基于 austinEng/webgpu-samples: WebGPU Samples (github.com) 仓库中的代码进行说明。

前言

接上文，在Hello WebGPU —— Compute Shader基础 - 掘金 (juejin.cn)这篇文章中，我们初步利用WebGPU提供的计算管线完成了矩阵的计算，今天我们开始做一点点稍微有难度的事情：模拟自然界中的鸟群运动方式。我们先来看最终的实现效果：

这些三角形表示就是一个鸟类个体，它们一开始都在杂乱无章的飞。

经过一段时间后，它们的运动会变成下面这样：

我们可以看到，随着程序的运行，鸟类个体会逐渐的形成群体性的运动。我们这样定义每个鸟类个体的运动规律：

它们彼此希望离更大的群体更近一些
如果它们两个之间靠的太近，它们又希望分离的稍微远一些
如果它们两个之间离的又太远，它们又会互相的靠近一点。

今天我们就利用WebGPU提供的通用计算能力来模拟这样的鸟类群体行为。

同样，今天的工作分为以下几步:

数据准备
编写compute shader
编写vertex shader & fragment shader
创建计算管线
创建渲染管线
编写渲染流程

Coding

让我们开始今天的Coding，首先我们需要准备相关的数据。

数据准备

首先，我们先确定鸟群中每个个体的初始位置信息：

  const numParticles = 1500;
  const initialParticleData = new Float32Array(numParticles * 4);
  for (let i = 0; i < numParticles; ++i) {
    initialParticleData[4 * i + 0] = 2 * (Math.random() - 0.5);
    initialParticleData[4 * i + 1] = 2 * (Math.random() - 0.5);
    initialParticleData[4 * i + 2] = 2 * (Math.random() - 0.5) * 0.1;
    initialParticleData[4 * i + 3] = 2 * (Math.random() - 0.5) * 0.1;
  }

为我们上面提到的运动规则设置参数，并为其创建Buffer，在每次渲染时进行写入

const simParams = {
    deltaT: 0.04,
    rule1Distance: 0.1, // 如果两个个体之间的距离小于0.1，我们认为他们是一个群体
    rule2Distance: 0.025, // 如果两个个体之间的距离小于0.025，则认为他们靠的太近，需要分开一点点
    rule3Distance: 0.03, // 如果两个个体之间的距离小于0.03，则认为他们离的太远，希望他们靠近彼此一些
    rule1Scale: 0.02, // 规则1的权重
    rule2Scale: 0.05, // 规则2的权重
    rule3Scale: 0.005, // 规则3的权重
  };
  
  

  const simParamBufferSize = 7 * Float32Array.BYTES_PER_ELEMENT;
  const simParamBuffer = device.createBuffer({
    size: simParamBufferSize,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });

然后，我们需要2套 GPUBuffer 和 GPUBindGroup 对象，其中的一套用于存储当前的鸟群信息（包括位置、速度），另一套用于存储计算的结果。计算完毕后我们需要将两套对象交换一下顺序，也就是说，拿第一次的结果当做第二次的输入，用第一次的输入来接收第二次的计算结果。


  const particleBuffers: GPUBuffer[] = new Array(2);
  const particleBindGroups: GPUBindGroup[] = new Array(2);
  for (let i = 0; i < 2; ++i) {
    particleBuffers[i] = device.createBuffer({
      size: initialParticleData.byteLength,
      usage: GPUBufferUsage.VERTEX | GPUBufferUsage.STORAGE,
      mappedAtCreation: true,
    });
    new Float32Array(particleBuffers[i].getMappedRange()).set(
      initialParticleData
    );
    particleBuffers[i].unmap();
  }

  for (let i = 0; i < 2; ++i) {
    particleBindGroups[i] = device.createBindGroup({
      layout: computePipeline.getBindGroupLayout(0),
      entries: [
        {
          binding: 0,
          resource: {
            buffer: simParamBuffer,
          },
        },
        {
          binding: 1,
          resource: {
            buffer: particleBuffers[i],
            offset: 0,
            size: initialParticleData.byteLength,
          },
        },
        {
          binding: 2,
          resource: {
            buffer: particleBuffers[(i + 1) % 2],
            offset: 0,
            size: initialParticleData.byteLength,
          },
        },
      ],
    });
  }

这里注意，由于 particleBuffers 在计算着色器和顶点着色器中都需要使用，所以usage为： GPUBufferUsage.VERTEX | GPUBufferUsage.STORAGE. OK，到目前为止，需要参与计算的数据已经准备完毕，我们现在来准备一下渲染鸟群需要用的数据。我们来为鸟群中的每个个体准备顶点数据，这里简单起见，我们就使用一个简单的三角形来表示一个鸟类个体。

  const vertexBufferData = new Float32Array([
    -0.01, -0.02, 0.01,
    -0.02, 0.0, 0.02,
  ]);

  const spriteVertexBuffer = device.createBuffer({
    size: vertexBufferData.byteLength,
    usage: GPUBufferUsage.VERTEX,
    mappedAtCreation: true,
  });
  new Float32Array(spriteVertexBuffer.getMappedRange()).set(vertexBufferData);
  spriteVertexBuffer.unmap();

编写compute shader

struct Particle {
  pos : vec2<f32>;
  vel : vec2<f32>;
};
struct SimParams {
  deltaT : f32;
  rule1Distance : f32;
  rule2Distance : f32;
  rule3Distance : f32;
  rule1Scale : f32;
  rule2Scale : f32;
  rule3Scale : f32;
};
struct Particles {
  particles : array<Particle>;
};
@binding(0) @group(0) 
var<uniform> params : SimParams;
@binding(1) @group(0) 
var<storage, read> particlesA : Particles;
@binding(2) @group(0) 
var<storage, read_write> particlesB : Particles;

@stage(compute) @workgroup_size(64)
fn main(@builtin(global_invocation_id) GlobalInvocationID : vec3<u32>) {
  var index : u32 = GlobalInvocationID.x;

  var vPos = particlesA.particles[index].pos;
  var vVel = particlesA.particles[index].vel;
  var cMass = vec2<f32>(0.0, 0.0);
  var cVel = vec2<f32>(0.0, 0.0);
  var colVel = vec2<f32>(0.0, 0.0);
  var cMassCount : u32 = 0u;
  var cVelCount : u32 = 0u;
  var pos : vec2<f32>;
  var vel : vec2<f32>;
  // 遍历鸟群中的所有个体，与当前计算的鸟类个体按规则进行运动
  for (var i : u32 = 0u; i < arrayLength(&particlesA.particles); i = i + 1u) {
    if (i == index) {
      continue;
    }

    pos = particlesA.particles[i].pos.xy;
    vel = particlesA.particles[i].vel.xy;
    if (distance(pos, vPos) < params.rule1Distance) {
      // 统计相距较近的鸟类数量
      cMass = cMass + pos;
      cMassCount = cMassCount + 1u;
    }
    if (distance(pos, vPos) < params.rule2Distance) {
      // 靠得太近就分开一些
      colVel = colVel - (pos - vPos);
    }
    if (distance(pos, vPos) < params.rule3Distance) {
      // 离开的太远就靠近一些
      cVel = cVel + vel;
      cVelCount = cVelCount + 1u;
    }
  }
  if (cMassCount > 0u) {
    // 类似于计算某一小群鸟类的聚集中心
    var temp = f32(cMassCount);
    cMass = (cMass / vec2<f32>(temp, temp)) - vPos;
  }
  if (cVelCount > 0u) {
    var temp = f32(cVelCount);
    cVel = cVel / vec2<f32>(temp, temp);
  }
  vVel = vVel + (cMass * params.rule1Scale) + (colVel * params.rule2Scale) +
      (cVel * params.rule3Scale);

  // 限制速度范围
  vVel = normalize(vVel) * clamp(length(vVel), 0.0, 0.1);
  // 更新位置信息
  vPos = vPos + (vVel * params.deltaT);
  // 越界处理
  if (vPos.x < -1.0) {
    vPos.x = 1.0;
  }
  if (vPos.x > 1.0) {
    vPos.x = -1.0;
  }
  if (vPos.y < -1.0) {
    vPos.y = 1.0;
  }
  if (vPos.y > 1.0) {
    vPos.y = -1.0;
  }
  // 将计算结果写回Buffer中
  particlesB.particles[index].pos = vPos;
  particlesB.particles[index].vel = vVel;
}

这里再次讲解一下关于 workgroup_size的问题，它表示的是 workgroup 的大小，表示的是一次性并行执行shader的核心数量。那么我们在ts代码中调用的 dispatch API 则是指定了 workgroup的个数。所以shader程序执行的次数等于 workgroup_size * workgroup_count，在这个例子中，我们的workgroup_size的大小是(64, 1, 1), 我们调用 dispatch 时， workgroup_count 为24。所以，我们实际调用 compute shader 的次数为： 24 x 64 = 1536 次。

后续的代码逻辑可以查看上述代码中的注释，此处就不再过多的赘述了。

编写vertex shader & fragment shader

接下来，我们编写用于渲染的 shader 代码

@stage(vertex)
fn vert_main(@location(0) a_particlePos : vec2<f32>,
             @location(1) a_particleVel : vec2<f32>,
             @location(2) a_pos : vec2<f32>) -> @builtin(position) vec4<f32> {
  let angle = -atan2(a_particleVel.x, a_particleVel.y);
  let pos = vec2<f32>(
      (a_pos.x * cos(angle)) - (a_pos.y * sin(angle)),
      (a_pos.x * sin(angle)) + (a_pos.y * cos(angle)));
  return vec4<f32>(pos + a_particlePos, 0.0, 1.0);
}

@stage(fragment)
fn frag_main() -> @location(0) vec4<f32> {
  return vec4<f32>(1.0, 1.0, 1.0, 1.0);
}

首先，我们根据当前鸟类个体的速度求出它的朝向（角度），然后让顶点旋转这个角度，最后再加上我们在之前的计算shader中求得的最新的位置即可。

对于片元着色器器这里写的就更简单了，设置为白色即可。

创建计算管线


  const computePipeline = device.createComputePipeline({
    compute: {
      module: device.createShaderModule({
        code: updateSpritesWGSL,
      }),
      entryPoint: 'main',
    },
  });

这里创建计算管线的代码比较简单，不过多赘述。

创建渲染管线

const spriteShaderModule = device.createShaderModule({ code: spriteWGSL });
  const renderPipeline = device.createRenderPipeline({
    vertex: {
      module: spriteShaderModule,
      entryPoint: 'vert_main',
      buffers: [
        {
          // instanced particles buffer
          arrayStride: 4 * 4,
          stepMode: 'instance',
          attributes: [
            {
              // instance position
              shaderLocation: 0,
              offset: 0,
              format: 'float32x2',
            },
            {
              // instance velocity
              shaderLocation: 1,
              offset: 2 * 4,
              format: 'float32x2',
            },
          ],
        },
        {
          // vertex buffer
          arrayStride: 2 * 4,
          stepMode: 'vertex',
          attributes: [
            {
              // vertex positions
              shaderLocation: 2,
              offset: 0,
              format: 'float32x2',
            },
          ],
        },
      ],
    },
    fragment: {
      module: spriteShaderModule,
      entryPoint: 'frag_main',
      targets: [
        {
          format: presentationFormat,
        },
      ],
    },
    primitive: {
      topology: 'triangle-list',
    },
  });

这里大部分的代码与之前讲解的渲染流程几乎一致，但是这里有一个需要注意的地方，我们为buffers 数组中的第一个 buffer 指定了 stepMode 为 instance 。

stepMode 一共有两个模式： vertex 和 instance。

vertex: 顶点数据的地址基于 arrayStride 不断的进行累加，但是在两个实例之间，顶点数据的地址会被重置
instance: 顶点数据的地址基于 arrayStride 不断的进行累加，但是在两个实例之间，但是顶点数据的地址不会被重置

也就是说如果我们采用了多实例绘制的方法来绘制图形，并且我们希望不同的实例之间读取到的buffer数据是不同的，我们就需要使用到这种 instance 的模式。

编写渲染流程

现在我们的准备工作基本已经准备就绪，可以开始编写渲染流程了。

首先，在渲染之前，先进行计算的流程。整套流程使用一个 commandEncoder 即可。

// 这里用一个 t 临时变量来交换两个bindGroup和 particleBuffer
let t = 0;
const commandEncoder = device.createCommandEncoder();
{
  const passEncoder = commandEncoder.beginComputePass();
  passEncoder.setPipeline(computePipeline);
  passEncoder.setBindGroup(0, particleBindGroups[t % 2]);
  passEncoder.dispatch(Math.ceil(numParticles / 64));
  passEncoder.endPass();
}

计算流程与渲染流程不同的地方在于:

passEncoder的创建不再是通过 beginRenderPass 创建，而是通过 beginComputePass 创建了。
不通过调用 draw 命令进行渲染，而是调用 dispatch API来进行计算。

计算完毕后，我们可以进行渲染了。

{
  const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
  passEncoder.setPipeline(renderPipeline);
  passEncoder.setVertexBuffer(0, particleBuffers[(t + 1) % 2]);
  passEncoder.setVertexBuffer(1, spriteVertexBuffer);
  passEncoder.draw(3, numParticles, 0, 0);
  passEncoder.endPass();
}

这里有一点值得说明一下，还记得我们之前创建渲染管线对象的时候，我们有3个 attributes，并且绑定了3个location, 但是这里我们却只设置了2个 GPUBuffer，这是因为 particleBuffer 同时包含了 position 和 velocity 两种数据。我们需要注意 setVertexBuffer 这个函数的第一个参数并不是与 location 一一对应的，而是与 PSO 中的 buffers 中的索引对应起来的。

后续，我们只需要提交我们的 commandEncoder 就可以了，然后更新我们的临时变量 t ，再开启下一次渲染

device.queue.submit([commandEncoder.finish()]);
++t;
requestAnimationFrame(frame);

总结

让我们来简单回顾一下今天学习的内容：

首先我们为鸟群指定了飞行的规则；
再利用Compute Shader进行鸟群位置的计算，将计算的结果写入到 GPUBuffer 中，再在渲染流程中复用了包含计算结果的 GPUBuffer 对其进行渲染。
学习到了渲染流程中为 GPUBuffer 不同的 stepMode，对于多实例绘制时不同的实例应用不同的顶点数据时，需要使用到 instance的 stepMode。

OK，今天的内容就是这么多了，如果您觉得本文有用，请个作者点个赞哦~ 您的赞赏就是作者更新的动力~

Hello WebGPU —— 模拟鸟群

前言