Tweak Particle Sample (#492)

I was investigating why this sample didn't work on compat mode in Chrome.

I found 1 issue in the code which is the shader that makes the probability map
only worked if the source image is was a multiple of 64 pixels wide.
This was because it was not conditionally skipping out of bounds pixels
and so processing those pixels would end up overwriting other results.

This issue does not come up in the sample but if you change the image to
some image that is not a multiple of 64 pixels across the issue comes up.

In debugging why GL was failing, one my paths was to make a 16x16 pixel
`F` using the canvas 2D api and passing that in so that I could print
all of the data. With that I ran into this issue.

So, I added the skipping. Thank's to Peter McNeeley for finding that issue.

While pursing the code I find somethings I though maybe should be refactored?

* Calculating the number of mip level was done in a loop.

  Changed it to just math.

* Moving some calculations to use properties on the texture.

  I can revert these changes. They were useful when I switched the
  source image to a canvas since then the code didn't depend on variables that
  were no longer available.

* Getting rid of `struct Buffer`

  I can revert these changes. I get the impression that this was
  a limitation of an earlier version of WGSL? It seems more confusing
  to me personally to have a struct with a single array in it than
  to just have the array.

* Simpler math

  I can revert this but, changed `sum = dot(vec4f(a, b, c, d), vec4f(1.0))`
  to `sum = a + b + c + d`.

Note that, even with the change for the first pass, the code only works
for square power-of-2 images. Both the probability generation and the
usage of that data require each mip level to be exactly 1/2 the size
of the previous level. So, added some asserts so people copying this
won't stumble into issues.

Author: greggman

sample/particles/main.ts

@@ -182,39 +182,30 @@ quadVertexBuffer.unmap();
 //////////////////////////////////////////////////////////////////////////////
 // Texture
 //////////////////////////////////////////////////////////////////////////////
-let texture: GPUTexture;
-let textureWidth = 1;
-let textureHeight = 1;
-let numMipLevels = 1;
-{
-  const response = await fetch('../../assets/img/webgpu.png');
-  const imageBitmap = await createImageBitmap(await response.blob());
-
-  // Calculate number of mip levels required to generate the probability map
-  while (
-    textureWidth < imageBitmap.width ||
-    textureHeight < imageBitmap.height
-  ) {
-    textureWidth *= 2;
-    textureHeight *= 2;
-    numMipLevels++;
-  }
-  texture = device.createTexture({
-    size: [imageBitmap.width, imageBitmap.height, 1],
-    mipLevelCount: numMipLevels,
-    format: 'rgba8unorm',
-    usage:
-      GPUTextureUsage.TEXTURE_BINDING |
-      GPUTextureUsage.STORAGE_BINDING |
-      GPUTextureUsage.COPY_DST |
-      GPUTextureUsage.RENDER_ATTACHMENT,
-  });
-  device.queue.copyExternalImageToTexture(
-    { source: imageBitmap },
-    { texture: texture },
-    [imageBitmap.width, imageBitmap.height]
-  );
-}
+const isPowerOf2 = (v: number) => Math.log2(v) % 1 === 0;
+const response = await fetch('../../assets/img/webgpu.png');
+const imageBitmap = await createImageBitmap(await response.blob());
+assert(imageBitmap.width === imageBitmap.height, 'image must be square');
+assert(isPowerOf2(imageBitmap.width), 'image must be a power of 2');
+
+// Calculate number of mip levels required to generate the probability map
+const mipLevelCount =
+  (Math.log2(Math.max(imageBitmap.width, imageBitmap.height)) + 1) | 0;
+const texture = device.createTexture({
+  size: [imageBitmap.width, imageBitmap.height, 1],
+  mipLevelCount,
+  format: 'rgba8unorm',
+  usage:
+    GPUTextureUsage.TEXTURE_BINDING |
+    GPUTextureUsage.STORAGE_BINDING |
+    GPUTextureUsage.COPY_DST |
+    GPUTextureUsage.RENDER_ATTACHMENT,
+});
+device.queue.copyExternalImageToTexture(
+  { source: imageBitmap },
+  { texture: texture },
+  [imageBitmap.width, imageBitmap.height]
+);
 
 //////////////////////////////////////////////////////////////////////////////
 // Probability map generation
@@ -247,22 +238,22 @@ let numMipLevels = 1;
     usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
   });
   const buffer_a = device.createBuffer({
-    size: textureWidth * textureHeight * 4,
+    size: texture.width * texture.height * 4,
     usage: GPUBufferUsage.STORAGE,
   });
   const buffer_b = device.createBuffer({
-    size: textureWidth * textureHeight * 4,
+    size: buffer_a.size,
     usage: GPUBufferUsage.STORAGE,
   });
   device.queue.writeBuffer(
     probabilityMapUBOBuffer,
     0,
-    new Int32Array([textureWidth])
+    new Uint32Array([texture.width])
   );
   const commandEncoder = device.createCommandEncoder();
-  for (let level = 0; level < numMipLevels; level++) {
-    const levelWidth = textureWidth >> level;
-    const levelHeight = textureHeight >> level;
+  for (let level = 0; level < texture.mipLevelCount; level++) {
+    const levelWidth = Math.max(1, texture.width >> level);
+    const levelHeight = Math.max(1, texture.height >> level);
     const pipeline =
       level == 0
         ? probabilityMapImportLevelPipeline.getBindGroupLayout(0)
@@ -472,3 +463,9 @@ function frame() {
 }
 configureContext();
 requestAnimationFrame(frame);
+
+function assert(cond: boolean, msg = '') {
+  if (!cond) {
+    throw new Error(msg);
+  }
+}

sample/particles/probabilityMap.wgsl

@@ -2,17 +2,12 @@ struct UBO {
   width : u32,
 }
 
-struct Buffer {
-  weights : array<f32>,
-}
-
 @binding(0) @group(0) var<uniform> ubo : UBO;
-@binding(1) @group(0) var<storage, read> buf_in : Buffer;
-@binding(2) @group(0) var<storage, read_write> buf_out : Buffer;
+@binding(1) @group(0) var<storage, read> buf_in : array<f32>;
+@binding(2) @group(0) var<storage, read_write> buf_out : array<f32>;
 @binding(3) @group(0) var tex_in : texture_2d<f32>;
 @binding(3) @group(0) var tex_out : texture_storage_2d<rgba8unorm, write>;
 
-
 ////////////////////////////////////////////////////////////////////////////////
 // import_level
 //
@@ -21,9 +16,13 @@ struct Buffer {
 ////////////////////////////////////////////////////////////////////////////////
 @compute @workgroup_size(64)
 fn import_level(@builtin(global_invocation_id) coord : vec3u) {
-  _ = &buf_in;
+  _ = &buf_in; // so the bindGroups are similar.
+  if (!all(coord.xy < vec2u(textureDimensions(tex_in)))) {
+    return;
+  }
+
   let offset = coord.x + coord.y * ubo.width;
-  buf_out.weights[offset] = textureLoad(tex_in, vec2i(coord.xy), 0).w;
+  buf_out[offset] = textureLoad(tex_in, vec2i(coord.xy), 0).w;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -36,19 +35,21 @@ fn import_level(@builtin(global_invocation_id) coord : vec3u) {
 ////////////////////////////////////////////////////////////////////////////////
 @compute @workgroup_size(64)
 fn export_level(@builtin(global_invocation_id) coord : vec3u) {
-  if (all(coord.xy < vec2u(textureDimensions(tex_out)))) {
-    let dst_offset = coord.x    + coord.y    * ubo.width;
-    let src_offset = coord.x*2u + coord.y*2u * ubo.width;
+  if (!all(coord.xy < vec2u(textureDimensions(tex_out)))) {
+    return;
+  }
 
-    let a = buf_in.weights[src_offset + 0u];
-    let b = buf_in.weights[src_offset + 1u];
-    let c = buf_in.weights[src_offset + 0u + ubo.width];
-    let d = buf_in.weights[src_offset + 1u + ubo.width];
-    let sum = dot(vec4f(a, b, c, d), vec4f(1.0));
+  let dst_offset = coord.x    + coord.y    * ubo.width;
+  let src_offset = coord.x*2u + coord.y*2u * ubo.width;
 
-    buf_out.weights[dst_offset] = sum / 4.0;
+  let a = buf_in[src_offset + 0u];
+  let b = buf_in[src_offset + 1u];
+  let c = buf_in[src_offset + 0u + ubo.width];
+  let d = buf_in[src_offset + 1u + ubo.width];
+  let sum = a + b + c + d;
 
-    let probabilities = vec4f(a, a+b, a+b+c, sum) / max(sum, 0.0001);
-    textureStore(tex_out, vec2i(coord.xy), probabilities);
-  }
+  buf_out[dst_offset] = sum / 4.0;
+
+  let probabilities = vec4f(a, a+b, a+b+c, sum) / max(sum, 0.0001);
+  textureStore(tex_out, vec2i(coord.xy), probabilities);
 }