GlslAd.kt

src/jvmMain/kotlin/afield/lang/stacksl/cpu/GlslAd.kt · 744 lines

package afield.lang.stacksl.cpu

import afield.math.autodiff.AD
import afield.math.autodiff.ADTape
import kotlin.math.absoluteValue

Reverse-mode automatic differentiation for the GLSL interpreter: d(pixel)/d(literal).

GlslAdScalars is a GlslScalars instance, so a shader differentiated here runs through the interpreter's one copy of the control flow, over the same parse tree the Float reference path walks. Nothing about a GLSL builtin's meaning is restated; what is added is a local partial beside each operation.

Why reverse mode

An image loss over a wire vector is many-inputs-one-output, which is reverse mode's shape: one forward pass plus one backward sweep answers for every wire at once, so the cost does not grow with the number of free parameters. The loss is a sum over pixels, so the gradient is a sum of per-pixel backward sweeps, one short tape per pixel, thrown away before the next. That bounds the tape: a 64x64 fit never holds more than one pixel's operations at a time, where a whole-image tape would hold 4096 times as many.

The primal is bit-identical to the Float path

AD computes in Double, and a Double shader is the wrong shader. GlslValue.kt's overview sets out why: widening an intermediate makes the CPU result more accurate than the GPU's and therefore a worse reference, and fract of a large product amplifies a last-bit difference to O(1). On a hash-bearing program a Double primal would not be slightly different, it would be a different image.

So every node this file pushes stores its value rounded to Float, at exactly the points the Float path rounds. Double rounding is safe for + - * / and sqrt: Double has more than twice Float's significand, so rounding a Double-exact sum to Float gives the same answer as the Float operation. The differentiable pass's primal image is therefore the same image the reference path produces, which StackSLGradientTest asserts bit for bit over a named population, while the adjoint is the derivative of the smooth relaxation, which is the quantity a fit needs.

Derivative conventions

• a branch, a loop bound and an index carry no derivative: differentiated as taken, which is correct for the piecewise function away from the switch-over. autodiff.kt says the same.
• floor, ceil, trunc, round, roundEven, sign, step and int(...) truncation carry exactly zero derivative. fract therefore carries derivative one, since it is x - floor(x).
• abs, min and max split evenly at an exact tie, which is what a central difference reports at a kink, so the analytic gradient and the finite-difference oracle agree there.
• a comparison, any, all, isnan and isinf produce a derivative-free bool.
• texture() is refused. A bilinear CPU texture fetch is differentiable in principle, but the two feedback gallery programs are already excluded from every asserted comparison in this backend, so a derivative through a sampler would have no oracle to check it against. StackSLDifferentiabilityTest counts it as unsupported rather than as zero.

internal class AdValue(val type: GlslType, val c: Array<AD>)

A GLSL value whose components are AD nodes on one tape: the differentiable GlslValue.

init

    init {
        require(c.size == type.components) { "$type needs ${type.components} components, got ${c.size}" }
    }

fun primal(): GlslValue = GlslValue(type, FloatArray(c.size) { c[it].value.toFloat() })

The primal, as the Float path would hold it — which is the point of fl.

override fun toString(): String = "ad:" + primal().toString()

Round a Double to the nearest Float and back — where the Float path rounds, this rounds.

private fun fl(v: Double): Double = v.toFloat().toDouble()

internal class GlslAdScalars

The differentiable instance of GlslScalars, over one ADTape.

One instance per tape, and one tape per pixel; the instance is cheap (it holds the tape and a surrogate mode) and is emphatically not thread-safe, exactly like GlslInvocation.

surrogate selects which of the jump-discontinuous operations get a mollified-delta partial in place of their exact zero — see AdSurrogate. It is a mode of this instance, not a second evaluator: the same parse tree, the same GlslMachine control flow, the same primal values. With the default AdSurrogate.OFF every partial and every tape node is what it was before the mode existed, and the tape does not even change length.

internal class GlslAdScalars(
    val tape: ADTape,
    val surrogate: AdSurrogate = AdSurrogate.OFF
) : GlslScalars<AdValue> {

tape primitives

fun constant(v: Double): AD = tape.constant(fl(v))

A constant: no parents, so no derivative flows into it.

fun variable(v: Double): AD = tape.variable(fl(v))

A free parameter — a wire. Its adjoint after a backward sweep is what a fit reads.

One node with one parent: value with ∂value/∂x = d.

    private fun node(x: AD, value: Double, d: Double): AD = AD(tape, tape.push(value, x.index, d, -1, 0.0))

One node with two parents.

    private fun node(x: AD, y: AD, value: Double, dx: Double, dy: Double): AD =
        AD(tape, tape.push(value, x.index, dx, y.index, dy))

A node whose value is rounded to Float — one node per Float operation, not two.

Pushing the rounded value directly (rather than an exact node followed by a rounding node) halves the tape and keeps the local partials those of the unrounded operation, which is what the smooth relaxation's derivative is.

    private fun fnode(x: AD, value: Double, d: Double): AD = node(x, fl(value), d)

    private fun fnode(x: AD, y: AD, value: Double, dx: Double, dy: Double): AD = node(x, y, fl(value), dx, dy)

Derivative-free: a value the gradient must not flow through (floor, a comparison, int).

    private fun frozen(value: Double): AD = tape.constant(fl(value))

the jump sites: exactly zero, or surrogate's mollified delta

A piecewise-constant operation with one parent: the exact value, and either its exact zero partial or AdSurrogate's mollified delta.

u is the signed distance from the argument to the nearest jump and height is that jump's size. When family is not surrogated — or when the kernel underflows, which it does a few temperatures away from a jump — this falls back to frozen, so the off path pushes the same parentless constant it always did and the tape's length is unchanged.

    private fun jumpNode(
        x: AD,
        value: Double,
        family: AdSurrogateFamily,
        u: Double,
        height: Double = 1.0
    ): AD {
        if (!surrogate.covers(family)) return frozen(value)
        val d = surrogate.kernel(u, height)
        return if (d == 0.0) frozen(value) else node(x, fl(value), d)
    }

Signed distance to the nearest integer — where floor, ceil and trunc jump.

    private fun toNearestInteger(x: Double): Double = x - Math.rint(x)

Signed distance to the nearest half-integer — where round and roundEven jump.

    private fun toNearestHalfInteger(x: Double): Double = (x - 0.5) - Math.rint(x - 0.5)

step(edge, x): one jump of height 1 at x == edge, so the partial is +K in x and −K in edge. Two parents, which is why it cannot go through jumpNode.

The isMask branch is the one place this file departs from "one jump, one kernel", and it does so because the composition it fires on is not a second jump. See maskNodes.

    private fun stepAd(e: AD, x: AD): AD {
        val value = if (x.value < e.value) 0.0 else 1.0
        if (!surrogate.covers(AdSurrogateFamily.STEP)) return frozen(value)
        // Re-thresholding a mask is the IDENTITY, not a jump: for `x` in {0, 1} and an edge
        // strictly inside (0, 1), `step(edge, x) == x` exactly. So the local partial is 1 in `x`
        // and — because moving the edge anywhere inside (0, 1) changes nothing — 0 in `edge`, and
        // the kernel belongs to whichever real jump produced `x` rather than being applied twice.
        if (isMask(x) && e.value > 0.0 && e.value < 1.0) return markMask(node(x, fl(value), 1.0))
        val d = surrogate.kernel(x.value - e.value, 1.0)
        return markMask(if (d == 0.0) frozen(value) else node(x, e, fl(value), d, -d))
    }

boolean-ish masks, which is what StackSL's predicate words actually produce

Tape indices whose value is 0.0 or 1.0 by construction — a mask in the sense public/docs/stacksl_cheatsheet.md fixes as the contract of select, if, and, or, not, any, all and band: a condition is read "boolean-ish, with threshold 0.5".

Why this set exists, and why a value test could not replace it

select/if lowers to mix(f, t, step(0.5, abs(cond))). When cond is itself a mask — the common case, since threshold lowers to step — that step(0.5, abs(·)) pair is the identity on the values cond can take, so its correct mollified partial is 1, and the kernel that belongs to the chain is the inner jump's alone. Treated as a jump of its own instead, its argument sits at distance exactly 0.5 from its edge forever, K_t(0.5) decays like exp(−1/2t), and a gradient that has to cross two of them is multiplied by it — so select switches itself off as a schedule cools. That was measured at 1.45e+11× across dev_docs/STACKSL_SURROGATE.md's sweep before this set existed.

The tempting cheap version is to test the value: pass the partial through whenever the argument happens to be 0.0 or 1.0. That is wrong, and not in a corner — a wire seeded at exactly 1.0 under w 0.5 threshold would get a partial of 1 where the honest answer is K_t(0.5) ≈ 0, turning a wall into a cliff. The question is not what the argument is here, it is what the argument ranges over, and only the expression that built it knows that. So membership is structural: pushed at the operations whose range really is {0, 1}, and propagated only through the ones that map {0, 1} into itself.

Every pass-through is then exact rather than approximate, which is why this is a fix and not a second surrogate with a second bias to measure.

Extent, deliberately small

step produces a mask; abs, min and max propagate one (|0|=0, |1|=1, and both reducers of {0,1} stay in {0,1}) — which covers threshold/step/select/if and the and/or that lower to min/max of two of them. not's vec4(1.0) − mask is not tracked: it would put a mask test on sub, the hottest arithmetic site in the file, to recover one word. band is step(a) − step(b) and is outside for the same reason. So x not select still nests, and AdSurrogateFamily.STEP's KDoc and the surrogate doc both say so rather than leaving it to be discovered.

A constant is never marked, even one whose value is 0.0 or 1.0. That is what stops the minAd/maxAd propagation leaking through clamp: saturate of a mask is min(max(mask, 0.0), 1.0), and since the bounds are unmarked the result is not a mask either. Marking them would be sound — a constant's range is a single point inside {0, 1} and it carries no derivative regardless — but it buys nothing measured, and the smaller surface is the one whose every pass-through is easy to check by hand.

Only populated when AdSurrogateFamily.STEP is surrogated: with AdSurrogate.OFF a step is a parentless constant, nothing downstream of it carries a derivative at all, and this file's promise that the off path pushes the same tape it always did is kept by allocating nothing.

    private var maskNodes: java.util.BitSet? = null

True when a mask's provenance can change a partial — i.e. when step is surrogated.

    private val trackMasks: Boolean = surrogate.covers(AdSurrogateFamily.STEP)

Record a as 0/1-valued by construction, and return it unchanged.

    private fun markMask(a: AD): AD {
        if (trackMasks) (maskNodes ?: java.util.BitSet().also { maskNodes = it }).set(a.index)
        return a
    }

Is a a mask?

The value test is redundant with the invariant and kept anyway: it is the cheap guard that makes a mis-propagated mark unable to produce a wrong partial, only a missed one.

    private fun isMask(a: AD): Boolean =
        trackMasks && (a.value == 0.0 || a.value == 1.0) && maskNodes?.get(a.index) == true

Float-faithful elementary arithmetic

    private fun add(a: AD, b: AD): AD = fnode(a, b, a.value + b.value, 1.0, 1.0)
    private fun sub(a: AD, b: AD): AD = fnode(a, b, a.value - b.value, 1.0, -1.0)
    private fun mul(a: AD, b: AD): AD = fnode(a, b, a.value * b.value, b.value, a.value)

    private fun div(a: AD, b: AD): AD {
        val inv = 1.0 / b.value
        return fnode(a, b, a.value * inv, inv, -a.value * inv * inv)
    }

    private fun neg(a: AD): AD = node(a, -a.value, -1.0)

    private fun mulK(a: AD, k: Double): AD = fnode(a, a.value * k, k)
    private fun subFromK(k: Double, a: AD): AD = fnode(a, k - a.value, -1.0)

floor, as the Float path computes it, with the zero derivative its jumps require — or AdSurrogate's mollified delta train when the AdSurrogateFamily.STAIRCASE family is on.

fract and mod route through here, so a surrogated floor mollifies the sawtooth's −1 jump and mod's too, with no separate site. That is also the reason the staircase family is dangerous: _noise's lattice cell index is a floor, and a gradient through that one is _hash21's, which is chaotic.

    private fun floorAd(a: AD): AD = jumpNode(
        a, kotlin.math.floor(a.value), AdSurrogateFamily.STAIRCASE, toNearestInteger(a.value)
    )

mod(x, y) = x - y * floor(x / y) — the same four Float roundings, in the same places.

    private fun modAd(x: AD, y: AD): AD = sub(x, mul(y, floorAd(div(x, y))))

mix(x, y, a) = x * (1 - a) + y * a — the spec's form, not the x + a*(y-x) rewrite.

    private fun mixAd(x: AD, y: AD, a: AD): AD = add(mul(x, subFromK(1.0, a)), mul(y, a))

max, splitting the derivative evenly at an exact tie, as autodiff.kt does.

max of two masks is a mask, which is what carries or through to a select: or lowers to max(step(0.5, abs(a)), step(0.5, abs(b))). See maskNodes.

    private fun maxAd(a: AD, b: AD): AD = maybeMask(a, b, when {
        a.value > b.value -> node(a, b, a.value, 1.0, 0.0)
        a.value < b.value -> node(a, b, b.value, 0.0, 1.0)
        else -> node(a, b, a.value, 0.5, 0.5)
    })

min, and the and half of the pair maxAd documents.

    private fun minAd(a: AD, b: AD): AD = maybeMask(a, b, when {
        a.value < b.value -> node(a, b, a.value, 1.0, 0.0)
        a.value > b.value -> node(a, b, b.value, 0.0, 1.0)
        else -> node(a, b, a.value, 0.5, 0.5)
    })

Mark out a mask when both reducer inputs were — {0,1} is closed under min and max.

    private fun maybeMask(a: AD, b: AD, out: AD): AD =
        if (trackMasks && isMask(a) && isMask(b)) markMask(out) else out

    private fun clampAd(x: AD, lo: AD, hi: AD): AD = minAd(maxAd(x, lo), hi)

abs, splitting evenly at the origin — except on a mask, where it is the identity.

The exception matters as much as step's does, and for the same reason: select lowers to step(0.5, abs(cond)), so abs sees a mask whose false lanes are exactly 0.0, and the even split reports 0 there. Left alone, the abs would kill every false lane's gradient even after stepAd stopped squaring the kernel. |0| = 0 and |1| = 1, so the honest partial on a mask is 1 and the whole step(0.5, abs(·)) pair passes a gradient straight through.

    private fun absAd(a: AD): AD {
        if (isMask(a)) return markMask(node(a, a.value, 1.0))
        return node(a, a.value.absoluteValue, if (a.value > 0.0) 1.0 else if (a.value < 0.0) -1.0 else 0.0)
    }

    private fun sqrtAd(a: AD): AD {
        val root = kotlin.math.sqrt(a.value)
        return fnode(a, root, if (root == 0.0 || root.isNaN()) 0.0 else 0.5 / root)
    }

    private fun powAd(x: AD, y: AD): AD {
        val p = Math.pow(x.value, y.value)
        val dx = if (x.value == 0.0) 0.0 else y.value * Math.pow(x.value, y.value - 1.0)
        val dy = if (x.value > 0.0) p * kotlin.math.ln(x.value) else 0.0
        return fnode(x, y, p, guard(dx), guard(dy))
    }

    private fun guard(d: Double): Double = if (d.isNaN() || d.isInfinite()) 0.0 else d

smoothstep(e0, e1, x): clamp then t*t*(3 - 2t), operation for operation.

    private fun smoothstepAd(e0: AD, e1: AD, x: AD): AD {
        val t = clampAd(div(sub(x, e0), sub(e1, e0)), constant(0.0), constant(1.0))
        return mul(mul(t, t), subFromK(3.0, mulK(t, 2.0)))
    }

GlslScalars

override fun lift(v: GlslValue): AdValue

    override fun lift(v: GlslValue): AdValue =
        AdValue(v.type, Array(v.c.size) { constant(v.c[it].toDouble()) })

fun liftVariables(v: GlslValue): AdValue

A wire: the free parameters a fit descends over.

    fun liftVariables(v: GlslValue): AdValue =
        AdValue(v.type, Array(v.c.size) { variable(v.c[it].toDouble()) })

override fun typeOf(v: AdValue): GlslType = v.type

override fun zeros(type: GlslType): AdValue

    override fun zeros(type: GlslType): AdValue = AdValue(type, Array(type.components) { constant(0.0) })

override fun ones(type: GlslType): AdValue

    override fun ones(type: GlslType): AdValue = AdValue(type, Array(type.components) { constant(1.0) })

override fun boolOf(b: Boolean): AdValue

    override fun boolOf(b: Boolean): AdValue =
        AdValue(GlslType.BOOL, arrayOf(frozen(if (b) 1.0 else 0.0)))

override fun asBool(v: AdValue): Boolean = v.c[0].value != 0.0

override fun asInt(v: AdValue): Int = Math.round(v.c[0].value.toFloat())

override fun negate(v: AdValue): AdValue = AdValue(v.type, Array(v.c.size) { neg(v.c[it]) })

override fun binary(op: String, a: AdValue, b: AdValue): AdValue

    override fun binary(op: String, a: AdValue, b: AdValue): AdValue = when (op) {
        "+" -> pair(a, b) { x, y -> add(x, y) }
        "-" -> pair(a, b) { x, y -> sub(x, y) }
        "*" -> matmul(a, b)
        "/" -> if (a.type.base == GlslBase.INT && b.type.base == GlslBase.INT) {
            pair(a, b) { x, y -> frozen((x.value.toInt() / y.value.toInt()).toDouble()) }
        } else {
            pair(a, b) { x, y -> div(x, y) }
        }

        "%" -> pair(a, b) { x, y -> modAd(x, y) }
        "^^" -> boolOf(asBool(a) != asBool(b))
        "<" -> boolOf(a.c[0].value < b.c[0].value)
        ">" -> boolOf(a.c[0].value > b.c[0].value)
        "<=" -> boolOf(a.c[0].value <= b.c[0].value)
        ">=" -> boolOf(a.c[0].value >= b.c[0].value)
        "&", "|", "^", "<<", ">>" -> throw GlslUnsupported(
            "bit operator '$op' has no derivative — differentiating a program that packs bits " +
                "(toBits16, fromBits16) is not supported"
        )

        else -> throw GlslUnsupported("binary operator '$op' is not in the StackSL subset")
    }

override fun equalComponents(a: AdValue, b: AdValue): Boolean

    override fun equalComponents(a: AdValue, b: AdValue): Boolean {
        if (a.c.size != b.c.size) return false
        return a.c.indices.all { a.c[it].value.toFloat() == b.c[it].value.toFloat() }
    }

override fun construct(type: GlslType, args: List<AdValue>): AdValue

    override fun construct(type: GlslType, args: List<AdValue>): AdValue {
        if (args.isEmpty()) throw GlslEvalError("$type() needs arguments")
        if (args.size == 1 && args[0].type.isScalar) {
            val v = args[0].c[0]
            if (type.isMatrix) {
                val out = Array(type.components) { constant(0.0) }
                for (i in 0 until type.rows) out[i * type.rows + i] = v
                return AdValue(type, out)
            }
            val broadcast = if (type.base == GlslBase.INT) truncateToInt(v) else v
            return AdValue(type, Array(type.components) { broadcast })
        }
        if (args.size == 1 && args[0].type.isMatrix && type.isMatrix) {
            val src = args[0]
            val out = Array(type.components) { constant(0.0) }
            for (col in 0 until type.cols) for (row in 0 until type.rows) {
                out[col * type.rows + row] = when {
                    col < src.type.cols && row < src.type.rows -> src.c[col * src.type.rows + row]
                    col == row -> constant(1.0)
                    else -> constant(0.0)
                }
            }
            return AdValue(type, out)
        }
        val flat = ArrayList<AD>(type.components)
        args.forEach { a -> a.c.forEach { flat += it } }
        if (flat.size < type.components) {
            throw GlslEvalError("$type(...) got ${flat.size} components, needs ${type.components}")
        }
        return AdValue(type, Array(type.components) {
            if (type.base == GlslBase.INT) truncateToInt(flat[it]) else flat[it]
        })
    }

GLSL's float-to-int conversion truncates towards zero, and quantisation has no derivative.

    private fun truncateToInt(v: AD): AD = frozen(v.value.toFloat().toInt().toDouble())

override fun swizzle(v: AdValue, field: String): AdValue

    override fun swizzle(v: AdValue, field: String): AdValue {
        val idx = swizzleIndices(field, v.c.size)
        return AdValue(GlslType.of(v.type.base, idx.size), Array(idx.size) { v.c[idx[it]] })
    }

override fun component(v: AdValue, index: Int): AdValue

    override fun component(v: AdValue, index: Int): AdValue =
        AdValue(GlslType.of(v.type.base, 1), arrayOf(v.c[index]))

override fun column(v: AdValue, index: Int): AdValue

    override fun column(v: AdValue, index: Int): AdValue {
        val rows = v.type.rows
        return AdValue(GlslType.of(v.type.base, rows), Array(rows) { v.c[index * rows + it] })
    }

override fun withComponents(v: AdValue, indices: IntArray, src: AdValue): AdValue

    override fun withComponents(v: AdValue, indices: IntArray, src: AdValue): AdValue {
        val next = v.c.copyOf()
        for (i in indices.indices) next[indices[i]] = src.c[if (src.c.size == 1) 0 else i]
        return AdValue(v.type, next)
    }

override fun withComponent(v: AdValue, index: Int, src: AdValue): AdValue

    override fun withComponent(v: AdValue, index: Int, src: AdValue): AdValue {
        val next = v.c.copyOf()
        next[index] = src.c[0]
        return AdValue(v.type, next)
    }

override fun coerce(v: AdValue, type: GlslType): AdValue

    override fun coerce(v: AdValue, type: GlslType): AdValue = when {
        v.type == type -> v
        v.type.components == type.components && v.type.cols == type.cols ->
            if (type.base == GlslBase.INT) AdValue(type, Array(v.c.size) { truncateToInt(v.c[it]) })
            else AdValue(type, v.c.copyOf())

        v.type.isScalar -> construct(type, listOf(v))
        else -> throw GlslEvalError("cannot store ${v.type} in a $type")
    }

override fun callBuiltin(name: String, args: List<AdValue>, samplers: GlslSamplers): AdValue?

    override fun callBuiltin(name: String, args: List<AdValue>, samplers: GlslSamplers): AdValue? =
        builtin(name, args)

    private fun builtin(name: String, a: List<AdValue>): AdValue? = when (name) {
        "radians" -> unary(a[0]) { x -> fnode(x, x.value * Math.PI / 180.0, Math.PI / 180.0) }
        "degrees" -> unary(a[0]) { x -> fnode(x, x.value * 180.0 / Math.PI, 180.0 / Math.PI) }
        "sin" -> unary(a[0]) { x -> fnode(x, kotlin.math.sin(x.value), kotlin.math.cos(x.value)) }
        "cos" -> unary(a[0]) { x -> fnode(x, kotlin.math.cos(x.value), -kotlin.math.sin(x.value)) }
        "tan" -> unary(a[0]) { x ->
            val t = kotlin.math.tan(x.value)
            fnode(x, t, 1.0 + t * t)
        }

        "asin" -> unary(a[0]) { x ->
            fnode(x, kotlin.math.asin(x.value), guard(1.0 / kotlin.math.sqrt(1.0 - x.value * x.value)))
        }

        "acos" -> unary(a[0]) { x ->
            fnode(x, kotlin.math.acos(x.value), guard(-1.0 / kotlin.math.sqrt(1.0 - x.value * x.value)))
        }

        "atan" -> if (a.size == 2) {
            pair(a[0], a[1]) { y, x ->
                val d = x.value * x.value + y.value * y.value
                fnode(y, x, kotlin.math.atan2(y.value, x.value), guard(x.value / d), guard(-y.value / d))
            }
        } else {
            unary(a[0]) { x -> fnode(x, kotlin.math.atan(x.value), 1.0 / (1.0 + x.value * x.value)) }
        }

        "sinh" -> unary(a[0]) { x -> fnode(x, kotlin.math.sinh(x.value), kotlin.math.cosh(x.value)) }
        "cosh" -> unary(a[0]) { x -> fnode(x, kotlin.math.cosh(x.value), kotlin.math.sinh(x.value)) }
        "tanh" -> unary(a[0]) { x ->
            val t = kotlin.math.tanh(x.value)
            fnode(x, t, 1.0 - t * t)
        }

        "pow" -> pair(a[0], a[1]) { x, y -> powAd(x, y) }
        "exp" -> unary(a[0]) { x ->
            val e = kotlin.math.exp(x.value)
            fnode(x, e, e)
        }

        "log" -> unary(a[0]) { x -> fnode(x, kotlin.math.ln(x.value), guard(1.0 / x.value)) }
        "exp2" -> unary(a[0]) { x ->
            val e = Math.pow(2.0, x.value)
            fnode(x, e, e * LN2)
        }

        "log2" -> unary(a[0]) { x -> fnode(x, kotlin.math.log2(x.value), guard(1.0 / (x.value * LN2))) }
        "sqrt" -> unary(a[0]) { x -> sqrtAd(x) }
        "inversesqrt" -> unary(a[0]) { x ->
            val r = 1.0 / kotlin.math.sqrt(x.value)
            fnode(x, r, guard(-0.5 * r / x.value))
        }

        "abs" -> unary(a[0]) { x -> absAd(x) }
        // `sign`'s single jump at the origin is the only one in the subset that is not height 1:
        // it goes from -1 to +1, so its distributional derivative is 2·delta.
        "sign" -> unary(a[0]) { x ->
            jumpNode(x, kotlin.math.sign(x.value), AdSurrogateFamily.SIGN, x.value, height = 2.0)
        }

        "floor" -> unary(a[0]) { x -> floorAd(x) }
        "trunc" -> unary(a[0]) { x ->
            val t = x.value.toFloat().let { if (it < 0f) kotlin.math.ceil(it) else kotlin.math.floor(it) }.toDouble()
            // `trunc` is flat across the origin — it is 0 on the whole of (-1, 1) — so the integer
            // at zero is not a jump. An infinite distance is how the kernel is told there is none.
            val nearest = Math.rint(x.value)
            val u = if (nearest == 0.0) Double.POSITIVE_INFINITY else x.value - nearest
            jumpNode(x, t, AdSurrogateFamily.STAIRCASE, u)
        }

        "round" -> unary(a[0]) { x ->
            jumpNode(
                x, Math.round(x.value.toFloat()).toDouble(), AdSurrogateFamily.STAIRCASE,
                toNearestHalfInteger(x.value)
            )
        }

        "roundEven" -> unary(a[0]) { x ->
            jumpNode(x, Math.rint(x.value), AdSurrogateFamily.STAIRCASE, toNearestHalfInteger(x.value))
        }

        "ceil" -> unary(a[0]) { x ->
            jumpNode(
                x, kotlin.math.ceil(x.value), AdSurrogateFamily.STAIRCASE, toNearestInteger(x.value)
            )
        }
        "fract" -> unary(a[0]) { x -> sub(x, floorAd(x)) }
        "mod" -> pair(a[0], a[1]) { x, y -> modAd(x, y) }
        "min" -> pair(a[0], a[1]) { x, y -> minAd(x, y) }
        "max" -> pair(a[0], a[1]) { x, y -> maxAd(x, y) }
        "clamp" -> triple(a) { x, lo, hi -> clampAd(x, lo, hi) }

        // `mix(x, y, bvec)` is a SELECT, not an interpolation, and `_safeDiv` depends on it.
        "mix" -> if (a[2].type.base == GlslBase.BOOL) {
            val n = maxOf(a[0].c.size, a[1].c.size, a[2].c.size)
            AdValue(GlslType.of(GlslBase.FLOAT, n), Array(n) { i ->
                val pick = a[2].c[if (a[2].c.size == 1) 0 else i].value != 0.0
                val src = if (pick) a[1] else a[0]
                src.c[if (src.c.size == 1) 0 else i]
            })
        } else {
            triple(a) { x, y, t -> mixAd(x, y, t) }
        }

        // Three StackSL words lower to `step`: `threshold` is `step(edge, value)`, `step` is
        // itself, and `select`/`if` is `mix(f, t, step(0.5, abs(cond)))` — so this one site is also
        // what lets a gradient reach a `select`'s condition.
        "step" -> pair(a[0], a[1]) { e, x -> stepAd(e, x) }
        "smoothstep" -> triple(a) { e0, e1, x -> smoothstepAd(e0, e1, x) }

        "length" -> scalar(lengthAd(a[0]))
        "distance" -> scalar(lengthAd(pair(a[0], a[1]) { x, y -> sub(x, y) }))
        "dot" -> scalar(dotAd(a[0], a[1]))
        "cross" -> {
            val (p, q) = a
            AdValue(
                GlslType.VEC3,
                arrayOf(
                    sub(mul(p.c[1], q.c[2]), mul(p.c[2], q.c[1])),
                    sub(mul(p.c[2], q.c[0]), mul(p.c[0], q.c[2])),
                    sub(mul(p.c[0], q.c[1]), mul(p.c[1], q.c[0]))
                )
            )
        }

        "normalize" -> {
            val len = lengthAd(a[0])
            AdValue(a[0].type, Array(a[0].c.size) { div(a[0].c[it], len) })
        }

        "faceforward" -> if (dotAd(a[2], a[1]).value < 0.0) a[0] else negate(a[0])
        "reflect" -> {
            val d = dotAd(a[0], a[1])
            pair(a[0], a[1]) { i, n -> sub(i, mul(mulK(d, 2.0), n)) }
        }

        "refract" -> {
            val eta = a[2].c[0]
            val dotni = dotAd(a[1], a[0])
            val k = sub(constant(1.0), mul(mul(eta, eta), subFromK(1.0, mul(dotni, dotni))))
            if (k.value < 0.0) zeros(a[0].type) else {
                val sk = sqrtAd(k)
                pair(a[0], a[1]) { i, n -> sub(mul(eta, i), mul(add(mul(eta, dotni), sk), n)) }
            }
        }

        "lessThan" -> boolPair(a[0], a[1]) { x, y -> x < y }
        "lessThanEqual" -> boolPair(a[0], a[1]) { x, y -> x <= y }
        "greaterThan" -> boolPair(a[0], a[1]) { x, y -> x > y }
        "greaterThanEqual" -> boolPair(a[0], a[1]) { x, y -> x >= y }
        "equal" -> boolPair(a[0], a[1]) { x, y -> x == y }
        "notEqual" -> boolPair(a[0], a[1]) { x, y -> x != y }
        "any" -> boolOf(a[0].c.any { it.value != 0.0 })
        "all" -> boolOf(a[0].c.all { it.value != 0.0 })
        "not" -> AdValue(
            GlslType.of(GlslBase.BOOL, a[0].c.size),
            Array(a[0].c.size) { frozen(if (a[0].c[it].value == 0.0) 1.0 else 0.0) }
        )

        "isnan" -> AdValue(
            GlslType.of(GlslBase.BOOL, a[0].c.size),
            Array(a[0].c.size) { frozen(if (a[0].c[it].value.isNaN()) 1.0 else 0.0) }
        )

        "isinf" -> AdValue(
            GlslType.of(GlslBase.BOOL, a[0].c.size),
            Array(a[0].c.size) { frozen(if (a[0].c[it].value.isInfinite()) 1.0 else 0.0) }
        )

        "matrixCompMult" -> pair(a[0], a[1]) { x, y -> mul(x, y) }
        "transpose" -> {
            val m = a[0]
            val out = Array(m.c.size) { constant(0.0) }
            for (col in 0 until m.type.cols) for (row in 0 until m.type.rows) {
                out[row * m.type.cols + col] = m.c[col * m.type.rows + row]
            }
            AdValue(GlslType.of(GlslBase.FLOAT, m.type.cols, m.type.rows), out)
        }

        "texture", "textureLod", "textureSize", "textureGather" -> throw GlslUnsupported(
            "'$name' has no derivative in this backend: a CPU texture fetch is differentiable in " +
                "principle, but the feedback programs it serves are excluded from every asserted " +
                "comparison here, so there would be no oracle for the answer"
        )

        else -> null
    }

shape plumbing, mirroring GlslValue.kt's

    private inline fun unary(v: AdValue, op: (AD) -> AD): AdValue =
        AdValue(v.type, Array(v.c.size) { op(v.c[it]) })

    private inline fun pair(a: AdValue, b: AdValue, op: (AD, AD) -> AD): AdValue {
        val an = a.c.size
        val bn = b.c.size
        val type = if (an >= bn) a.type else b.type
        val n = maxOf(an, bn)
        if (an != bn && an != 1 && bn != 1) {
            throw GlslEvalError("cannot combine ${a.type} with ${b.type} component-wise")
        }
        return AdValue(type, Array(n) { i -> op(a.c[if (an == 1) 0 else i], b.c[if (bn == 1) 0 else i]) })
    }

    private inline fun triple(a: List<AdValue>, op: (AD, AD, AD) -> AD): AdValue {
        val n = maxOf(a[0].c.size, a[1].c.size, a[2].c.size)
        val type = listOf(a[0].type, a[1].type, a[2].type).maxByOrNull { it.components }!!
        val out = Array(n) { i ->
            op(
                a[0].c[if (a[0].c.size == 1) 0 else i],
                a[1].c[if (a[1].c.size == 1) 0 else i],
                a[2].c[if (a[2].c.size == 1) 0 else i]
            )
        }
        return AdValue(if (type.components == n) type else GlslType.of(GlslBase.FLOAT, n), out)
    }

    private inline fun boolPair(a: AdValue, b: AdValue, op: (Double, Double) -> Boolean): AdValue {
        val n = maxOf(a.c.size, b.c.size)
        return AdValue(GlslType.of(GlslBase.BOOL, n), Array(n) { i ->
            val x = a.c[if (a.c.size == 1) 0 else i].value
            val y = b.c[if (b.c.size == 1) 0 else i].value
            frozen(if (op(x, y)) 1.0 else 0.0)
        })
    }

    private fun scalar(v: AD): AdValue = AdValue(GlslType.FLOAT, arrayOf(v))

dot accumulates in Float, one rounding per multiply and one per add.

    private fun dotAd(a: AdValue, b: AdValue): AD {
        if (a.c.size != b.c.size) throw GlslEvalError("dot(${a.type}, ${b.type})")
        var s = constant(0.0)
        for (i in a.c.indices) s = add(s, mul(a.c[i], b.c[i]))
        return s
    }

    private fun lengthAd(v: AdValue): AD = sqrtAd(dotAd(v, v))

*: matrix times vector, matrix times matrix, and the broadcasts * also has to cover.

    private fun matmul(a: AdValue, b: AdValue): AdValue = when {
        a.type.isMatrix && b.type.isMatrix -> {
            if (a.type.cols != b.type.rows) throw GlslEvalError("${a.type} * ${b.type}")
            val rows = a.type.rows
            val cols = b.type.cols
            val out = Array(rows * cols) { constant(0.0) }
            for (col in 0 until cols) for (row in 0 until rows) {
                var s = constant(0.0)
                for (k in 0 until a.type.cols) {
                    s = add(s, mul(a.c[k * rows + row], b.c[col * b.type.rows + k]))
                }
                out[col * rows + row] = s
            }
            AdValue(GlslType.of(GlslBase.FLOAT, rows, cols), out)
        }

        a.type.isMatrix && b.type.isScalar -> unary(a) { mul(it, b.c[0]) }
        a.type.isScalar && b.type.isMatrix -> unary(b) { mul(it, a.c[0]) }

        a.type.isMatrix -> {
            if (a.type.cols != b.c.size) throw GlslEvalError("${a.type} * ${b.type}")
            val out = Array(a.type.rows) { constant(0.0) }
            for (row in 0 until a.type.rows) {
                var s = constant(0.0)
                for (k in 0 until a.type.cols) s = add(s, mul(a.c[k * a.type.rows + row], b.c[k]))
                out[row] = s
            }
            AdValue(GlslType.of(GlslBase.FLOAT, a.type.rows), out)
        }

        b.type.isMatrix -> {
            if (b.type.rows != a.c.size) throw GlslEvalError("${a.type} * ${b.type}")
            val out = Array(b.type.cols) { constant(0.0) }
            for (col in 0 until b.type.cols) {
                var s = constant(0.0)
                for (k in 0 until b.type.rows) s = add(s, mul(a.c[k], b.c[col * b.type.rows + k]))
                out[col] = s
            }
            AdValue(GlslType.of(GlslBase.FLOAT, b.type.cols), out)
        }

        else -> pair(a, b) { x, y -> mul(x, y) }
    }

    private companion object {
        val LN2 = kotlin.math.ln(2.0)
    }
}