Jump flood
A territory map of 48 drifting towns. Ten dispatches rebuild their regions each frame in one compute pass, each with its own step size, using immediates when the device supports them and a bind group per dispatch otherwise. One more dispatch records which regions touch, so the draw measures each pixel's exact distance to its region's edges, and a 'use gpu' paint function inks smooth borders, colour bands and town marks on paper.
Files, in reading order
- scene.tsCreates the scene: its resources, its settings and what each frame records.
- jump-flood.tsFinds nearest seeds by jump flooding.
- map.tsA territory map on paper, drawn as an atlas colours its countries, for a grid of
sizecells andseedstowns.
Kit modules owner. Also @typegpu/noise, typegpu.
scene.ts
import { owner, type Frame, type Surface } from '@typegpu/kit';import { d, type TgpuRoot } from 'typegpu';import type { Settings } from '../settings.ts';import { createJumpFlood } from './jump-flood.ts';import { createMap } from './map.ts';/** Cells along each side of the grid. */const SIZE = 512;const SEEDS = 48;export interface SceneOptions { /** Uses immediates for dispatch values. Defaults to device support. */ readonly immediates?: boolean;}/** * A territory map whose regions are the Voronoi cells of drifting towns, rebuilt each frame by * jump flooding in one compute pass. */export async function createScene(root: TgpuRoot, surface: Surface, options: SceneOptions = {}) { const time = root.createUniform(d.f32); const own = owner('The jump flood was destroyed.'); const map = createMap(root, SIZE, SEEDS); own(map.destroy); const build = async (immediates?: boolean) => { const built = createJumpFlood(root, { size: SIZE, seeds: SEEDS, format: surface.format, time, target: surface, paint: map.paint, ...(immediates === undefined ? {} : { immediates }), }); await built.initAsync(); return built; }; let flood = await build(options.immediates); const values = { passes: flood.passes, immediates: flood.immediates }; const supported = root.enabledWgslLanguageFeatures.has('immediate_address_space'); own(() => flood.destroy()); const settings: Settings<typeof values> = { values, knobs: { passes: { label: 'flood dispatches', min: 0, max: flood.passes, step: 1 }, immediates: supported ? {} : { disabled: 'unsupported' }, }, /** Changing `immediates` rebuilds the flood. Frames read `passes` directly. */ async write(key) { if (key === 'passes' || values.immediates === flood.immediates) return; const next = await build(values.immediates); if (own.destroyed || next.immediates !== values.immediates) return next.destroy(); flood.destroy(); flood = next; }, }; return { settings, frame(frame: Frame) { time.write(frame.time); const compute = frame.encoder.beginComputePass(); flood.flood(compute, values.passes); compute.end(); const output = frame.encoder.beginRenderPass({ colorAttachments: { view: surface.view(frame) }, }); flood.draw(output); output.end(); }, destroy: own.destroy, };}jump-flood.ts
import { owner, type SizeSource } from '@typegpu/kit';import { randf } from '@typegpu/noise';import { common, d, std, tgpu, type TgpuComputePass, type TgpuRenderCommands, type TgpuRoot, type TgpuUniform,} from 'typegpu';const WORKGROUP = 8;const SEED_WORKGROUP = 64;const TAU = Math.PI * 2;/** * Seed circles have centres anywhere on the grid, radii in `CIRCLES` and speeds up to `DRIFT` * radians per second in either direction. Radii use fractions of grid size. */const CIRCLES = [0.02, 0.1] as const;const DRIFT = 0.12;/** Right, down and both diagonals below, which meet each touching pair of cells once. */const TOUCHING = tgpu.const(d.arrayOf(d.vec2i, 4), [ d.vec2i(1, 0), d.vec2i(-1, 1), d.vec2i(0, 1), d.vec2i(1, 1),]);/** Sentinel for an unreached cell. */export const NONE = 0xffffffff;/** Dispatch step distance and input-buffer half. */const Pass = d.struct({ step: d.u32, source: d.u32 });const Seed = d.struct({ position: d.vec2f });/** What the draw knows of a pixel. Distances and positions are in cells. */export const Cell = d.struct({ /** The pixel's position on the grid. */ point: d.vec2f, /** The nearest seed's index, or `NONE` where no seed has reached yet. */ seed: d.u32, /** The nearest seed's position on the grid. */ site: d.vec2f, /** Distance to the nearest seed. */ toSite: d.f32, /** The second-nearest seed's index, or `NONE`. */ second: d.u32, /** Distance to the second-nearest seed. */ toSecond: d.f32, /** Distance to the nearest edge of the seed's region. */ toEdge: d.f32, /** Cells one pixel spans, for anti-aliasing. */ footprint: d.f32,});/** Colours a pixel of the diagram. Runs on the GPU. */export type Paint = (cell: d.Infer<typeof Cell>) => d.v3f;const View = d.struct({ viewport: d.vec2f, result: d.u32 });export interface JumpFloodOptions { /** Cells along each side of the square grid. A power of two. */ readonly size: number; readonly seeds: number; readonly format: GPUTextureFormat; /** Frame-clock time in seconds for seed drift. */ readonly time: TgpuUniform<d.F32>; /** What the diagram is drawn into, such as the canvas surface. */ readonly target: SizeSource; readonly paint: Paint; /** Uses immediates for dispatch values. Defaults to device support. */ readonly immediates?: boolean;}/** * Finds nearest seeds by jump flooding. Each dispatch checks eight cells at the current step * distance, then the distance halves. Owns buffers and records into caller passes. */export function createJumpFlood( root: TgpuRoot, { size, seeds: count, format, time, target, paint, immediates = root.enabledWgslLanguageFeatures.has('immediate_address_space'), }: JumpFloodOptions,) { // Repeat step 1 after halving to correct most remaining errors const steps: number[] = []; for (let step = size / 2; step >= 1; step /= 2) steps.push(step); steps.push(1); const cellCount = size * size; const cellGroups = Math.ceil(size / WORKGROUP); const words = Math.ceil(count / 32); const own = owner(); // Read and write separate halves so dispatches cannot race with their own reads const cells = own(root.createBuffer(d.arrayOf(d.u32, cellCount * 2)).$usage('storage')); const seeds = own(root.createBuffer(d.arrayOf(Seed, count)).$usage('storage')); const [writableCells, readableCells] = [cells.as('mutable'), cells.as('readonly')]; const [writableSeeds, readableSeeds] = [seeds.as('mutable'), seeds.as('readonly')]; // A row of bits per seed. The flood's result sets the bits of seeds whose regions touch, so // the draw measures exact distances to a region's edges. Fragments read it atomically too const touching = own( root.createBuffer(d.arrayOf(d.atomic(d.u32), count * words)).$usage('storage'), ); const neighbours = touching.as('mutable'); const view = root.createUniform(View, { viewport: [target.width, target.height], result: 0 }); // Queue writes cannot change uniforms between dispatches in one compute pass. Set an // immediate on the pass, or bind a separate uniform per dispatch. `current` reads either // source and defaults to the bind-group entry const layout = tgpu.bindGroupLayout({ pass: { uniform: Pass } }); const current = tgpu.accessor(Pass, () => { 'use gpu'; return Pass(layout.$.pass); }); const immediate = immediates ? tgpu['~unstable'].immediateVar(Pass) : undefined; const stepValues = steps.map((step, index) => new Uint32Array([step, index % 2])); const uniforms = immediate ? [] : steps.map((step, index) => root.createBuffer(Pass, { step, source: index % 2 }).$usage('uniform'), ); const groups = uniforms.map((uniform) => root.createBindGroup(layout, { pass: uniform })); const clearKernel = tgpu.computeFn({ workgroupSize: [WORKGROUP, WORKGROUP], in: { gid: d.builtin.globalInvocationId }, })(({ gid }) => { 'use gpu'; if (gid.x >= size || gid.y >= size) return; const index = gid.y * size + gid.x; writableCells.$[index] = NONE; if (index < count * words) std.atomicStore(neighbours.$[index]!, 0); }); const placeKernel = tgpu.computeFn({ workgroupSize: [SEED_WORKGROUP], in: { gid: d.builtin.globalInvocationId }, })(({ gid }) => { 'use gpu'; if (gid.x >= count) return; randf.seed2(d.vec2f(d.f32(gid.x), 0.5)); const center = d.vec2f(randf.sample(), randf.sample()); const radius = std.mix(CIRCLES[0], CIRCLES[1], randf.sample()); const speed = (randf.sample() * 2 - 1) * DRIFT; const angle = randf.sample() * TAU + time.$ * speed; const circling = center + d.vec2f(std.cos(angle), std.sin(angle)) * radius; const position = std.clamp(circling, d.vec2f(0), d.vec2f(1 - 0.5 / size)); writableSeeds.$[gid.x] = Seed({ position }); const cell = d.vec2u(position * size); writableCells.$[cell.y * size + cell.x] = gid.x; }); const floodKernel = tgpu.computeFn({ workgroupSize: [WORKGROUP, WORKGROUP], in: { gid: d.builtin.globalInvocationId }, })(({ gid }) => { 'use gpu'; if (gid.x >= size || gid.y >= size) return; const pass = current.$; const here = d.vec2f(gid.xy) + 0.5; let nearest = d.u32(NONE); let shortest = d.f32(size * 2); for (let y = -1; y <= 1; y++) { for (let x = -1; x <= 1; x++) { const at = d.vec2i(gid.xy) + d.vec2i(x, y) * d.i32(pass.step); if (at.x >= 0 && at.y >= 0 && at.x < size && at.y < size) { // A dispatch may not bind one buffer both read-only and read-write, so the source half // is read through the mutable binding too const seed = writableCells.$[d.u32(at.y) * size + d.u32(at.x) + pass.source * cellCount]; if (seed !== NONE) { const distance = std.distance(readableSeeds.$[seed].position * size, here); if (distance < shortest) { shortest = distance; nearest = seed; } } } } } writableCells.$[gid.y * size + gid.x + (1 - pass.source) * cellCount] = nearest; }); const link = (from: number, to: number) => { 'use gpu'; const word = from * words + (to >>> 5); const bit = d.u32(1) << (to & 31); if ((std.atomicLoad(neighbours.$[word]!) & bit) === 0) std.atomicOr(neighbours.$[word]!, bit); }; const neighboursKernel = tgpu.computeFn({ workgroupSize: [WORKGROUP, WORKGROUP], in: { gid: d.builtin.globalInvocationId }, })(({ gid }) => { 'use gpu'; if (gid.x >= size || gid.y >= size) return; const half = view.$.result * cellCount; const seed = readableCells.$[gid.y * size + gid.x + half]; if (seed === NONE) return; for (let i = 0; i < 4; i++) { const at = d.vec2i(gid.xy) + TOUCHING.$[i]!; if (at.x >= 0 && at.x < size && at.y < size) { const other = readableCells.$[d.u32(at.y) * size + d.u32(at.x) + half]; if (other !== NONE && other !== seed) { link(seed, other); link(other, seed); } } } }); const fragment = tgpu.fragmentFn({ in: { uv: d.vec2f }, out: d.vec4f })(({ uv }) => { 'use gpu'; // Fit the grid across the target's longer side const scale = view.$.viewport / std.max(view.$.viewport.x, view.$.viewport.y); const point = ((uv - 0.5) * scale + 0.5) * size; const cell = d.vec2u(std.clamp(point, d.vec2f(0), d.vec2f(size - 1))); const footprint = size / std.max(view.$.viewport.x, view.$.viewport.y); const half = view.$.result * cellCount; let seed = readableCells.$[cell.y * size + cell.x + half]; if (seed === NONE) { const unreached = Cell({ point, seed, site: d.vec2f(), toSite: 0, second: seed, toSecond: 0, toEdge: 0, footprint, }); return d.vec4f(paint(unreached), 1); } // Each neighbour bounds the distance to the edge by its bisector. Near an edge a neighbour // may be nearer to the pixel than the cell's seed; the scan then repeats from that neighbour let site = d.vec2f(readableSeeds.$[seed].position * size); let toSite = std.distance(site, point); let toEdge = d.f32(size); let second = d.u32(NONE); let toSecond = d.f32(size * 2); for (let scan = 0; scan < 2; scan++) { toEdge = d.f32(size); second = d.u32(NONE); toSecond = d.f32(size * 2); let secondSite = d.vec2f(); for (let word = d.u32(0); word < words; word++) { let bits = std.atomicLoad(neighbours.$[seed * words + word]!); while (bits !== 0) { const other = word * 32 + std.countTrailingZeros(bits); bits &= bits - 1; const there = readableSeeds.$[other].position * size; const across = std.normalize(there - site); toEdge = std.min(toEdge, std.dot((site + there) * 0.5 - point, across)); const toThere = std.distance(there, point); if (toThere < toSecond) { second = other; toSecond = toThere; secondSite = d.vec2f(there); } } } if (scan > 0 || toSecond >= toSite) break; seed = second; site = d.vec2f(secondSite); toSite = toSecond; } return d.vec4f( paint(Cell({ point, seed, site, toSite, second, toSecond, toEdge, footprint })), 1, ); }); const clearPipeline = root.createComputePipeline({ compute: clearKernel }); const placePipeline = root.createComputePipeline({ compute: placeKernel }); const neighboursPipeline = root.createComputePipeline({ compute: neighboursKernel }); const floodPipeline = (immediate ? root.with(current, immediate) : root).createComputePipeline({ compute: floodKernel, }); const drawPipeline = root.createRenderPipeline({ vertex: common.fullScreenTriangle, fragment, targets: { format }, }); let width = target.width; let height = target.height; let result = 0; return { /** Nearest seed per cell, or `NONE`. Each half contains `size * size` cells. */ cells, seeds, /** Bit rows of touching regions, `ceil(seeds / 32)` words per seed. */ neighbours: touching, /** Whether dispatch values use immediates. */ immediates: immediate !== undefined, /** Maximum flood-dispatch count. */ passes: steps.length, initAsync: () => Promise.all( [clearPipeline, placePipeline, floodPipeline, neighboursPipeline, drawPipeline].map((it) => it.initAsync(), ), ), /** Places seeds, records the first `passes` flood dispatches and finds touching regions. */ flood(pass: TgpuComputePass, passes = steps.length) { pass.setPipeline(clearPipeline); pass.dispatchWorkgroups(cellGroups, cellGroups); pass.setPipeline(placePipeline); pass.dispatchWorkgroups(Math.ceil(count / SEED_WORKGROUP)); pass.setPipeline(floodPipeline); for (let index = 0; index < passes; index++) { if (immediate) pass.setImmediates(immediate, stepValues[index]!); else pass.setBindGroup(layout, groups[index]!); pass.dispatchWorkgroups(cellGroups, cellGroups); } // The draw reads the half the last dispatch wrote if (passes % 2 !== result) { result = passes % 2; view.patch({ result }); } pass.setPipeline(neighboursPipeline); pass.dispatchWorkgroups(cellGroups, cellGroups); }, draw(pass: TgpuRenderCommands) { if (target.width !== width || target.height !== height) { width = target.width; height = target.height; view.patch({ viewport: d.vec2f(width, height) }); } pass.setPipeline(drawPipeline); pass.draw(3); }, destroy: own.destroy, [Symbol.dispose]: own.destroy, };}map.ts
import { perlin2d, randf } from '@typegpu/noise';import { d, std, tgpu, type TgpuRoot } from 'typegpu';import { Cell, NONE } from './jump-flood.ts';/** Display colours. */const PAPER = d.vec3f(0.96, 0.91, 0.8);const INK = d.vec3f(0.22, 0.17, 0.14);/** Washes: sage, ochre, rose, slate, lilac and teal. */const WASHES = tgpu.const(d.arrayOf(d.vec3f, 6), [ d.vec3f(0.46, 0.66, 0.38), d.vec3f(0.94, 0.66, 0.26), d.vec3f(0.9, 0.5, 0.42), d.vec3f(0.36, 0.62, 0.88), d.vec3f(0.64, 0.52, 0.78), d.vec3f(0.34, 0.66, 0.62),]);/** Distances in cells: border half-width and the colour band inside it. */const BORDER = 0.26;const BAND = 16;/** Noise frequencies per cell for the paper and the washes, and how far apart each wash samples. */const MOTTLE = 0.008;const UNEVEN = { every: 0.02, apart: 7.31 };/** Town marks in cells. Capitals carry a ring around their dot. */const TOWN = { dot: 1.15, ring: 2.5, stroke: 0.24, clear: 0.8 };/** How much of a pixel `distance` cells from a shape lies within `radius` of it. */const cover = (distance: number, radius: number, footprint: number) => { 'use gpu'; return std.saturate((radius - distance) / footprint + 0.5);};/** 1 for a capital, 0 for a town. */const capital = (seed: number) => { 'use gpu'; randf.seed(d.f32(seed) + 0.5); randf.sample(); return std.select(d.f32(0), d.f32(1), randf.sample() < 0.3);};/** Ink coverage and clearing of the mark of a town `toSite` cells away. */const mark = (toSite: number, seed: number, footprint: number) => { 'use gpu'; if (toSite > TOWN.ring + TOWN.clear + footprint) return d.vec2f(); const big = capital(seed); const ring = TOWN.ring * big; const dot = cover(toSite, TOWN.dot, footprint); const circle = cover(std.abs(toSite - ring), TOWN.stroke, footprint) * big; const clear = cover(toSite, std.max(ring, TOWN.dot) + TOWN.clear, footprint); return d.vec2f(std.max(dot, circle), clear);};/** Paints a pixel of the map, for a grid of `size` cells. */const mapOf = (size: number) => (cell: d.Infer<typeof Cell>) => { 'use gpu'; const f = cell.footprint; // Paper with a broad mottle, lit from the top left and darker towards its edges const uv = cell.point / size; const mottle = perlin2d.sample(cell.point * MOTTLE) * 0.035; const fromMiddle = std.length((uv - 0.5) * d.vec2f(1, 1.4)); const light = 1 + 0.04 * (0.5 - uv.x) + 0.04 * (0.5 - uv.y); const paper = PAPER * (1 + mottle) * light * (1 - fromMiddle * fromMiddle * 0.28); if (cell.seed === NONE) return paper; randf.seed(d.f32(cell.seed) + 0.5); const wash = WASHES.$[d.u32(randf.sample() * 6) % 6]!; // A pale fill, deepening into a band along the border, unevenly as a wash dries const band = std.pow(1 - std.saturate(cell.toEdge / BAND), 3); const uneven = perlin2d.sample(cell.point * UNEVEN.every + d.f32(cell.seed) * UNEVEN.apart) * 0.07; const strength = std.saturate(0.3 + band * 0.42 + uneven); let color = paper * std.mix(d.vec3f(1), wash, strength); color = std.mix(color, INK, cover(cell.toEdge, std.max(BORDER, f * 0.5), f) * 0.9); // Clear the paper around towns, then ink their marks, the second-nearest's too where marks meet const own = mark(cell.toSite, cell.seed, f); let ink = own.x; let clear = own.y; if (cell.second !== NONE) { const other = mark(cell.toSecond, cell.second, f); ink = std.max(ink, other.x); clear = std.max(clear, other.y); } color = std.mix(color, paper * std.mix(d.vec3f(1), wash, 0.12), clear * 0.8); return std.mix(color, INK, ink * 0.95);};/** * A territory map on paper, drawn as an atlas colours its countries, for a grid of `size` cells * and `seeds` towns. Owns the noise gradients its paint reads. */export function createMap(root: TgpuRoot, size: number, seeds: number) { // Gradients for every lattice point the paper and the washes sample, so none wraps const reach = Math.ceil(size * UNEVEN.every + (seeds - 1) * UNEVEN.apart) + 2; const gradients = perlin2d.staticCache({ root, size: d.vec2u(reach, reach) }); const paint = tgpu .fn( [Cell], d.vec3f, )(mapOf(size)) .with(perlin2d.getJunctionGradientSlot, gradients.getJunctionGradient); return { paint, destroy: () => gradients.destroy() };}