33:T7e1,This notebook contains my first real attempt at implementing something nontrivial in WebGPU. It is a direct port of [Line Sweep Ambient Occlusion in JavaScript](https://observablehq.com/@rreusser/line-sweep-ambient-occlusion) to a WebGPU compute shader. The original algorithm iterates over all scanlines, first left, then right, then up, then down (and maybe diagonal). The nice thing about the [Line-Sweep Ambient Occlusion algorithm](https://karim.naaji.fr/lsao.html) is that for each direction's pass, we track horizon visibility on a stack and shade as we go so that we only need to touch each pixel exactly once. On the CPU we process these scanlines one at a time. We should be able to parallelize that on the GPU, but years of working with WebGL has taught me that while GPUs are good at evaluating pixels independently, they're not a good fit for this type of sweep algorithm where each pixel depends on all previous pixels. But hey, WebGPU has arrived and compute shaders play by different rules, so let's go for it. The implementation below dispatches workgroups of ${workgroupSize} invocations at a time. Each invocation implements a single scanline in a single direction. That means within a single shader invocation, we march all the way across an entire scanline, sampling the terrain data 512 times and shading all 512 pixels of the row or column. I suspect it's a moderately suboptimal for reasons I don't understand very well, but my hasty benchmarks indicate it's about 20x faster than [the CPU version of the algorithm](https://observablehq.com/@rreusser/line-sweep-ambient-occlusion) and, upon ramping it up to a hundred or so (redundant) passes and taking an average, each scanline processes in less than 0.5ms. Overall, I'm very happy with the success of translation to WebGPU, and I'm moderately happy with the performance boost. It's not up to feature parity with [Line-Sweep Terrain Lighting](https://observablehq.com/@rreusser/line-sweep-terrain-lighting), but I'd like to get it there.34:T864,computeShader = createShaderModule(device, { code: ` override tileSize = 512u; override tileBuffer = 1u; override workgroupSize = 128; struct UniformStruct { tilesize: vec2, step: vec2, buffer: i32, pixelSize: f32, normalization: f32 } @binding(0) @group(0) var uniforms: UniformStruct; @binding(1) @group(0) var terrainData: array; @binding(2) @group(0) var outputData: array; fn unbufferedIndex(ij: vec2) -> u32 { return (u32(ij.x) % tileSize) + u32(ij.y) * tileSize; } fn bufferedIndex(ij: vec2) -> u32 { let w = tileSize + 2u * tileBuffer; let ijbuf = vec2(ij + i32(tileBuffer)); return (ijbuf.x % w) + (ijbuf.y) * w; } @compute @workgroup_size(workgroupSize) fn main( @builtin(global_invocation_id) coord: vec3u) { var ij = vec2(i32(coord.x), i32(coord.x)); var dijz: vec3; if (uniforms.step.y == 0) { ij.x = select(0i, i32(tileSize - 1), uniforms.step.x < 0i); } else if (uniforms.step.x == 0) { ij.y = select(0i, i32(tileSize - 1), uniforms.step.y < 0i); } var hull: array, 64>; hull[0] = vec3(vec2(ij - uniforms.step), terrainData[bufferedIndex(ij - uniforms.step)]); var hullPtr = 0u; for (var i = 0u; i < tileSize; i = i + 1u) { let uidx = unbufferedIndex(ij); let bidx = bufferedIndex(ij); let z = terrainData[bidx]; let ijz = vec3(vec2(ij), z); dijz = hull[hullPtr] - ijz; var s0 = dijz.z * dijz.z / dot(dijz, dijz); s0 = select(-s0, s0, dijz.z > 0.0); while(hullPtr > 0) { dijz = hull[hullPtr - 1] - ijz; var s1 = dijz.z * dijz.z / dot(dijz, dijz); s1 = select(-s1, s1, dijz.z > 0.0); if (s0 > s1) { break; } s0 = s1; hullPtr -= 1u; } dijz = hull[hullPtr] - ijz; dijz = vec3(dijz.xy, dijz.z / uniforms.pixelSize); outputData[uidx] = outputData[uidx] + uniforms.normalization * exp(-dijz.z / length(dijz)); // Fail silently but gracefully if we overflow the stack hullPtr = min(hullPtr + 1u, 63u); hull[hullPtr] = ijz; ij = ij + uniforms.step; } }` })35:T4c1,renderToAO = { rerun; dirs.forEach((dir, i) => device.queue.writeBuffer( uniformBuffer, i * uniformSizePerDispatch, packUniforms( [tile.tileSize, tile.tileSize], dir, tile.buffer, pixelSize, 1 / dirs.length ) ) ); const encoder = device.createCommandEncoder(); // Clear the buffer which then accumulates the contribution from all directions encoder.clearBuffer(computationBuffer); const opts = {}; if (supportsTimestamps) opts.timestampWrites = { querySet, beginningOfPassWriteIndex: 0, endOfPassWriteIndex: 1 }; const pass = encoder.beginComputePass(opts); pass.setPipeline(computePipeline); // Invoke all four directions in one compute pass for (let i = 0; i < dirs.length; i++) { pass.setBindGroup(0, bindGroup, [i * uniformSizePerDispatch]); pass.dispatchWorkgroups(Math.floor(tile.tileSize / workgroupSize)); } pass.end(); if (supportsTimestamps) { // Measure execution time encoder.resolveQuerySet(querySet, 0, 2, timestampBuffer, 0); encoder.copyBufferToBuffer(timestampBuffer, 0, readbackBuffer, 0, 2 * 8); } device.queue.submit([encoder.finish()]); }36:[] 31:{"title":"Line-Sweep Ambient Occlusion with WebGPU","cells":[{"id":0,"value":"# Line-Sweep Ambient Occlusion with WebGPU","pinned":false,"hidden":false,"mode":"md"},{"id":1094,"value":"$33","pinned":false,"hidden":false,"mode":"md"},{"id":5,"value":"canvas = {\n const el = DOM.canvas(shape[0], shape[1]);\n const maxSize = Math.min(width, shape[0] / 1);\n el.style.width = `${Math.min(maxSize, shape[0] / 1)}px`;\n el.style.height = `${Math.min(\n (maxSize * shape[1]) / shape[0],\n shape[1] / 1\n )}px`;\n return el;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":1285,"value":"viewof brightness = Inputs.range([-1, 1], { label: \"Brightness\" })","pinned":false,"hidden":false,"mode":"ojs"},{"id":1121,"value":"viewof contrast = Inputs.range([0, 1], { label: \"Contrast\" })","pinned":false,"hidden":false,"mode":"ojs"},{"id":193,"value":"viewof workgroupSize = Inputs.select(\n [...Array(9).keys()].map((i) => 2 ** i),\n { value: 128, label: \"Workgroup size\" }\n)","pinned":false,"hidden":false,"mode":"ojs"},{"id":897,"value":"{\n if (!supportsTimestamps) return html``;\n renderToAO;\n await readbackBuffer.mapAsync(GPUMapMode.READ);\n const timestamps = new BigUint64Array(readbackBuffer.getMappedRange());\n const start = timestamps[0];\n const end = timestamps[1];\n readbackBuffer.unmap();\n\n // Query the timestamp period in nanoseconds\n const timestampPeriod = adapter.limits.timestampPeriod ?? 1;\n\n const durationNs = Number(end - start) * timestampPeriod;\n return html`GPU execution time (${\n dirs.length\n } direction passes) = ${(\n durationNs / 1e6\n ).toFixed(1)} ms`;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":1232,"value":"viewof rerun = Inputs.button('Re-run')","pinned":false,"hidden":false,"mode":"ojs"},{"id":136,"value":"## Get terrain data","pinned":false,"hidden":false,"mode":"md"},{"id":126,"value":"function getTerrainTile({ x, y, z }) {\n const ACCESS_TOKEN_BUT_PLS_USE_YOUR_OWN_AND_NOT_MINE =\n \"pk.eyJ1IjoicnNyZXVzc2VyIiwiYSI6ImNtMjZmejNqczBmYzIya3B6ZjBzZnYwOGgifQ.CWSOe_wStIypz4_snROphQ\";\n const SKU = \"101pM5uZK0TZt\";\n return new Promise((resolve, reject) => {\n const img = new Image();\n img.src = `https://api.mapbox.com/raster/v1/mapbox.mapbox-terrain-dem-v1/${z}/${x}/${y}.webp?sku=${SKU}&access_token=${ACCESS_TOKEN_BUT_PLS_USE_YOUR_OWN_AND_NOT_MINE}`;\n img.crossOrigin = \"Anonymous\";\n img.onload = () => {\n const tileSize = Math.pow(2, Math.floor(Math.log2(img.width)));\n const buffer = (img.width - tileSize) / 2;\n resolve({ img, width: img.width, height: img.height, tileSize, buffer });\n };\n img.onerror = () => reject(\"Failed to load terrain tile\");\n });\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":146,"value":"function readImageData(img) {\n const canvas = DOM.element(\"canvas\");\n canvas.width = img.width;\n canvas.height = img.height;\n const ctx = canvas.getContext(\"2d\");\n ctx.drawImage(img, 0, 0);\n return ctx.getImageData(0, 0, img.width, img.height).data;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":140,"value":"function decodeTerrainData(data) {\n const out = new Float32Array(data.length / 4);\n for (let i = 0; i < data.length; i += 4) {\n out[i >> 2] =\n -10000 + 0.1 * ((data[i] << 16) | (data[i + 1] << 8) | data[i + 2]);\n }\n return out;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":158,"value":"terrainData = decodeTerrainData(readImageData(tile.img))","pinned":false,"hidden":false,"mode":"ojs"},{"id":129,"value":"xyz = ({ x: 164, y: 363, z: 10 })","pinned":false,"hidden":false,"mode":"ojs"},{"id":132,"value":"tile = getTerrainTile(xyz)","pinned":false,"hidden":false,"mode":"ojs"},{"id":628,"value":"MAXEXTENT = 20037508.342789244","pinned":false,"hidden":false,"mode":"ojs"},{"id":627,"value":"pixelSize = (2 * MAXEXTENT) / tile.tileSize / Math.pow(2, xyz.z)","pinned":false,"hidden":false,"mode":"ojs"},{"id":162,"value":"## Create a WebGPU context","pinned":false,"hidden":false,"mode":"md"},{"id":31,"value":"shape = [512, 512]","pinned":false,"hidden":false,"mode":"ojs"},{"id":10,"value":"adapter = navigator.gpu?.requestAdapter()","pinned":false,"hidden":false,"mode":"ojs"},{"id":1206,"value":"supportsTimestamps = adapter.features.has(\"timestamp-query\")","pinned":false,"hidden":false,"mode":"ojs"},{"id":15,"value":"device = adapter?.requestDevice({\n requiredFeatures: supportsTimestamps ? [\"timestamp-query\"] : []\n})","pinned":false,"hidden":false,"mode":"ojs"},{"id":18,"value":"context = {\n const ctx = canvas.getContext(\"webgpu\");\n ctx.configure({ device, format });\n return ctx;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":969,"value":"format = navigator.gpu.getPreferredCanvasFormat()","pinned":false,"hidden":false,"mode":"ojs"},{"id":313,"value":"import { createShaderModule } from \"@rreusser/webgpu-utils\"","pinned":false,"hidden":false,"mode":"ojs"},{"id":168,"value":"## Ambient occlusion algorithm","pinned":false,"hidden":false,"mode":"md"},{"id":82,"value":"terrainDataBuffer = {\n const buffer = device.createBuffer({\n size: terrainData.byteLength,\n label: \"Terrain data buffer\",\n usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,\n mappedAtCreation: true\n });\n new Float32Array(buffer.getMappedRange()).set(terrainData);\n buffer.unmap();\n invalidation.then(() => buffer.destroy());\n return buffer;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":228,"value":"computationBuffer = {\n const buffer = device.createBuffer({\n size: tile.tileSize ** 2 * Uint32Array.BYTES_PER_ELEMENT,\n label: \"Computation buffer\",\n usage:\n GPUBufferUsage.STORAGE |\n GPUBufferUsage.COPY_SRC |\n GPUBufferUsage.COPY_DST,\n mappedAtCreation: true\n });\n new Uint32Array(buffer.getMappedRange()).fill(0);\n buffer.unmap();\n invalidation.then(() => buffer.destroy());\n return buffer;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":48,"value":"$34","pinned":true,"hidden":false,"mode":"ojs"},{"id":821,"value":"dirs = [\n [1, 0],\n [-1, 0],\n [0, 1],\n [0, -1]\n]","pinned":false,"hidden":false,"mode":"ojs"},{"id":816,"value":"uniformSizePerDispatch = 256","pinned":false,"hidden":false,"mode":"ojs"},{"id":299,"value":"uniformBuffer = {\n const buffer = device.createBuffer({\n size: uniformSizePerDispatch * dirs.length,\n usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST\n });\n invalidation.then(() => buffer.destroy());\n return buffer;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":405,"value":"function packUniforms(\n [width, height],\n [dx, dy],\n buffer,\n pixelSize,\n normalization\n) {\n const bytes = new Uint8Array(uniformSizePerDispatch);\n const i32 = new Int32Array(bytes.buffer);\n const u32 = new Uint32Array(bytes.buffer);\n const f32 = new Float32Array(bytes.buffer);\n u32[0] = width;\n u32[1] = height;\n i32[2] = dx;\n i32[3] = dy;\n u32[4] = buffer;\n f32[5] = pixelSize;\n f32[6] = normalization;\n return bytes;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":55,"value":"bindGroupLayout = device.createBindGroupLayout({\n entries: [\n {\n binding: 0,\n visibility: GPUShaderStage.COMPUTE,\n buffer: { type: \"uniform\", hasDynamicOffset: true }\n },\n {\n binding: 1,\n visibility: GPUShaderStage.COMPUTE,\n buffer: { type: \"read-only-storage\" }\n },\n {\n binding: 2,\n visibility: GPUShaderStage.COMPUTE,\n buffer: { type: \"storage\" }\n }\n ]\n})","pinned":false,"hidden":false,"mode":"ojs"},{"id":80,"value":"bindGroup = device.createBindGroup({\n layout: bindGroupLayout,\n entries: [\n {\n binding: 0,\n resource: {\n buffer: uniformBuffer,\n offset: 0,\n size: uniformSizePerDispatch\n }\n },\n { binding: 1, resource: { buffer: terrainDataBuffer } },\n { binding: 2, resource: { buffer: computationBuffer } }\n ]\n})","pinned":false,"hidden":false,"mode":"ojs"},{"id":72,"value":"computePipeline = device.createComputePipeline({\n layout: device.createPipelineLayout({\n bindGroupLayouts: [bindGroupLayout]\n }),\n compute: {\n module: computeShader,\n entryPoint: \"main\",\n constants: {\n workgroupSize,\n tileSize: tile.tileSize,\n tileBuffer: tile.buffer\n }\n }\n})","pinned":false,"hidden":false,"mode":"ojs"},{"id":856,"value":"## Measure execution time","pinned":false,"hidden":false,"mode":"md"},{"id":863,"value":"queryCount = 2","pinned":false,"hidden":false,"mode":"ojs"},{"id":857,"value":"querySet = {\n if (!supportsTimestamps) return null;\n const set = device.createQuerySet({\n type: \"timestamp\",\n count: queryCount\n });\n invalidation.then(() => set.destroy());\n return set;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":870,"value":"timestampBuffer = {\n const buffer = device.createBuffer({\n size: queryCount * 8, // 2 timestamps * 8 bytes\n usage: GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC\n });\n invalidation.then(() => buffer.destroy());\n return buffer;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":872,"value":"readbackBuffer = {\n const buffer = device.createBuffer({\n size: queryCount * 8,\n usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST\n });\n invalidation.then(() => buffer.destroy());\n return buffer;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":946,"value":"## Render to screen","pinned":false,"hidden":false,"mode":"md"},{"id":1111,"value":"rtsUniformBuffer = {\n const buffer = device.createBuffer({\n size: 32,\n usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST\n });\n invalidation.then(() => buffer.destroy());\n return buffer;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":959,"value":"rtsShader = createShaderModule(device, {\n code: `\nstruct Uniforms {\n contrast: f32,\n brightness: f32\n}\n\n@group(0) @binding(0) var values: array;\n@group(0) @binding(1) var uniforms: Uniforms;\n\nconst PI = ${Math.PI};\n\nfn getIndex(coord: vec2) -> u32 {\n return coord.y * 512u + coord.x;\n}\n\n@vertex\nfn vs_main(@builtin(vertex_index) idx: u32) -> @builtin(position) vec4 {\n var pos = array, 3>(\n vec2(-4.0, -4.0),\n vec2(4.0, -4.0),\n vec2(0.0, 4.0)\n );\n return vec4(pos[idx], 0.0, 1.0);\n}\n\n\n@fragment\nfn fs_main(@builtin(position) pos: vec4) -> @location(0) vec4 {\n let coord = vec2(pos.xy);\n var value = values[getIndex(coord)];\n value = 0.5 + 0.5 * tanh((value - 1.0) * uniforms.contrast + uniforms.brightness);\n value = pow(value, 0.454);\n return vec4(value, value, value, 1.0);\n}\n`\n})","pinned":false,"hidden":false,"mode":"ojs"},{"id":966,"value":"rtsPipeline = device.createRenderPipeline({\n layout: \"auto\",\n vertex: { module: rtsShader, entryPoint: \"vs_main\" },\n fragment: {\n module: rtsShader,\n entryPoint: \"fs_main\",\n targets: [{ format }]\n },\n primitive: { topology: \"triangle-list\" }\n})","pinned":false,"hidden":false,"mode":"ojs"},{"id":964,"value":"rtsBindGroup = device.createBindGroup({\n layout: rtsPipeline.getBindGroupLayout(0),\n entries: [\n { binding: 0, resource: { buffer: computationBuffer } },\n { binding: 1, resource: { buffer: rtsUniformBuffer } }\n ]\n})","pinned":false,"hidden":false,"mode":"ojs"},{"id":1081,"value":"## Final render","pinned":false,"hidden":false,"mode":"md"},{"id":70,"value":"$35","pinned":false,"hidden":false,"mode":"ojs"},{"id":978,"value":"renderToScreen = {\n renderToAO;\n const encoder = device.createCommandEncoder();\n\n device.queue.writeBuffer(\n rtsUniformBuffer,\n 0,\n new Float32Array([\n Math.min(100, (7 * contrast) / (1 - contrast)),\n Math.max(-30, Math.min(30, Math.tan((Math.PI / 2) * brightness)))\n ])\n );\n\n const pass = encoder.beginRenderPass({\n colorAttachments: [\n {\n view: context.getCurrentTexture().createView(),\n loadOp: \"clear\",\n clearValue: [0, 0, 0, 1],\n storeOp: \"store\"\n }\n ]\n });\n\n pass.setPipeline(rtsPipeline);\n pass.setBindGroup(0, rtsBindGroup);\n pass.draw(3);\n pass.end();\n\n device.queue.submit([encoder.finish()]);\n}","pinned":false,"hidden":false,"mode":"ojs"}],"files":"$W36","stdlib":"1","titling":"auto"}