Last active
April 1, 2024 23:31
-
-
Save greggman/e20f28aab7427c68dc4f427c6dc32654 to your computer and use it in GitHub Desktop.
WebGPU Picking (pickBuffer = size of canvas)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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; | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
<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> |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// Note: This example is inefficent | |
// 1. There is no reason to copy the texture texture to a buffer. | |
// Instead, we can just copy the pixel under the mouse to a buffer. | |
// | |
// See this example: https://jsgist.org/?src=f5cb75bd70e715415683abcbf5679c44 | |
// | |
// 2. There is no reason to draw the entire scene to a canvas size texture | |
// for picking. Maybe if you want N point picking (like 10 fingers on mobile) | |
// but for single point picking we can just adjust the projection math and | |
// draw to a single pixel texture. | |
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', | |
}, | |
}; | |
let pickBuffer; | |
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(); | |
const bytesPerRow = ((canvas.width * 4 + 255) / 256 | 0) * 256; | |
const sizeNeeded = canvas.height * bytesPerRow; | |
if (!pickBuffer || pickBuffer.size !== sizeNeeded) { | |
if (pickBuffer) { | |
pickBuffer.destroy(); | |
} | |
pickBuffer = device.createBuffer({ | |
size: sizeNeeded, | |
usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ, | |
}); | |
} | |
commandEncoder.copyTextureToBuffer({ | |
texture: pickTexture, | |
}, { | |
buffer: pickBuffer, | |
bytesPerRow, | |
rowsPerImage: canvas.height, | |
}, { | |
width: canvas.width, | |
height: canvas.height, | |
// depth: 1, | |
}); | |
device.queue.submit([commandEncoder.finish()]); | |
const startTime = performance.now(); | |
await pickBuffer.mapAsync(GPUMapMode.READ); | |
const ids = new Uint32Array(pickBuffer.getMappedRange()); | |
pickMapElaspedTime = performance.now() - startTime; | |
const offset = pixelY * bytesPerRow / 4 + pixelX; | |
const id = ids[offset]; | |
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(); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{"name":"WebGPU Picking (pickBuffer = size of canvas)","settings":{},"filenames":["index.html","index.css","index.js"]} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment