Generating a Quadtree on the GPU with WGSL

A pointer-based quadtree — nodes holding four child references, grown by recursive insertion — cannot be built efficiently on a GPU, because insertion is serial and the tree shape is data-dependent. The parallel alternative is a linear quadtree: every point is assigned a Morton code (the bit-interleaving of its quantized x and y), the codes are sorted so spatially adjacent points land at adjacent indices, and nodes become nothing more than contiguous ranges of that sorted array whose codes share a common prefix. There is no tree structure to allocate and no recursion — construction is a quantize pass, a radix sort, and a range-finding pass, all of which parallelize across every point. This page gives the complete WGSL and a TypeScript driver for building that linear quadtree, so a million points are indexed in a handful of compute passes rather than a serial descent. It is the concrete tree-construction case of the build-scan-compact skeleton in WGSL spatial algorithms on the GPU.

Runnable reference implementation

The pipeline has three WGSL entry points. build_codes quantizes each point to a 32-bit Morton key and packs it with the point’s original index so the sort can carry the payload. histogram and scatter are one digit-pass of a least-significant-digit radix sort over those keys — repeated for each 8-bit digit, four passes cover a 32-bit code. mark_nodes then finds, for a chosen tree level, the boundaries where the Morton prefix changes, which delimit the quadtree cells at that level. The Morton part of the code reuses the standard bit-spreading trick.

// ---- Pass 1: quantize points to Morton keys, packed with payload index ----
struct KeyIndex { key: u32, idx: u32 };

@group(0) @binding(0) var<storage, read>       positions: array<vec2<f32>>;
@group(0) @binding(1) var<storage, read_write> keys:      array<KeyIndex>;
@group(0) @binding(2) var<uniform>             extent:    vec4<f32>;   // min_x,min_y,max_x,max_y

fn part1by1(x: u32) -> u32 {
    var v = x & 0x0000ffffu;
    v = (v | (v << 8u)) & 0x00ff00ffu;
    v = (v | (v << 4u)) & 0x0f0f0f0fu;
    v = (v | (v << 2u)) & 0x33333333u;
    v = (v | (v << 1u)) & 0x55555555u;
    return v;
}

@compute @workgroup_size(256)
fn build_codes(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= arrayLength(&positions)) { return; }          // guard ragged tail
    let span = vec2<f32>(extent.z - extent.x, extent.w - extent.y);
    let n = (positions[i] - extent.xy) / max(span, vec2<f32>(1e-9));
    let qx = u32(clamp(n.x, 0.0, 0.999999) * 65536.0);     // 16-bit per axis
    let qy = u32(clamp(n.y, 0.0, 0.999999) * 65536.0);
    keys[i] = KeyIndex(part1by1(qx) | (part1by1(qy) << 1u), i);
}

// ---- Pass 2a: per-digit histogram (one 256-bucket count per workgroup) ----
const RADIX: u32 = 256u;
@group(0) @binding(0) var<storage, read>       keysIn:  array<KeyIndex>;
@group(0) @binding(1) var<storage, read_write> counts:  array<atomic<u32>>;  // groups * 256
@group(0) @binding(2) var<uniform>             shift:   u32;                 // 0, 8, 16, 24

var<workgroup> local_hist: array<atomic<u32>, RADIX>;

@compute @workgroup_size(256)
fn histogram(@builtin(global_invocation_id) gid: vec3<u32>,
             @builtin(local_invocation_index) lid: u32,
             @builtin(workgroup_id) wg: vec3<u32>) {
    atomicStore(&local_hist[lid], 0u);
    workgroupBarrier();
    let i = gid.x;
    if (i < arrayLength(&keysIn)) {
        let digit = (keysIn[i].key >> shift) & (RADIX - 1u);
        atomicAdd(&local_hist[digit], 1u);                 // tally this workgroup's digits
    }
    workgroupBarrier();
    atomicStore(&counts[wg.x * RADIX + lid], atomicLoad(&local_hist[lid]));
}

// ---- Pass 2b: stable scatter into sorted order using exclusive digit offsets ----
@group(0) @binding(0) var<storage, read>       keysIn:  array<KeyIndex>;
@group(0) @binding(1) var<storage, read_write> keysOut: array<KeyIndex>;
@group(0) @binding(2) var<storage, read_write> offsets: array<atomic<u32>>; // prefixed digit bases
@group(0) @binding(3) var<uniform>             shift:   u32;

@compute @workgroup_size(256)
fn scatter(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= arrayLength(&keysIn)) { return; }
    let digit = (keysIn[i].key >> shift) & (RADIX - 1u);
    let dst = atomicAdd(&offsets[digit], 1u);              // claim next slot for this digit
    keysOut[dst] = keysIn[i];                              // stable within a digit bucket
}

// ---- Pass 3: mark node starts at a chosen tree level ----
@group(0) @binding(0) var<storage, read>       sorted:   array<KeyIndex>;
@group(0) @binding(1) var<storage, read_write> nodeFlag: array<u32>;
@group(0) @binding(2) var<uniform>             level:    u32;                // 1..16

@compute @workgroup_size(256)
fn mark_nodes(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= arrayLength(&sorted)) { return; }
    // A node at `level` covers `2*(16-level)` low bits; shift them off to get the prefix.
    let drop = 2u * (16u - level);
    let prefix = sorted[i].key >> drop;
    let prev = select(sorted[i - 1u].key >> drop, 0xffffffffu, i == 0u);
    nodeFlag[i] = select(0u, 1u, i == 0u || prefix != prev); // 1 marks a cell boundary
}

The TypeScript driver quantizes, runs four radix digit-passes (offsets rebuilt from the histogram prefix between each), then marks nodes. The offset buffer must be reset to each digit’s exclusive base before every scatter.

async function buildQuadtree(
  device: GPUDevice,
  positions: Float32Array,   // packed vec2<f32>
  pipelines: {
    build: GPUComputePipeline; hist: GPUComputePipeline;
    scatter: GPUComputePipeline; mark: GPUComputePipeline;
  },
  bindFor: (shift: number) => { hist: GPUBindGroup; scatter: GPUBindGroup },
): Promise<void> {
  const n = positions.length / 2;
  const groups = Math.ceil(n / 256);
  const enc = device.createCommandEncoder();

  // Pass 1: keys.
  const p1 = enc.beginComputePass();
  p1.setPipeline(pipelines.build);
  p1.dispatchWorkgroups(groups);
  p1.end();

  // Passes 2: four LSD radix digit-passes over the 32-bit Morton key.
  for (let shift = 0; shift < 32; shift += 8) {
    const b = bindFor(shift);
    const h = enc.beginComputePass();
    h.setPipeline(pipelines.hist);
    h.setBindGroup(0, b.hist);
    h.dispatchWorkgroups(groups);           // per-workgroup 256-bucket histograms
    h.end();
    // (prefix-sum the histograms into `offsets` here — same scan as the parent guide)
    const s = enc.beginComputePass();
    s.setPipeline(pipelines.scatter);
    s.setBindGroup(0, b.scatter);
    s.dispatchWorkgroups(groups);           // stable scatter into keysOut
    s.end();
  }

  // Pass 3: mark node boundaries at the target level.
  const m = enc.beginComputePass();
  m.setPipeline(pipelines.mark);
  m.dispatchWorkgroups(groups);
  m.end();

  device.queue.submit([enc.finish()]);
  await device.queue.onSubmittedWorkDone();
}

Once mark_nodes has flagged the boundaries, one more prefix-sum-and-scatter — the same compaction described in the parent guide — turns the flag array into a compact node table where each entry stores a cell’s Morton prefix and the [start, end) range of sorted point indices it covers. That table is the queryable index: a range query resolves the viewport’s Morton bounds and binary-searches the sorted keys for the covering nodes, and a nearest-neighbour query walks outward from the query point’s cell along z-order. Because the whole structure is two flat buffers — the sorted KeyIndex array and the node table — it uploads, binds, and streams like any other storage buffer, with no pointer chasing and no per-node allocation. The memory footprint is 8N bytes for the sorted keys plus a node table bounded by the number of occupied cells at the target level, which for real point distributions is far smaller than the 4^level worst case.

The construction cost is dominated by the radix sort: four digit-passes, each a histogram, a prefix sum, and a scatter over all N points. The quantize and mark passes are one coalesced sweep each. Rebuilding the tree every frame is affordable for a few million moving points; for a static cloud, build once and keep the sorted buffer resident, rebuilding only when the extent changes.

Parameter reference

Every tunable value in the implementation above, with guidance for point-cloud and vector workloads. These tables scroll horizontally on narrow viewports.

Failure modes

• Codes collapse to one bucket. Un-normalized coordinates or a wrong extent uniform push every quantized value to 0 or the maximum, so all points share a Morton code and the sort produces one giant node. Detection: mark_nodes flags only index 0. Fix: verify the extent matches the data’s projected bounds and that positions are projected to one linear space, per memory alignment for spatial data buffers.
• Unstable sort scrambles ties. If the scatter offsets are not reset to the exclusive digit base before each pass, or the digit-passes are not run least-significant-first, equal keys reorder between passes and the final order is not z-order. Detection: adjacent sorted entries have non-monotonic keys. Fix: rebuild offsets from the histogram prefix before every scatter, and iterate shift from 0 upward.
• Odd number of digit-passes leaves data in the wrong buffer. Ping-pong sorts swap keysIn/keysOut each pass; an odd count leaves the sorted result in the scratch buffer. Detection: the marked nodes reference the pre-sort order. Fix: use an even number of passes, or copy back after an odd count.
• Device lost on a large single dispatch. Sorting tens of millions of points in one submission can trip the driver TDR watchdog. Detection: device.lost resolves mid-build. Fix: cap per-dispatch point counts and recover through the browser support fallback routing strategies.

Backend / Python interop note

Quantization is stable only if the client and server agree on the projected extent. Compute the extent once on the Python side over the whole dataset and ship it alongside the points, rather than deriving it per tile — a per-tile extent makes the same point quantize to different codes in neighbouring tiles, so nodes do not stitch across tile seams. Emit positions as a contiguous C-ordered float32 array (packed vec2<f32>) so a writeBuffer lands directly in the storage buffer with no client repacking, exactly as the streaming path in reducing GPU memory fragmentation during spatial aggregation prescribes.

Related

• WGSL Spatial Algorithms on the GPU — the build-scan-compact skeleton this quadtree specializes.
• On-GPU Viewport Culling for Vector Tiles — query the node ranges built here against the view extent.
• Optimizing Workgroup Sizes for Vector Geometry Filtering — sizing the histogram and scatter dispatches for throughput.
• Choosing Workgroup Size for GPU Occupancy — the occupancy trade-off behind the radix sort’s block size.

Up: WGSL Spatial Algorithms on the GPU