Sorting Water

Part one and part two were about gravity, where every body pulls on every other body no matter how far apart they are. Most of the work went into summarizing the bodies that are far away.

The water at the top of this page works the other way around. A bit of water only interacts with the bits right next to it: it pushes on its neighbours, gets pushed back, and tens of thousands of those small pushes add up to pressure and waves. Collisions and molecular dynamics work the same way, with forces that drop to zero past some cutoff distance. That sounds easier than gravity, and it is once you know which particles are your neighbours. Finding that out fast on a GPU took a counting sort and a parallel prefix sum, and the rest of this page is how they fit together.

Step 1 · ContactsMostly checking for nothing

Start with a box of discs under gravity. Two discs push apart when they overlap, using a stiff spring with a little damping, and don't feel each other at all otherwise. The obvious way to write it is the same pair of nested loops as the galaxy:

for (let i = 0; i < count; i++) {
  for (let j = i + 1; j < count; j++) {
    if (distance(i, j) < 2 * radius) {   // almost always false
      applyContactForce(i, j);
    }
  }
}

The counter under the demo shows what that costs. At the default 2,000 discs it's about two million distance checks per step, a few milliseconds of JavaScript. Only around 8,000 of those pairs are touching while the pile is still settling, and closer to 13,000 once it has packed down. Either way it's well under 1%. With gravity, every term in the O(n²) sum added something. Here more than 99% of the checks compute a force of zero.

Your cursor shoves the discs near it, and that costs the same two million checks. Part one's fix doesn't help here. A Barnes–Hut tree answers "what does everything far away add up to?", and for contact forces the answer is always zero. This needs the opposite kind of structure, one that answers "who is near me?" and skips everything else.

Step 2 · BucketsA uniform grid

The standard answer is a grid. Lay a uniform grid over the world, with cells at least as wide as the interaction distance, and put each particle in its cell. Anything close enough to touch a particle has to be in the same cell or one of the eight around it. So the neighbour query becomes: look in nine cells and test what's in them. If the particles are spread fairly evenly, that's O(n) work in total.

The demo below is the same kind of pile, now 3,500 discs on a 24×24 grid. Your cursor runs the query live. The yellow 3×3 window is every cell it would search, and the yellow discs are the only candidates it would test. When the cursor isn't over the box, the query parks itself inside the pile. Near the top of the pile the window holds about a hundred discs, and a few hundred down near the floor, where the pile is packed tightest. The rest are never looked at. Over a whole step, the counter shows a bit over 400,000 candidate pairs, against 6.1 million for checking every pair.

The simplest CPU version is short. Turn each position into one integer cell key, append the particle's index to that cell's bucket, and then visit the 3×3 block around each particle:

const buckets = Array.from({ length: grid * grid }, () => []);

for (let i = 0; i < particles.length; i++) {
  buckets[cellOf(particles[i].position)].push(i);
}

function testNeighbours(i) {
  const [cellX, cellY] = cellCoord(particles[i].position);
  for (let offsetY = -1; offsetY <= 1; offsetY++) {
    for (let offsetX = -1; offsetX <= 1; offsetX++) {
      const cell = cellAt(cellX + offsetX, cellY + offsetY);
      if (cell === undefined) continue;           // outside the box
      for (const j of buckets[cell]) {
        if (j > i) testPair(i, j);                // each pair once
      }
    }
  }
}

(The demo above already uses the sorted layout from the next section, but this is the version you'd write first.)

The query part is fine. The problem for the GPU is the innocent-looking push. Appending to a list means allocation, pointers, and a write position that depends on who appended before you. That's the same serial dependency that made the Barnes–Hut tree hard to build in parallel, and thousands of threads taking turns to append would spend most of their time waiting. Whatever structure I use has to be built the GPU way, with every thread working at once.

Part one got around this by making the tree implicit, so there was nothing to link. For the grid, a sort does the same job.

Step 3 · SortingCounting sort and the prefix sum

What the query needs is each cell's particles stored next to each other, plus the index where each cell's group starts. Sorting the particles by cell number gives you that: each cell becomes a contiguous run in one flat array. So the pointer structure I couldn't build in parallel can be replaced by an ordering, and GPUs can sort.

Sorting by cell number doesn't even need comparisons. The method is counting sort. Count how many particles land in each cell. Then, for every cell, add up the counts of all the cells before it; that running total is where the cell's group starts. Finally, move every particle straight to its slot.

counts[cell]++                       // 1. histogram    (one pass)
starts[c] = counts[0..c-1] summed    // 2. prefix sum   (the hard part)
slot = starts[cell] + claimed++      // 3. scatter      (one pass)

In serial JavaScript the whole sort looks like this. The GPU version keeps the same arrays and the same three phases. The loops become compute dispatches and the two increments become atomics.

const counts = new Uint32Array(cellCount);
for (const p of particles) counts[cellOf(p)]++;          // histogram

const starts = new Uint32Array(cellCount);
for (let c = 1; c < cellCount; c++) {
  starts[c] = starts[c - 1] + counts[c - 1];             // exclusive scan
}

const cursor = starts.slice();
const sorted = new Array(particles.length);
for (const p of particles) {
  const cell = cellOf(p);
  sorted[cursor[cell]++] = p;                             // scatter
}

The second phase, turning per-cell counts into running totals, looks like the one part that has to run in order. A running total seems to need a single accumulator moving left to right, since you can't know the total through cell 100 until you've finished cell 99.

The parallel version I used is the one Daniel Hillis and Guy Steele described in 1986, and it works by doubling. The scan runs in rounds, and in each round every cell updates at the same time. In round one, each cell adds in the value one place to its left, so it now holds the sum of two cells. In round two, each cell adds the value two places to its left. That value already covers two cells, so each cell now holds the sum of four. Round three reaches back four places and covers eight. After log₂ n rounds (ten for a thousand cells), every position holds the sum of everything up to and including itself. Within a round, no thread waits on any other.

That gives totals that include each cell's own count. Shift them one place to the right and put a zero in front, and each cell holds the index where its run starts.

Here's the whole process on 18 particles and 6 cells, scrubbed by your scroll position: the histogram, the scan rounds, the shifted "starts" row, and the scatter. Go slowly through the scan rounds. They're the part worth rewinding.

And here it is a bit bigger, running on its own: 56 particles in 16 cells, four scan passes, then the shift, then the scatter. The buttons step through it one phase at a time, toggle auto-play, and shuffle in a new set of particles.

This operation, called a scan or prefix sum, turned out to be useful far beyond sorting. Guy Blelloch argued that scans should count as basic parallel operations and showed how much can be built from them: radix sort, stream compaction, and handing out memory to many threads at once all come down to scans. I think of it as the parallel version of any loop that carries a running total forward. If the operation in that loop is associative, a scan can run it in logarithmic depth.

pass.setPipeline(this.pipes.scan_blocks);  pass.dispatchWorkgroups(BLOCKS);  // BLOCKS = 256
pass.setPipeline(this.pipes.scan_sums);    pass.dispatchWorkgroups(1);
pass.setPipeline(this.pipes.scan_add);     pass.dispatchWorkgroups(BLOCKS);

var<workgroup> sa: array<u32, 256>;
var<workgroup> sb: array<u32, 256>;

// inclusive scan of the 256 values staged in sa
fn scanShared(lid: u32) -> u32 {
  var fromA = true;
  var d = 1u;
  loop {
    if (d >= 256u) { break; }
    if (fromA) {
      var v = sa[lid];
      if (lid >= d) { v += sa[lid - d]; }
      sb[lid] = v;
    } else {
      var v = sb[lid];
      if (lid >= d) { v += sb[lid - d]; }
      sa[lid] = v;
    }
    workgroupBarrier();
    fromA = !fromA;
    d = d << 1u;
  }
  return sa[lid];  // 8 rounds, so the result lands back in sa
}

Step 4 · AssemblyThe whole sort on the GPU

The particles never stop moving, so the full pipeline runs from scratch every step:

clear    → zero the cell counters
count    → every particle atomically bumps its cell    (histogram)
scan     → starts[] from counts[]                      (3 dispatches)
scatter  → every particle copies itself to its slot    (one atomicAdd each)
forces   → every particle reads its 9 cells, pushes, integrates  (two passes for water)

The histogram and scatter kernels are a few lines of WGSL each. The atomics don't put the particles in any particular order within a cell. They only make sure no two particles get the same slot.

@compute @workgroup_size(256)
fn count(@builtin(global_invocation_id) gid: vec3u) {
  if (gid.x >= GP.count) { return; }
  atomicAdd(&counts[cellOf(partsIn[gid.x].xy)], 1u);
}

@compute @workgroup_size(256)
fn scatter(@builtin(global_invocation_id) gid: vec3u) {
  if (gid.x >= GP.count) { return; }
  let p = partsIn[gid.x];
  let slot = atomicAdd(&cursor[cellOf(p.xy)], 1u);
  sorted[slot] = p;
}

The scatter copies each particle's whole state into the sorted array, not just its index. After the sort, particles that are near each other in space are near each other in memory, so the force pass reads short contiguous runs instead of jumping around. In part one I listed Morton ordering as something to try for the same reason; here that kind of locality comes with the sort. The sort writes into a second buffer, the forces run on that one in place, and the next step sorts it back into the first. The only coordination anywhere is atomic increments, which the hardware resolves without locks.

Here it is running as dry sand: 12,000 grains by default (the slider goes from 2,000 to 30,000), re-sorted every substep, five substeps per frame. Each grain is a disc one grid cell across, so any grain it can touch is in the 3×3 block around it. Drag your cursor through the pile.

Look at what the slider costs now. The grain size is fixed, so more grains means a deeper pile, and the count, scatter and force passes each do work in proportion to the number of grains. The scan doesn't grow at all, since it's always 65,536 cells. There's no n² term left anywhere in the pipeline.

Step 5 · WaterSmoothed particles

Sand grains only interact when they collide. Water needs more than that. It resists being compressed everywhere, including in the still middle of the tank. The method I used is smoothed particle hydrodynamics (SPH). The fluid is still a set of particles, but each one is treated as a small blurry blob of matter instead of a point, and the density anywhere is the sum of the blobs overlapping that spot. Above the rest density the fluid is squeezed and pushes its neighbours apart. Below it, the fluid pulls them back in. Pressure ends up as one more short-range force, so the neighbour search handles it without changes.

// pass 1, density: how crowded is each particle?
ρ_i = Σ over neighbours j:  (1 − r/h)²

// pass 2, pressure force: push from crowded toward roomy
P_i  = k · (ρ_i − ρ_rest)
acc -= r̂ · (P_i + P_j)/2 · (1 − r/h)

The smoothing radius h is one grid cell, so the 3×3 block always covers it, and the sort turns "over neighbours" into nine ordinary integer ranges. This is the inner loop of the density pass. It visits nine cells, reads each one's contiguous run, and skips anything outside h:

for (var oy = -1; oy <= 1; oy++) {
  for (var ox = -1; ox <= 1; ox++) {
    let c = cc + vec2i(ox, oy);
    if (c.x < 0 || c.y < 0 || c.x >= i32(SP.grid) || c.y >= i32(SP.grid)) { continue; }
    let ci = u32(c.y) * SP.grid + u32(c.x);
    let s = cellStart[ci];               // where this cell's run begins
    let n = cellCount[ci];               // and how long it is
    for (var k = s; k < s + n; k++) {
      let d = parts[k].xy - pi;
      let r2 = dot(d, d);
      if (r2 < h2) {
        let q = sqrt(r2) / h;
        let w = 1.0 - q;
        rho += w * w;
        rhoNear += w * w * w;            // for near pressure, see below
      }
    }
  }
}

It starts as a wide block of water on the left side of the tank, a loose version of the classic dam-break test. There are 50,000 particles by default (the slider goes from 10,000 to 80,000), sorted into the 256×256 grid and run through the density and force passes five times per frame. The particle size is fixed, so the slider changes how much water there is, not how finely it's resolved.

EpilogueWhat carries over

A few things from this part that apply well beyond water:

• On a GPU, a data structure can often be an ordering. Per-cell lists couldn't be built in parallel, but sorting the particles by cell holds the same information in plain layout. Part one's pyramid replaced pointers with arithmetic, and counting sort replaces them with a permutation.
• Scan is worth knowing well. It's what made the sort parallel, and it handles compaction, allocation and a lot of other steps that look serial at first.
• Most of the pipeline is moving data. In the water demo, five dispatches organize data (count, three scan passes, scatter) and two compute forces. The physics is a few lines. Most of the work went into arranging memory so those lines read contiguous runs.
• Near and far forces want different structures. Gravity needed a hierarchy, because distant things still matter in total. Contacts needed flat buckets, because distant things don't matter at all.

Things I'd try next: a third dimension, which is the same sort with a third coordinate folded into the cell index; radix sort, which is counting sort applied a digit at a time, for keys too large for one histogram; and both structures in one simulation, with part one's mass pyramid for gravity and this grid for collisions. That last one is the start of a planet-formation simulator, and it's what part four builds.

Next: Part four, Worlds from Dust. Part one's gravity pyramid and this part's grid run over the same particles, and a ring of dust around a star clumps into moonlets.