WebGPU Picking

README.md

WebGPU Picking

view on jsgist

index.css

html, body { margin: 0; height: 100%; font-family: monospace; }
canvas { width: 100%; height: 100%; display: block; }
#info {
  position: absolute;
  right: 0;
  bottom: 0;
  padding: 0.5em;
  background-color: rgba(0, 0, 0, 0.9);
  color: white;
  min-width: 7em;
}
#fail {
  position: fixed;
  left: 0;
  top: 0;
  width: 100%;
  height: 100%;
  display: flex;
  justify-content: center;
  align-items: center;
  background: red;
  color: white;
  font-weight: bold;
  font-family: monospace;
  font-size: 16pt;
  text-align: center;
}
.shape:hover {
  fill: red;
}

index.html

<canvas></canvas>
<div id="info">
  <label><input type=checkbox id="move" checked>animate</label>
  <div id="pick"><div>
</div>
<div id="fail" style="display: none">
  <div class="content"></div>
</div>
<script src="https://mrdoob.github.io/stats.js/build/stats.min.js"></script>

index.js

import {vec3, mat4} from 'https://webgpufundamentals.org/3rdparty/wgpu-matrix.module.js';
import * as twgl from 'https://twgljs.org/dist/4.x/twgl-full.module.js';

const degToRad = d => d * Math.PI / 180;
const rand = (min, max) => Math.random() * (max - min) + min;

async function main() {
  const gpu = navigator.gpu;
  if (!gpu) {
    fail('this browser does not support webgpu');
    return;
  }

  const adapter = await gpu.requestAdapter();
  if (!adapter) {
    fail('this browser appears to support WebGPU but it\'s disabled');
    return;
  }

  const device = await adapter.requestDevice();

  const canvas = document.querySelector('canvas');
  const context = canvas.getContext('webgpu');
  const presentationFormat = navigator.gpu.getPreferredCanvasFormat(adapter);

  const numCircles = 1000;

  const shaderSrc = `
  struct VSUniforms {
    worldViewProjection: mat4x4f,
  };
  @group(0) @binding(0) var<uniform> vsUniforms: VSUniforms;

  struct MyVSInput {
      @location(0) position: vec4f,
      @location(1) normal: vec3f,
  };

  struct MyVSOutput {
    @builtin(position) position: vec4f,
    @location(0) color: vec4f,
  };

  @vertex
  fn myVSMain(v: MyVSInput) -> MyVSOutput {
    var vsOut: MyVSOutput;
    vsOut.position = vsUniforms.worldViewProjection * v.position;
    vsOut.color = vec4f(v.normal * 0.5 + 0.5, 1.0) * 0.0 + 1.0;
    return vsOut;
  }

  struct FSUniforms {
    colorMult: vec4f,
  };

  struct PickUniforms {
    id: u32,
  };

  @group(0) @binding(1) var<uniform> fsUniforms: FSUniforms;
  @group(0) @binding(1) var<uniform> pickUniforms: PickUniforms;

  @fragment
  fn myFSMain(v: MyVSOutput) -> @location(0) vec4f {
    return v.color * fsUniforms.colorMult;
  }

  @fragment
  fn pickFSMain(v: MyVSOutput) -> @location(0) u32 {
    return pickUniforms.id;
  }

  `;

  const shaderModule = device.createShaderModule({code: shaderSrc});

  const pipeline = device.createRenderPipeline({
    vertex: {
      module: shaderModule,
      entryPoint: 'myVSMain',
      buffers: [
        // position
        {
          arrayStride: 3 * 4, // 3 floats, 4 bytes each
          attributes: [
            {shaderLocation: 0, offset: 0, format: 'float32x3'},
          ],
        },
        // normals
        {
          arrayStride: 3 * 4, // 3 floats, 4 bytes each
          attributes: [
            {shaderLocation: 1, offset: 0, format: 'float32x3'},
          ],
        },
      ],
    },
    fragment: {
      module: shaderModule,
      entryPoint: 'myFSMain',
      targets: [
        {format: presentationFormat},
      ],
    },
    layout: 'auto',
    primitive: {
      topology: 'triangle-list',
      cullMode: 'back',
    },
    depthStencil: {
      depthWriteEnabled: true,
      depthCompare: 'less',
      format: 'depth24plus',
    },
  });

  const pickPipeline = device.createRenderPipeline({
    vertex: {
      module: shaderModule,
      entryPoint: 'myVSMain',
      buffers: [
        // position
        {
          arrayStride: 3 * 4, // 3 floats, 4 bytes each
          attributes: [
            {shaderLocation: 0, offset: 0, format: 'float32x3'},
          ],
        },
        // normals
        {
          arrayStride: 3 * 4, // 3 floats, 4 bytes each
          attributes: [
            {shaderLocation: 1, offset: 0, format: 'float32x3'},
          ],
        },
      ],
    },
    fragment: {
      module: shaderModule,
      entryPoint: 'pickFSMain',
      targets: [
        {format: 'r32uint'},
      ],
    },
    layout: 'auto',
    primitive: {
      topology: 'triangle-list',
      cullMode: 'back',
    },
    depthStencil: {
      depthWriteEnabled: true,
      depthCompare: 'less',
      format: 'depth24plus',
    },
  });

  function createBuffer(device, data, usage) {
    const buffer = device.createBuffer({
      size: data.byteLength,
      usage,
      mappedAtCreation: true,
    });
    const dst = new data.constructor(buffer.getMappedRange());
    dst.set(data);
    buffer.unmap();
    return buffer;
  }

  function createGeo(device, vertices) {
    const d = twgl.primitives.deindexVertices(vertices);
    const f = twgl.primitives.flattenNormals(d);
    return {
      buffers: {
        position: createBuffer(device, f.position, GPUBufferUsage.VERTEX),
        normal: createBuffer(device, f.normal, GPUBufferUsage.VERTEX),
      },
      numVerts: f.position.length / 3,
    };
  }

  function updateUniformBuffer(device, ub) {
    device.queue.writeBuffer(
      ub.uniformBuffer,
      0,
      ub.values
    );        
  }

  const sphereBufferInfo = createGeo(device, twgl.primitives.createSphereVertices(10, 24, 12));

  const shapes = [
    sphereBufferInfo,
  ];

  const hsl = (h, s, l) => `hsl(${h * 360 | 0}, ${s * 100}%, ${l * 100 | 0}%)`;

  const cssColorToRGBA8 = (() => {
    const canvas = new OffscreenCanvas(1, 1);
    const ctx = canvas.getContext('2d', {willReadFrequently: true});
    return cssColor => {
      ctx.clearRect(0, 0, 1, 1);
      ctx.fillStyle = cssColor;
      ctx.fillRect(0, 0, 1, 1);
      return Array.from(ctx.getImageData(0, 0, 1, 1).data);
    };
  })();

  const cssColorToRGBA = cssColor => cssColorToRGBA8(cssColor).map(v => v / 255);

  const objects = [];
  const numObjects = numCircles;
  for (let ii = 0; ii < numObjects; ++ii) {
    const vUniformBufferSize = 1 * 16 * 4; // 1 mat4s * 16 floats per mat * 4 bytes per float
    const fUniformBufferSize = 4 * 4;      // 1 vec3 * 3 floats per vec3 * 4 bytes per float

    const vsUniformBuffer = device.createBuffer({
      size: vUniformBufferSize,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
    const fsUniformBuffer = device.createBuffer({
      size: fUniformBufferSize,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
    const pickFSUniformBuffer = device.createBuffer({
      size: fUniformBufferSize,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
    const vsUniformValues = new Float32Array(1 * 16); // 1 mat4s
    const worldViewProjection = vsUniformValues.subarray(0, 16);

    const fsUniformValues = new Float32Array(4);  // 1 vec4
    const colorMult = fsUniformValues.subarray(0, 4);

    const pickFSUniformValues = new Uint32Array(1);  // 1 uint32
    const id = pickFSUniformValues.subarray(0, 1);

    colorMult.set(cssColorToRGBA(hsl(Math.random(), 1, 0.75)));
    id.set([ii + 1]);

    const bindGroup = device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries: [
        { binding: 0, resource: { buffer: vsUniformBuffer } },
        { binding: 1, resource: { buffer: fsUniformBuffer } },
      ],
    });

    const pickBindGroup = device.createBindGroup({
      layout: pickPipeline.getBindGroupLayout(0),
      entries: [
        { binding: 0, resource: { buffer: vsUniformBuffer } },
        { binding: 1, resource: { buffer: pickFSUniformBuffer } },
      ],
    });

    const object = {
      v: {
        translation: [rand(-100, 100), rand(-100, 100), rand(-150, -50)],
        xRotationSpeed: rand(0.8, 1.2),
        yRotationSpeed: rand(0.8, 1.2),
        velocity: [rand(-10, 10), rand(-10, 10)],
      },
      vs: {
        uniformBuffer: vsUniformBuffer,
        values: vsUniformValues,
      },
      fs: {
        uniformBuffer: fsUniformBuffer,
        values: fsUniformValues,
      },
      pickFS: {
        uniformBuffer: pickFSUniformBuffer,
        values: pickFSUniformValues,
      },
      uniforms: {
        worldViewProjection,
        colorMult,
        id,
      },
      bufferInfo: shapes[ii % shapes.length],
      bindGroup,
      pickBindGroup,
    };
    updateUniformBuffer(device, object.fs);
    updateUniformBuffer(device, object.pickFS);

    objects.push(object);
  }

  function computeMatrix(viewProjectionMatrix, translation, xRotation, yRotation, dst) {
    mat4.translate(viewProjectionMatrix, translation, dst);
    mat4.rotateX(dst, xRotation, dst);
    mat4.rotateY(dst, yRotation, dst);
    mat4.scale(dst, [0.6, 0.6, 0.6], dst);
  }

  const renderPassDescriptor = {
    colorAttachments: [
      {
        view: undefined, // Assigned later
        clearValue: [1, 1, 1, 1],
        loadOp: 'clear',
        storeOp: 'store',
      },
    ],
    depthStencilAttachment: {
      view: undefined,  // Assigned later
      depthClearValue: 1,
      depthLoadOp: 'clear',
      depthStoreOp: 'store',
    },
  };
  
  let depthTexture;
  let depthTextureView;

  function resizeToDisplaySize(device) {
    const width = Math.min(device.limits.maxTextureDimension2D, canvas.clientWidth);
    const height = Math.min(device.limits.maxTextureDimension2D, canvas.clientHeight);

    const needResize = width !== canvas.width ||
                       height !== canvas.height ||
                       !depthTexture;
    if (needResize) {
      if (depthTexture) {
        depthTexture.destroy();
      }

      canvas.width = width;
      canvas.height = height;

      context.configure({
        device,
        format: presentationFormat,
        compositingAlphaMode: "premultiplied",
      });

      depthTexture = device.createTexture({
        size: [width, height],
        format: 'depth24plus',
        usage: GPUTextureUsage.RENDER_ATTACHMENT,
      });
      depthTextureView = depthTexture.createView();
    }
    return needResize;
  }

  const fieldOfViewRadians = degToRad(60);

  let mouseX = -1;
  let mouseY = -1;
  let oldPickNdx = -1;
  let oldPickColor = new Float32Array(4);
  let waitingForPreviousResults = false;

  const pickRenderPassDescriptor = {
    colorAttachments: [
      {
        view: undefined, // Assigned later
        clearValue: [0, 0, 0, 0],
        loadOp: 'clear',
        storeOp: 'store',
      },
    ],
    depthStencilAttachment: {
      view: undefined,  // Assigned later
      depthClearValue: 1,
      depthLoadOp: 'clear',
      depthStoreOp: 'store',
    },
  };

  const numPixels = 1;
  const pickBuffer = device.createBuffer({
    size: numPixels * 4,
    usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
  });
  let pickTexture;
  let pickDepthTexture;
  let pickMapElaspedTime;
  const pickTextureSize = [];
  const infoElem = document.querySelector('#pick');

  function resizeTexture(tex, oldSize, newSize, format) {
    if (!tex || oldSize[0] !== newSize[0] || oldSize[1] !== newSize[1]) {
      if (tex) {
        tex.destroy();
      }

      tex = device.createTexture({
          size: [...newSize, 1],
          format,
          usage:
            GPUTextureUsage.COPY_SRC | 
            GPUTextureUsage.RENDER_ATTACHMENT,
        });
    }
    return tex;
  }

  async function drawWithIdsForMaterialsAndGetPixelUnderMouse(pixelX, pixelY) {
    if (pixelX < 0 || pixelY < 1 || pixelX >= canvas.width || pixelY >= canvas.height) {
      return 0;
    }
    pickTexture = resizeTexture(pickTexture, pickTextureSize, [canvas.width, canvas.height], 'r32uint');
    pickDepthTexture = resizeTexture(pickDepthTexture, pickTextureSize, [canvas.width, canvas.height], 'depth24plus');

    pickRenderPassDescriptor.colorAttachments[0].view = pickTexture.createView();
    pickRenderPassDescriptor.depthStencilAttachment.view = pickDepthTexture.createView();

    const commandEncoder = device.createCommandEncoder();
    const passEncoder = commandEncoder.beginRenderPass(pickRenderPassDescriptor);
    passEncoder.setPipeline(pickPipeline);
    for (const object of objects) {
      passEncoder.setBindGroup(0, object.pickBindGroup);
      passEncoder.setVertexBuffer(0, object.bufferInfo.buffers.position);
      passEncoder.setVertexBuffer(1, object.bufferInfo.buffers.normal);
      passEncoder.draw(object.bufferInfo.numVerts);
    }
    passEncoder.end();
    commandEncoder.copyTextureToBuffer({
      texture: pickTexture,
      // mipLevel: 0,
      origin: {
        x: pixelX,
        y: pixelY,
      }

    }, {
      buffer: pickBuffer,
      bytesPerRow: ((numPixels * 4 + 255) | 0) * 256,
      rowsPerImage: 1,
    }, {
      width: numPixels,
      // height: 1,
      // depth: 1,
    });
    device.queue.submit([commandEncoder.finish()]);
    const startTime = performance.now();
    await pickBuffer.mapAsync(GPUMapMode.READ, 0, 4 * numPixels);
    pickMapElaspedTime = performance.now() - startTime;
    const ids = new Uint32Array(pickBuffer.getMappedRange(0, 4 * numPixels));
    const id = ids[0];
    pickBuffer.unmap();
    return id;
  }

  async function processPicking() {
    if (!waitingForPreviousResults) {
      waitingForPreviousResults = true;
      const pixelX = mouseX * canvas.width / canvas.clientWidth;
      const pixelY = mouseY * canvas.height / canvas.clientHeight;
      const id = await drawWithIdsForMaterialsAndGetPixelUnderMouse(pixelX, pixelY);
      if (oldPickNdx >= 0) {
        const object = objects[oldPickNdx];
        object.uniforms.colorMult.set(oldPickColor);
        updateUniformBuffer(device, object.fs);
        oldPickNdx = -1;
      }

      if (id > objects.length) {
        console.error(`id > numObject: ${id}`);
      } else if (id > 0) {
        const pickNdx = id - 1;
        oldPickNdx = pickNdx;
        const object = objects[pickNdx];
        oldPickColor.set(object.uniforms.colorMult);
        object.uniforms.colorMult.set([1, 0, 0, 1]);
        updateUniformBuffer(device, object.fs);
      }

      infoElem.textContent = `\
obj#: ${id ? id : 'none'}
mapAsyncTime: ${(pickMapElaspedTime / 1000).toFixed(3)}ms
`;

      waitingForPreviousResults = false;
    }
  }

  let requestId;
  function requestRender() {
    if (!requestId) {
      requestId = requestAnimationFrame(render);
    }
  }

  window.addEventListener('mousemove', (e) => {
    const rect = canvas.getBoundingClientRect();
    mouseX = e.clientX - rect.left;
    mouseY = e.clientY - rect.top;
  });

  const stats = new Stats();
  document.body.appendChild(stats.dom);

  const moveElem = document.querySelector('#move');

  let then = 0;
  function render(time) {
    time *= 0.001;
  	stats.begin();
    const deltaTime = time - then;
    then = time;
    requestId = undefined;
    resizeToDisplaySize(device);

    const aspect = canvas.clientWidth / canvas.clientHeight;

    // we placed the circles in +/- 100 so adjust the projection so
    // that area fills the canvas
    const s = aspect > 1 ? 100 / aspect : 100;
    const projection = mat4.ortho(-aspect * s, aspect * s, -s, s, 1, 2000);
    const eye = [0, 0, 100];
    const target = [0, 0, 0];
    const up = [0, 1, 0];

    const view = mat4.lookAt(eye, target, up);
    const viewProjection = mat4.multiply(projection, view);

    const move = moveElem.checked;

    const h = s + 20;
    const w = s * aspect + 20;

    // Compute the matrices for each object.
    for (const object of objects) {
      const { translation, velocity } = object.v;
      if (move) {
        translation[0] = (translation[0] + w + velocity[0] * deltaTime) % (w * 2) - w;
        translation[1] = (translation[1] + h + velocity[1] * deltaTime) % (h * 2) - h;
      }
      computeMatrix(
          viewProjection,
          object.v.translation,
          object.v.xRotationSpeed * time * 0,
          object.v.yRotationSpeed * time * 0,
          object.uniforms.worldViewProjection);
      updateUniformBuffer(device, object.vs);
    }

    const colorTexture = context.getCurrentTexture();
    renderPassDescriptor.colorAttachments[0].view = colorTexture.createView();
    renderPassDescriptor.depthStencilAttachment.view = depthTextureView;

    const commandEncoder = device.createCommandEncoder();
    const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
    passEncoder.setPipeline(pipeline);
    for (const object of objects) {
      passEncoder.setBindGroup(0, object.bindGroup);
      passEncoder.setVertexBuffer(0, object.bufferInfo.buffers.position);
      passEncoder.setVertexBuffer(1, object.bufferInfo.buffers.normal);
      passEncoder.draw(object.bufferInfo.numVerts);
    }
    passEncoder.end();
    device.queue.submit([commandEncoder.finish()]);

    processPicking();
    requestRender();
    stats.end();
  }
  requestRender();
}

function fail(msg) {
  const elem = document.querySelector('#fail');
  const contentElem = elem.querySelector('.content');
  elem.style.display = '';
  contentElem.textContent = msg;
}

main();

jsGist.json

{"name":"WebGPU Picking","settings":{},"filenames":["index.html","index.css","index.js"]}