RectSLSimplify.kt

src/jvmMain/kotlin/afield/lang/rectsl/RectSLSimplify.kt · 164 lines

package afield.lang.rectsl

import kotlin.math.abs

An algebraic simplifier over the track IR, under a requirement that every rewrite be bit-exact.

The IR is data classes, so two tracks built by the same array body compare ==. The GLSL emitter's function deduplication rests on that, keying emitted functions by the track itself so that a hundred array copies share one function, and so does the min(x, x) rule below, which compares two subtrees directly. The narrower rewrites over the same IR are describe(), which prints a track, and substituteParameters, which sets named free parameters to numbers; this is the general one.

Nothing in afield.lang.rectsl writes a track out or reads one back, so a track cannot be serialized, although being data rather than closures would allow it.

What it is for

Not speed, primarily. RectSL's array words build per-copy channels by offsetting a constant, so a hundred-copy array arrives at a backend as a hundred BinaryDoubleTrack(const, const, ADD) nodes, and each backend has had to fold that itself: trackConstantValue exists because the GLSL emitter needed it, and the Fn lowering calls the same function for the same reason. A fold is the smallest rewrite there is, and both emitters had to reach for it themselves because there was nowhere over the IR to apply it once.

The bigger reason is that a simplified tree is a smaller Fn graph, and an Fn graph's size is the cost of both a sample and a gradient — Fn.eval walks every node, and gradientAt records every node on a tape. A fit runs that thousands of times.

The contract: every rule is bit-exact

A simplified track must sample bit-identically to the original, not merely equal-to-tolerance, which rules out several rewrites everyone writes by reflex. The gallery corpus caught the first one: 0 - x rewritten to -x disagrees at x = 0, because 0.0 - 0.0 is +0.0 and -(0.0) is -0.0.

A signed zero is not a curiosity in this IR. HeadingTrack takes atan2 of a velocity, and atan2(+0.0, -1.0) is +pi while atan2(-0.0, -1.0) is -pi — half a turn apart. A rewrite that "only" changes the sign of a zero can therefore spin a rectangle through 180 degrees. So the rules kept here are the ones exact for every double including both zeros, and the reflexive ones are dropped with their reasons:

Losing the additive identity costs less than it looks like it should: RectSL's array words build per-copy channels by offsetting a constant, and the constant fold takes those whole.

Free parameters are never folded, for the same reason trackConstantValue takes a flag: a FreeParamTrack holds a number but is not a number, and folding it away would delete the thing a fit is meant to solve for.

How it is checked

RectSLSimplifyTest runs it over the gallery and requires every rewritten channel to sample bit-identically to the original at 41 times each, and to be a fixed point of a second pass. The node-count reduction it prints is reported rather than asserted at any figure, because that figure moves with the gallery and with the IR — a key value became a node when TcbKey began holding a track, for instance. Run the test to read the current one.

A plausible-looking rule that is wrong in one corner is this file's characteristic failure, and bit-identity over the corpus already assembled for the differentiable backend is what catches it.

fun Track<Double>.nodeCount(): Int = 1 + trackChildren(this).sumOf { it.nodeCount() }

Total nodes in this track's tree, counting a shared subtree once per path — the size a backend pays.

fun Track<Double>.simplify(): Track<Double> = mapTrack(this) { node -> rewriteNode(node) }

This track with constants folded and the algebraic identities applied, bottom-up.

A pure function returning a new tree; nothing mutates. Sampling is preserved bit-for-bit on the whole gallery corpus (2754 channels, 16.5% fewer nodes, 112,914 samples compared) — see this file's KDoc for how that is checked, and for the table of rewrites left out, with their reasons.

True when track is the constant value, exactly. A free parameter is never "a constant".

private fun isConstant(track: Track<Double>, value: Double): Boolean =
    trackConstantValue(track, foldFreeParameters = false) == value

private fun rewriteNode(track: Track<Double>): Track<Double> {
    // The fold first: it subsumes every rule below for a subtree that does not vary with time, and it
    // is what turns a hundred-copy array's arithmetic into a hundred literals.
    trackConstantValue(track, foldFreeParameters = false)?.let { folded ->
        // A track that already *is* a bare constant is left alone rather than rebuilt, so simplify
        // is idempotent on its own output by identity and not merely by value.
        if (track is ConstantTrack<*>) return track
        return constantTrack(folded)
    }

    return when (track) {
        is BinaryDoubleTrack -> when (track.op) {
            // `x - 0` is the one additive identity that survives the signed-zero rule: subtracting
            // positive zero returns x untouched for every x, where *adding* it turns -0.0 into +0.0.
            TrackOp2.SUB -> if (isConstant(track.b, 0.0)) track.a else track
            TrackOp2.MUL -> when {
                isConstant(track.b, 1.0) -> track.a
                isConstant(track.a, 1.0) -> track.b
                else -> track
            }
            TrackOp2.DIV -> if (isConstant(track.b, 1.0)) track.a else track
            TrackOp2.POW -> if (isConstant(track.b, 1.0)) track.a else track
            // Structural equality doing real work: `min(x, x)` is `x` whatever x is.
            TrackOp2.MIN, TrackOp2.MAX -> if (track.a == track.b) track.a else track
            else -> track
        }

        is TernaryDoubleTrack -> when (track.op) {
            // No `mix` rule at all, which is the surprise. `mix(a, b, t)` evaluates as
            // `a + (b - a) * t`: at t = 1 that is `a + (b - a)`, which is simply not `b` in floating
            // point, and at t = 0 or with a == b it ends in `a + (+/-0.0)`, which is not `a` when a is
            // a negative zero. Every version of this rewrite is wrong somewhere.
            TrackOp3.CLAMP -> {
                val lo = trackConstantValue(track.b, foldFreeParameters = false)
                val hi = trackConstantValue(track.c, foldFreeParameters = false)
                // Equal bounds pin the value; and a clamp of something already inside its own bounds
                // cannot be recognised in general, so that case is left alone.
                if (lo != null && hi != null && lo == hi) constantTrack(lo) else track
            }
            else -> track
        }

        is MappedDoubleTrack -> {
            val inner = track.source
            when {
                // Involutions and idempotents. `neg neg` cancels; `abs` and `saturate` are already
                // themselves after one application.
                track.op == TrackOp1.NEG && inner is MappedDoubleTrack && inner.op == TrackOp1.NEG ->
                    inner.source
                track.op == TrackOp1.ABS && inner is MappedDoubleTrack && inner.op == TrackOp1.ABS ->
                    inner
                track.op == TrackOp1.SATURATE && inner is MappedDoubleTrack &&
                    inner.op == TrackOp1.SATURATE -> inner
                else -> track
            }
        }

        is TimeWarpTrack -> when (track.warp) {
            TimeWarpKind.SHIFT -> if (track.amount == 0.0) track.source else track
            TimeWarpKind.SCALE -> if (track.amount == 1.0) track.source else track
            // A degenerate period is the identity in both the CPU sampler and the GLSL, so the node
            // is not doing anything.
            TimeWarpKind.LOOP -> if (abs(track.amount) <= 1e-12) track.source else track
            TimeWarpKind.PINGPONG -> if (abs(track.amount * 2.0) <= 1e-12) track.source else track
            else -> track
        }

        else -> track
    }
}