GpuPressureProjectionProbe.kt
src/jvmMain/kotlin/afield/examples/four/GpuPressureProjectionProbe.kt · 517 lines
package afield.examples.four
import afield.graphics.vulkan.Renderer
import afield.graphics.vulkan.util.GpuProbe
import afield.graphics.vulkan.util.Options
import afield.graphics.vulkan.util.Window
import afield.graphics.vulkan.util.gpuProbe
import kotlin.math.PI
import kotlin.math.abs
import kotlin.math.cos
import kotlin.math.sin
import kotlin.system.exitProcess
FIELD_TEST id: four.gpu-probe.pressure-chain requiresDisplay: true runner: jvm-main entrypoint: afield.examples.four.GpuPressureProjectionProbe tags: [smoke, four, gpu, compute, probe, fluid, pressure, chain] timeoutSec: 600 expect: exitCode: 0 stdoutContains: - "gpu pressure projection chain PASSED"
object GpuPressureProjectionProbe
That a chain of compute dispatches — divergence, then Jacobi pressure, then project — leaves a velocity field that is measurably divergence-free, and that it gets more so with iterations.
Why a chain and not a dispatch
Every GPU probe in this repository grades a single dispatch. Real GPU work is a pipeline: the workspace document this facility was built for chains advect -> divergence -> pressure x 24 -> project, and the bugs are not inside any one pass, they are between passes — a missing barrier so pass N reads what pass N-1 has not finished writing, a ping-pong that reads the buffer it just wrote, a pass silently ordered before the one it depends on. None of those are visible from inside one dispatch, and a picture of a fluid cannot tell you that the divergence fell by 40x rather than by 1.4x.
So this is a probe over eighteen dispatches across three shaders, and what it asserts is a theorem rather than a threshold.
What it asserts
- The divergence pass agrees with a closed-form oracle. The same finite-difference stencil is written twice — once in GLSL, once in Kotlin over the uploaded fixture — and compared at every one of 4,096 cells. Not a sample: an off-by-one at a boundary is exactly what a sample misses.
- Helmholtz. Projecting a divergent field against the pressure solved from its own divergence must leave a field whose divergence is near zero. The "after" number is recomputed on the CPU from the read-back projected field, so the GPU cannot grade its own homework; the "before" number comes from the GPU's own divergence buffer, which the oracle above has already vouched for.
- The input is genuinely divergent. Without this the whole test passes on a field of zeros, which is the single most likely way for a fluid probe to be meaningless.
- The projection actually moved the field. A
projectpass that wrote its input unchanged would also, trivially, not increase the divergence. - More iterations help, monotonically. The chain is run for 1, 2, 4, 8 and 16 frames with the pressure re-zeroed each time, which is 16, 32, 64, 128 and 256 Jacobi sweeps. A solver whose residual does not fall with iteration count is not a solver — and a constant reduction across the sweep is the signature of a barrier that is not there, because the extra sweeps then all read the same stale pressure.
- The Poisson residual.
max |Lap p - div|againstmax |div|, recomputed on the CPU, which grades the Jacobi stage directly rather than through the projection. - Nothing is NaN and nothing is still sentinel. The output is prefilled with GpuProbe.SENTINEL before every sample, so a chain that silently stopped running cannot report agreement.
The discretisation, and why this one
A MAC (staggered) arrangement: u lives on the left x-face of each cell and v on the bottom
y-face, both stored in the x and y of one vec4 per cell. Divergence is a forward difference
and the pressure gradient is a backward difference, which makes them exact adjoints, so div .
grad is precisely the 5-point Laplacian the Jacobi stage inverts. The collocated alternative —
central differences for both on a cell-centred grid — has a checkerboard null space that the
pressure solve cannot see, and a probe built on it would report a small divergence for a field
carrying a large invisible one. The walls are solid: face velocities on the domain boundary are
zero, and pressure is mirrored outside it (Neumann), which is the boundary pair that makes the
projection well posed.
Running it
tools/run_field_tests.py --id four.gpu-probe.pressure-chain
Needs a GPU and a window: Field 4 has no headless mode, so a compute-only dispatch still
needs a Window for the device to exist. It also needs an awake display.
object GpuPressureProjectionProbe {
Cells per side. 64 x 64 = 4,096 cells, one workgroup of 64 per 64 cells.
private const val N = 64
private const val CELLS = N * N
Floats in a cell-per-vec4 field.
private const val FIELD_FLOATS = CELLS * 4
Jacobi sweeps per frame. One stage, this many dispatches, each followed by a barrier.
private const val JACOBI = 16
Frame counts sampled, so the effective sweep is this times JACOBI.
private val FRAME_SWEEP = listOf(1, 2, 4, 8, 16)
Workgroups needed to cover the grid at local_size_x = 64.
private const val GROUPS = CELLS / 64
The three shaders. One Field type, reused, because every field here is a vec4 per cell.
#version has to be the first thing in the source, so this literal starts on the opening
line rather than after a newline. A shader whose #version is not first reports a wall of
errors about everything else, which masks the real one.
private const val PREAMBLE = """#version 450
#extension GL_EXT_buffer_reference : require
layout(std430, buffer_reference) buffer Field { vec4 v[]; };
layout(local_size_x = 64) in;
"""
div[i,j] = (u[i+1,j] - u[i,j]) + (v[i,j+1] - v[i,j]), with solid walls.
Forward differences over face velocities. A face on the domain boundary reads zero rather than wrapping or clamping: a wall is a wall, and clamping would make the outermost ring report the divergence of a field that flows through it.
private val DIVERGENCE_GLSL = PREAMBLE + """
layout(push_constant) uniform Push {
Field velocity;
Field divergence;
uint side;
} pc;
void main() {
uint id = gl_GlobalInvocationID.x;
uint n = pc.side;
if (id >= n * n) return;
uint i = id % n;
uint j = id / n;
float uL = (i == 0u) ? 0.0 : pc.velocity.v[id].x;
float uR = (i + 1u >= n) ? 0.0 : pc.velocity.v[id + 1u].x;
float vB = (j == 0u) ? 0.0 : pc.velocity.v[id].y;
float vT = (j + 1u >= n) ? 0.0 : pc.velocity.v[id + n].y;
pc.divergence.v[id] = vec4((uR - uL) + (vT - vB), 0.0, 0.0, 0.0);
}
"""
One Jacobi sweep of Lap p = div, mirrored at the walls.
p = (pL + pR + pB + pT - div) / 4. Mirroring a missing neighbour onto the centre cell is
the discrete Neumann condition dp/dn = 0, which is the one that pairs with zero face
velocity on the boundary.
private val JACOBI_GLSL = PREAMBLE + """
layout(push_constant) uniform Push {
Field pressureIn;
Field pressureOut;
Field divergence;
uint side;
} pc;
void main() {
uint id = gl_GlobalInvocationID.x;
uint n = pc.side;
if (id >= n * n) return;
uint i = id % n;
uint j = id / n;
float pL = pc.pressureIn.v[(i == 0u) ? id : id - 1u].x;
float pR = pc.pressureIn.v[(i + 1u >= n) ? id : id + 1u].x;
float pB = pc.pressureIn.v[(j == 0u) ? id : id - n ].x;
float pT = pc.pressureIn.v[(j + 1u >= n) ? id : id + n ].x;
pc.pressureOut.v[id] =
vec4((pL + pR + pB + pT - pc.divergence.v[id].x) * 0.25, 0.0, 0.0, 0.0);
}
"""
u -= dp/dx, v -= dp/dy — backward differences, the adjoint of the divergence stencil.
Writes to a separate output field rather than in place, so the input fixture stays the thing the CPU oracle compares against and a sample can be repeated without re-uploading.
private val PROJECT_GLSL = PREAMBLE + """
layout(push_constant) uniform Push {
Field velocity;
Field pressure;
Field velocityOut;
uint side;
} pc;
void main() {
uint id = gl_GlobalInvocationID.x;
uint n = pc.side;
if (id >= n * n) return;
uint i = id % n;
uint j = id / n;
float pC = pc.pressure.v[id].x;
float pL = pc.pressure.v[(i == 0u) ? id : id - 1u].x;
float pB = pc.pressure.v[(j == 0u) ? id : id - n ].x;
float u = (i == 0u) ? 0.0 : pc.velocity.v[id].x - (pC - pL);
float v = (j == 0u) ? 0.0 : pc.velocity.v[id].y - (pC - pB);
pc.velocityOut.v[id] = vec4(u, v, 0.0, 0.0);
}
"""
Fixture and oracle, in Kotlin
A multi-scale divergent field: two low modes plus a high one on each component.
Multi-scale on purpose. Jacobi kills high-frequency error in a handful of sweeps and low
frequency in hundreds, so a field with only one scale in it would report either a
spectacular reduction or a negligible one, and neither would say anything about the chain. A
superposition makes the sweep in FRAME_SWEEP show the shape of the convergence rather than
one point on it.
private fun seedField(): FloatArray {
val out = FloatArray(FIELD_FLOATS)
for (j in 0 until N) for (i in 0 until N) {
val x = (i + 0.5) / N
val y = (j + 0.5) / N
val u = sin(2 * PI * x * 3) * cos(2 * PI * y * 2) + 0.5 * sin(2 * PI * x * 9)
val v = cos(2 * PI * x * 2) * sin(2 * PI * y * 3) + 0.5 * sin(2 * PI * y * 7)
val o = (j * N + i) * 4
out[o] = u.toFloat()
out[o + 1] = v.toFloat()
}
return out
}
DIVERGENCE_GLSL, in Kotlin, over a vec4-per-cell field. The oracle.
private fun divergenceOf(field: FloatArray): FloatArray {
val out = FloatArray(CELLS)
for (j in 0 until N) for (i in 0 until N) {
val id = j * N + i
val uL = if (i == 0) 0f else field[id * 4]
val uR = if (i + 1 >= N) 0f else field[(id + 1) * 4]
val vB = if (j == 0) 0f else field[id * 4 + 1]
val vT = if (j + 1 >= N) 0f else field[(id + N) * 4 + 1]
out[id] = (uR - uL) + (vT - vB)
}
return out
}
max |Lap p - div|, with the same mirrored walls the Jacobi stage uses.
private fun poissonResidual(pressure: FloatArray, divergence: FloatArray): Float {
var worst = 0f
for (j in 0 until N) for (i in 0 until N) {
val id = j * N + i
val pC = pressure[id * 4]
val pL = pressure[(if (i == 0) id else id - 1) * 4]
val pR = pressure[(if (i + 1 >= N) id else id + 1) * 4]
val pB = pressure[(if (j == 0) id else id - N) * 4]
val pT = pressure[(if (j + 1 >= N) id else id + N) * 4]
worst = maxOf(worst, abs((pL + pR + pB + pT - 4 * pC) - divergence[id]))
}
return worst
}
private fun maxAbs(a: FloatArray): Float {
var m = 0f
for (v in a) if (v.isFinite()) m = maxOf(m, abs(v))
return m
}
fun main(args: Array<String>)
body — 235 lines, from line 282
@JvmStatic
fun main(args: Array<String>) {
Options.enableVR = false
val window = Window("GPU Pressure Projection Probe")
val renderer = Renderer(window)
var failures = 0
try {
failures = gpuProbe(
renderer,
name = "pressure-chain",
pump = { window.pollEvents(); renderer.render() },
) { probe ->
run(probe)
}
} finally {
renderer.cleanup()
window.cleanup()
}
if (failures > 0) {
println("gpu pressure projection chain FAILED ($failures)")
exitProcess(1)
}
println("gpu pressure projection chain PASSED")
exitProcess(0)
}
private fun run(probe: GpuProbe): Int {
// --- fixtures. Every one of these is host-readable because `fixture` insists on it.
val velocity = probe.fixture("velocity", FIELD_FLOATS)
val velocityOut = probe.fixture("velocityOut", FIELD_FLOATS)
val divergence = probe.fixture("divergence", FIELD_FLOATS)
val pressureA = probe.fixture("pressureA", FIELD_FLOATS)
val pressureB = probe.fixture("pressureB", FIELD_FLOATS)
val seeded = seedField()
probe.upload(velocity, seeded)
probe.fill(pressureA, 0f)
probe.fill(pressureB, 0f)
// --- the chain. Three stages, three shader compiles, 1 + JACOBI + 1 dispatches.
//
// Stage order is registration order and dispatch order is insertion order, so this reads
// as the pipeline it is. The ping-pong is explicit and visible: iteration `k` reads the
// buffer iteration `k - 1` wrote, and because JACOBI is even the answer ends up back in
// pressureA, which is what `project` binds.
probe.stage("divergence", DIVERGENCE_GLSL).dispatch(groups = GROUPS) {
bind("velocity") { velocity }
bind("divergence") { divergence }
bind("side") { N }
}
val jacobi = probe.stage("jacobi", JACOBI_GLSL)
require(JACOBI % 2 == 0) { "JACOBI must be even so the result lands back in pressureA" }
for (k in 0 until JACOBI) {
val from = if (k % 2 == 0) pressureA else pressureB
val to = if (k % 2 == 0) pressureB else pressureA
jacobi.dispatch(groups = GROUPS, label = "sweep$k") {
bind("pressureIn") { from }
bind("pressureOut") { to }
bind("divergence") { divergence }
bind("side") { N }
}
}
probe.stage("project", PROJECT_GLSL).dispatch(groups = GROUPS) {
bind("velocity") { velocity }
bind("pressure") { pressureA }
bind("velocityOut") { velocityOut }
bind("side") { N }
}
println(
"--- chain: ${probe.stageList.size} stages, " +
"${probe.stageList.sumOf { it.dispatchCount }} dispatches over $CELLS cells " +
"(${N}x$N), $JACOBI Jacobi sweeps per frame"
)
// --- wait for the chain to exist at all. This is the wait that matters: a ComputeModule
// whose shader has not finished compiling dispatches as a silent no-op, so the question is
// "has the output been written", not "have four frames gone by".
// The two graded outputs, and deliberately NOT the pressure pair.
//
// Arming `pressureA`/`pressureB` looks strictly better and is wrong, measured: it hung
// for 1434 frames and then reported `9 of 16384 slots still hold the sentinel, from index
// 908`. The Jacobi stencil is an average, and a constant is a fixed point of an average:
// `(4 x -98765 - 0) / 4` is exactly `-98765`, so any interior cell whose four neighbours
// are all still sentinel and whose divergence is exactly zero reproduces the sentinel
// forever. A sentinel must not be a fixed point of the shader that is supposed to
// overwrite it — witness the END of a chain, not an intermediate that is fed back into
// itself.
//
// The Jacobi stage's execution is not left unwitnessed by this. A stage that never ran
// leaves the pressure at zero, the projection returns its input, and three of the claims
// below — the reduction, `extra sweeps helped` and `projection changed the field` — fail
// together rather than quietly.
probe.settle(velocityOut, divergence)
val warmFrames = probe.framesToLand
println(" cold start: landed after $warmFrames pumped frames (three shader compiles)")
for (s in probe.stageList) {
if (!s.valid) probe.fail("stage '${s.stageName}' never became valid")
}
// --- claim 1: the divergence pass, against the closed-form oracle, at every cell.
val gpuDiv = probe.read(divergence)
val cpuDiv = divergenceOf(seeded)
var worstDivError = 0f
var worstDivIndex = -1
for (id in 0 until CELLS) {
val e = abs(gpuDiv[id * 4] - cpuDiv[id])
if (e > worstDivError) { worstDivError = e; worstDivIndex = id }
}
val maxDivBefore = maxAbs(cpuDiv)
println(
(" divergence pass vs closed form: worst |error| %.3e at cell %d of %d; " +
"field carries max |div| %.4f")
.format(worstDivError, worstDivIndex, CELLS, maxDivBefore)
)
// Measured 1.192e-07 at cell 96, which is one ULP of a number of order 1 — the GLSL and
// the Kotlin sum the same four terms in the same order, so the only difference available
// is the compiler's contraction choice. 1e-6 is eight times the measurement, which leaves
// room for a fused multiply-add on one side and not the other, and nothing else.
probe.atMost(
"divergence pass vs closed form", worstDivError, 1e-6f,
buffer = "divergence", index = worstDivIndex
)
probe.allFinite("divergence pass", gpuDiv, buffer = "divergence")
// The trivial-pass guard: a field of zeros would satisfy everything below. The fixture is
// a fixed superposition computed on the CPU, so this is exactly reproducible — measured
// 1.7649, and a floor of 1.5 fails the moment somebody edits `seedField` into something
// flat rather than passing quietly on it.
probe.atLeast("input field max |div|", maxDivBefore, 1.5f, buffer = "velocity")
// --- claim 2: the sweep. Pressure re-zeroed each sample, so `frames x JACOBI` is the
// whole iteration count and the curve is comparable point to point.
println(" projected divergence against Jacobi sweeps (population: $CELLS cells each):")
data class Sample(val frames: Int, val sweeps: Int, val after: Float, val ratio: Float)
val samples = mutableListOf<Sample>()
var lastProjected = FloatArray(0)
for (frames in FRAME_SWEEP) {
probe.idle()
probe.fill(pressureA, 0f)
probe.fill(pressureB, 0f)
probe.arm(velocityOut, divergence)
while (probe.stillWaiting()) probe.pumpFrames(1)
val landed = probe.framesToLand
if (frames > landed) probe.pumpFrames(frames - landed)
val pumped = maxOf(frames, landed)
probe.idle()
val projected = probe.read(velocityOut)
lastProjected = projected
val after = maxAbs(divergenceOf(projected))
val sweeps = pumped * JACOBI
samples += Sample(pumped, sweeps, after, maxDivBefore / after)
println(
" %2d frames = %4d sweeps : max |div| after %.5f reduction %6.1fx".format(
pumped, sweeps, after, maxDivBefore / after
)
)
}
// Monotone, as a count rather than as a boolean, so a single non-monotone point is
// distinguishable from a flat curve.
var regressions = 0
for (k in 1 until samples.size) {
// 1.001 of slack: two adjacent points can agree to fp32 once converged.
if (samples[k].after > samples[k - 1].after * 1.001f) regressions++
}
probe.exactly("sweep is monotone (non-monotone steps)", regressions, 0)
val best = samples.last()
val worst = samples.first()
// Measured 2250.8x at 256 sweeps. Jacobi has no atomics and no race, the fixture is
// fixed, and the frame count at this end of the sweep is exactly 16, so this number is
// reproducible on a given GPU rather than approximately so. 1500 is a third below it; a
// chain that lost a stage or a barrier reports single digits, not 1400.
probe.atLeast(
"reduction at ${best.sweeps} sweeps", best.ratio, 1500.0f, buffer = "velocityOut"
)
// And that the *extra* sweeps did something. A chain missing a barrier between Jacobi
// iterations reads a stale pressure every time, so every sample comes back identical --
// which passes a reduction floor and a monotonicity check but not this.
//
// Measured 533x. The floor is 100 rather than 400 because the *first* sample's population
// is `max(1, framesToLand)`: on a contended GPU that can be 2 frames, which makes the
// first point 32 sweeps instead of 16 and legitimately shrinks this ratio to about 286.
// The sweep table prints the frame count for each point, so the population is visible.
probe.atLeast(
"extra sweeps helped (${worst.sweeps} -> ${best.sweeps})",
worst.after / best.after, 100.0f
)
// --- claim 3: the projection moved the field, and left it finite.
var worstMove = 0f
for (i in 0 until FIELD_FLOATS) worstMove = maxOf(worstMove, abs(lastProjected[i] - seeded[i]))
// Measured 1.434 against a field whose own max |div| is 1.765, i.e. the correction is the
// same order as the thing being corrected. A floor of 1.0 is a real constraint; 0.1 would
// have passed on a projection that barely touched the field.
probe.atLeast("projection changed the field", worstMove, 1.0f, buffer = "velocityOut")
probe.allFinite("projected field", lastProjected, buffer = "velocityOut")
probe.varies("projected field", lastProjected, 1.2f, buffer = "velocityOut")
// --- claim 4: the Jacobi stage directly, as a Poisson residual.
val pressure = probe.read(pressureA)
val residual = poissonResidual(pressure, cpuDiv)
println(
" Poisson residual max |Lap p - div| = %.5f against max |div| %.4f (%.1f%%)".format(
residual, maxDivBefore, 100f * residual / maxDivBefore
)
)
// Measured 4.45e-04 of max |div| after 256 sweeps. 5e-3 is an order of magnitude above it
// and still an order of magnitude below the 4.5e-2 that 64 sweeps would give, so this
// claim can tell the sweep lengths apart rather than accepting any of them.
probe.atMost("Poisson residual / max |div|", residual / maxDivBefore, 5e-3f, buffer = "pressureA")
probe.allFinite("pressure", pressure, buffer = "pressureA")
// Measured 11.91. The pressure field is the solution of a Poisson problem with an O(1)
// right-hand side over 64 cells, so it is an order of magnitude larger than the velocity;
// a floor of 8 catches a solve that collapsed to near-constant, which 0.05 would not.
probe.varies("pressure", pressure, 8.0f, buffer = "pressureA")
// --- and that the whole chain re-lands promptly once warm, which is the claim that the
// sentinel mechanism is measuring something rather than being always true.
//
// Measured 1. It is structurally 1: `landed()` waits *every* compute fence for this
// window, so one pumped frame is one completed chain no matter how many frames the
// renderer keeps in flight. The bound is 2 rather than 1 only so a scheduling hiccup is
// not a red test; a 3 here would mean the sentinel is surviving a completed frame, which
// is a real finding and should fail.
probe.atMost("frames to land once warm", probe.framesToLand, 2)
return probe.verdict("gpu pressure projection chain")
}
}