StackSLFit.kt

src/jvmMain/kotlin/afield/lang/stacksl/cpu/StackSLFit.kt · 392 lines

package afield.lang.stacksl.cpu

import afield.math.autodiff.AD
import afield.math.autodiff.ADGradient
import afield.math.autodiff.ADTape
import afield.math.autodiff.OptimizeMethod
import afield.math.autodiff.OptimizeOptions
import afield.math.autodiff.OptimizeRun
import afield.math.autodiff.minimizeWithGradient
import afield.math.autodiff.sq
import java.util.stream.IntStream

Gradient descent on a StackSL program's own literals: fit a shader to a picture, then write the fitted numbers back into the program.

Why a StackSL program is differentiable

StackSL is not geometry. It is a per-pixel field: main() runs once per pixel with fragUV as its only spatial input, and the output is an ordinary composition of arithmetic and GLSL builtins. So d(pixel)/d(wire) is the chain rule. There is no rasteriser in the pipeline, no coverage term, no approximation and so no bias to characterise.

RectSL, a shape language, is the harder case. There a pixel's value depends on a parameter through coverage, so d(pixel)/d(parameter) needs a differentiable rasteriser, and a coverage approximation is a modelling choice with a bias to characterise. RectSLSoftRaster.kt is that choice, made there.

Limits on what a fit can do

A fit can only move literal-backed wires. See StackSLWires: a wire noted "optional default palette a" is a default the emitter supplied, with no literal in the source to rewrite. StackSLFitResult.fittedIndices is the population that moved and WiredProgram.wires is the larger set that exists.
The template clamps. StackSLCpu.stillTemplate, the default, ends main() with outColor = vec4(clamp(c, 0.0, 1.0), 1.0), and clamp carries zero derivative outside its range. A pixel that is saturated at the seed therefore contributes nothing to the gradient, so a seed whose image is mostly blown out is a seed with almost no signal. StackSLFitResult.saturatedFraction measures how much of the image is in that state.
The default loss is a per-pixel sum of squares and rewards exactly that. It has no term for structure, so a fit can lower it a long way by matching a target's colour and luminance statistics while the image does not resemble the target at all. StackSLFitFigureTest writes a figure of the target, the seed and the fit side by side, which is the comparison the loss itself cannot make. StackSLFitLoss.LPIPS descends on the perceptual distance instead, with the gradient of the LPIPS network chained through the program (StackSLLpipsLoss); StackSLFitLoss.LPIPS_PIXEL on the mean of the two.

`object StackSLFit`

The fit itself: seed program in, fitted program out, with the descent's provenance attached.

fit is the whole surface; StackSLFitResult carries the fitted literals, the loss trajectory and the re-emitted source.

object StackSLFit {

`const val DEFAULT_DIM = 64`

The default fit resolution: 64 square, sized from the measured cost — see dev_docs.

`fun fit`

Fit program's literals to target.

target must be dim square. indices selects which wires may move; the default is every literal-backed one, which is the set StackSLWires.rewrite can write back.

annealing turns on the jump surrogate (AdSurrogate) with a cooling temperature, which is what lets a descent move a wire that sits behind a threshold, a posterize, a select or a bit-crush — operations whose exact gradient is zero everywhere a fit will ever stand. It is null by default, so a fit is exact unless asked otherwise, and dev_docs/STACKSL_SURROGATE.md says which family to anneal for which program: the AdSurrogateFamily.STAIRCASE family also softens _noise's lattice index, where zero is the true derivative, so switching everything on costs a noise-frequency wire its accuracy.

loss picks what is descended on: StackSLFitLoss.MSE by default, or LPIPS through StackSLLpipsLoss at lpipsSide (a multiple of dim) with an optional lpipsGrid of region terms. StackSLFitResult.seedLoss and StackSLFitResult.fittedLoss are then that loss's values.

body — 62 lines, from line 79

    fun fit(
        program: String,
        target: CpuRaster,
        template: String = StackSLCpu.stillTemplate(),
        dim: Int = DEFAULT_DIM,
        iterations: Int = 60,
        learningRate: Double = 0.05,
        method: OptimizeMethod = OptimizeMethod.ADAM,
        slot: String = "program",
        indices: List<Int>? = null,
        parallel: Boolean = true,
        annealing: AdAnnealing? = null,
        loss: StackSLFitLoss = StackSLFitLoss.MSE,
        lpipsSide: Int = maxOf(StackSLLpipsLoss.DEFAULT_SIDE / dim, 1) * dim,
        lpipsGrid: Int = 0
    ): StackSLFitResult {
        require(target.width == dim && target.height == dim) {
            "target is ${target.width}x${target.height}, fit is ${dim}x$dim"
        }
        val wired = StackSLWires.lift(program, slot)
        val image = StackSLImageLoss(wired, template, dim, target, parallel)
        val lpips = if (loss == StackSLFitLoss.MSE) null else StackSLLpipsLoss(image, lpipsSide, lpipsGrid, loss)
        val moving = indices ?: wired.literalWires.map { it.index }
        require(moving.isNotEmpty()) {
            "'$program' has no literal-backed wire to fit: ${wired.describe()}"
        }
        val full = wired.initial
        fun expand(p: DoubleArray): DoubleArray {
            val out = full.copyOf()
            moving.forEachIndexed { i, w -> out[w] = p[i] }
            return out
        }

        val run = minimizeWithGradient(
            DoubleArray(moving.size) { full[moving[it]] },
            OptimizeOptions(
                method = method,
                iterations = iterations,
                learningRate = learningRate,
                showProgress = false
            )
        ) { p ->
            // One `next()` per loss evaluation: the schedule's own KDoc says why that is the driver
            // and what LBFGS does to it. With `annealing == null` this is the exact gradient.
            val surrogate = annealing?.next() ?: AdSurrogate.OFF
            val whole = lpips?.valueAndGradient(expand(p), surrogate) ?: image.valueAndGradient(expand(p), surrogate)
            ADGradient(whole.value, DoubleArray(moving.size) { whole.gradient[moving[it]] })
        }

        val fitted = expand(run.x)
        return StackSLFitResult(
            wired = wired,
            loss = image,
            fittedIndices = moving,
            fittedWires = fitted,
            run = run,
            fittedProgram = StackSLWires.rewrite(program, wired.wires, fitted),
            annealing = annealing,
            lossKind = loss
        )
    }
}

`class StackSLImageLoss`

A mean-squared-error image loss over a wired program's parameter vector, with both an exact reverse-mode gradient and the plain Float render the gradient is checked against.

The shader is parsed once; a new parameter vector costs one setGlobal per wire, because StackSLWires.lower turned the wire array into scalar globals. Alpha is excluded from the loss: the gallery template writes a constant 1.0 there, and a constant contributes nothing but dilution.

class StackSLImageLoss(
    val wired: WiredProgram,
    template: String,
    val dim: Int,
    val target: CpuRaster,
    private val parallel: Boolean = true,
    private val outName: String = "outColor",
    private val fragUvName: String = "fragUV"
) {

`val source: String = StackSLCpu.substituteSlot(template, wired.glsl, wired.slot)`

The complete fragment shader, wires and all — the artifact both paths read.

    private val shader = parseGlslShader(source)

`init`

    init {
        require(target.width == dim && target.height == dim) {
            "target is ${target.width}x${target.height}, loss is ${dim}x$dim"
        }
    }

`fun render(values: DoubleArray): CpuRaster`

Render on the Float reference path at values.

    fun render(values: DoubleArray): CpuRaster =
        StackSLCpu.renderFragment(source, dim, dim, wired.globals(values), parallel = parallel)

`fun value(values: DoubleArray): Double = meanSquaredError(render(values), target)`

The loss at values, via the Float path — the finite-difference oracle's evaluator.

`fun valueAndGradient(values: DoubleArray, surrogate: AdSurrogate = AdSurrogate.OFF): ADGradient`

The loss and its exact gradient at values, by reverse-mode AD through the same parse tree the Float path walks.

One ADTape per pixel: the loss is a sum over pixels, so the gradient is a sum of per-pixel backward sweeps, and no tape outlives the pixel that built it. Rows run in parallel into private accumulators, so the sum is over a fixed partition and does not depend on thread scheduling — the answer is the same number every time, which a test is entitled to assume.

surrogate is the one knob that makes the gradient inexact, and only the gradient: an AdSurrogate other than AdSurrogate.OFF replaces the exact zero partial at a jump discontinuity with a mollified delta, leaving every primal value alone. So ADGradient.value is the true loss at both settings — it is the direction that is a surrogate. Default off, and every existing caller therefore gets exactly what it got before.

    fun valueAndGradient(values: DoubleArray, surrogate: AdSurrogate = AdSurrogate.OFF): ADGradient {
        val (total, gradient) = sweep(values, surrogate) { x, y, out, tape ->
            var pixel: AD = tape.constant(0.0)
            for (ch in 0 until CHANNELS) {
                pixel += (out.c[ch] - target.at(x, y, ch).toDouble()).sq()
            }
            pixel
        }
        val scale = 1.0 / (dim.toDouble() * dim.toDouble() * CHANNELS)
        for (k in gradient.indices) gradient[k] *= scale
        return ADGradient(total * scale, gradient)
    }

`fun seededGradient(values: DoubleArray, seed: FloatArray, surrogate: AdSurrogate = AdSurrogate.OFF): DoubleArray`

The vector-seeded backward pass: sum over pixels and channels of seed[ch, y, x] * d outch/ d wire, for an image-space gradient seed that something else computed.

This is how a loss the tape cannot see is descended on. For a loss L(image), pass seed = dL / d image and the result is dL / d wire by the chain rule: each pixel's tape is swept backward from sum_ch seed[ch, y, x] * out.cch` instead of from its squared error.outis the template's output *after* itsclamp, so a clamped channel passes back zero, as it should.seedis planar,3 x dim x dim, channel-major, row 0 at the top — the layoutLpips.distanceAndGradientPlanaruses; see [StackSLLpipsLoss.blockSum](../../../../../../../reference/source/jvmMain/afield/lang/stacksl/cpu/StackSLLpipsLoss.md#L144) for a seed computed on an enlarged image. Withseed = 2 (out - target) / (3 dim^2)this is exactly the [valueAndGradient](#L193) gradient, whichStackSLLpipsLossTest` asserts.

    fun seededGradient(values: DoubleArray, seed: FloatArray, surrogate: AdSurrogate = AdSurrogate.OFF): DoubleArray {
        require(seed.size == CHANNELS * dim * dim) { "seed holds ${seed.size} floats, expected ${CHANNELS * dim * dim}" }
        val plane = dim * dim
        return sweep(values, surrogate) { x, y, out, tape ->
            var pixel: AD = tape.constant(0.0)
            for (ch in 0 until CHANNELS) {
                val s = seed[ch * plane + y * dim + x].toDouble()
                if (s != 0.0) pixel += out.c[ch] * s
            }
            pixel
        }.second
    }

One backward sweep per pixel of pixelLoss, summed: the sum of the per-pixel values and the sum of their gradients with respect to every wire, unscaled.

source — 43 lines, from line 236

    private fun sweep(
        values: DoubleArray,
        surrogate: AdSurrogate,
        pixelLoss: (x: Int, y: Int, out: AdValue, tape: ADTape) -> AD
    ): Pair<Double, DoubleArray> {
        val n = wired.wires.size
        require(values.size == n) { "expected $n wire values, got ${values.size}" }
        val perRow = arrayOfNulls<DoubleArray>(dim)
        val rows = IntStream.range(0, dim)
        val stream = if (parallel) rows.parallel() else rows
        stream.forEach { y ->
            val acc = DoubleArray(n + 1)
            for (x in 0 until dim) {
                val tape = ADTape()
                val scalars = GlslAdScalars(tape, surrogate)
                val wireNodes = Array(n) { scalars.variable(values[it]) }
                val globals = HashMap<String, AdValue>(n + 1)
                for (k in 0 until n) {
                    globals[StackSLWires.globalName(wired.prefix, k)] =
                        AdValue(GlslType.FLOAT, arrayOf(wireNodes[k]))
                }
                globals[fragUvName] = scalars.lift(StackSLCpu.fragUvAt(x, y, dim, dim))
                val machine = shader.newMachine(scalars, globals, GlslSamplers.REFUSE)
                if (!machine.runMain()) {
                    throw GlslEvalError("pixel ($x, $y) executed discard; a fit has no answer for that")
                }
                val out = machine.global(outName) ?: throw GlslEvalError("shader never wrote '$outName'")
                val pixel = pixelLoss(x, y, out, tape)
                val adjoint = tape.backward(pixel)
                acc[0] += pixel.value
                for (k in 0 until n) acc[k + 1] += adjoint[wireNodes[k].index]
            }
            perRow[y] = acc
        }
        var total = 0.0
        val gradient = DoubleArray(n)
        perRow.forEach { row ->
            val r = row!!
            total += r[0]
            for (k in 0 until n) gradient[k] += r[k + 1]
        }
        return total to gradient
    }

`fun renderPrimal(values: DoubleArray, surrogate: AdSurrogate = AdSurrogate.OFF): CpuRaster`

The primal image the differentiable pass computed, as a raster.

Not a convenience: this is the evidence that the AD pass is evaluating the same shader. Every node GlslAdScalars pushes rounds its value to Float where the Float path rounds, so this raster is expected to be bit-identical to render's, and StackSLGradientTest asserts exactly that over a named population.

surrogate is accepted here for one reason: the surrogate's central claim is that it changes nothing in the forward pass, and a claim that cannot be measured is a hope. StackSLSurrogateTest renders this at every temperature it sweeps and asserts bit-equality, so a future surrogate that reaches into a primal value reddens instead of drifting.

    fun renderPrimal(values: DoubleArray, surrogate: AdSurrogate = AdSurrogate.OFF): CpuRaster {
        val n = wired.wires.size
        val rgba = FloatArray(dim * dim * 4)
        for (y in 0 until dim) {
            val tape = ADTape()
            val scalars = GlslAdScalars(tape, surrogate)
            val globals = HashMap<String, AdValue>(n + 1)
            for (k in 0 until n) {
                globals[StackSLWires.globalName(wired.prefix, k)] =
                    AdValue(GlslType.FLOAT, arrayOf(scalars.constant(values[k])))
            }
            for (x in 0 until dim) {
                globals[fragUvName] = scalars.lift(StackSLCpu.fragUvAt(x, y, dim, dim))
                val machine = shader.newMachine(scalars, globals, GlslSamplers.REFUSE)
                machine.runMain()
                val out = machine.global(outName)!!.primal()
                val base = (y * dim + x) * 4
                for (ch in 0 until minOf(4, out.c.size)) rgba[base + ch] = out.c[ch]
            }
        }
        return CpuRaster(dim, dim, rgba)
    }

`fun saturatedFraction(values: DoubleArray): Double`

Fraction of RGB samples that sit exactly on a clamp bound, and so carry no gradient.

    fun saturatedFraction(values: DoubleArray): Double {
        val raster = render(values)
        var saturated = 0
        for (i in 0 until dim * dim) for (ch in 0 until CHANNELS) {
            val v = raster.rgba[i * 4 + ch]
            if (v <= 0f || v >= 1f) saturated++
        }
        return saturated.toDouble() / (dim.toDouble() * dim * CHANNELS)
    }

`companion object`

    companion object {

`const val CHANNELS = 3`

RGB. Alpha is a constant in every gallery template, so comparing it is decoration.

`fun meanSquaredError(a: CpuRaster, b: CpuRaster): Double`

Mean squared error over the three colour channels of every pixel.

        fun meanSquaredError(a: CpuRaster, b: CpuRaster): Double {
            require(a.width == b.width && a.height == b.height) { "raster sizes differ" }
            var sum = 0.0
            for (i in 0 until a.width * a.height) for (ch in 0 until CHANNELS) {
                val d = (a.rgba[i * 4 + ch] - b.rgba[i * 4 + ch]).toDouble()
                sum += d * d
            }
            return sum / (a.width.toDouble() * a.height * CHANNELS)
        }
    }
}

`class StackSLFitResult`

What a fit produced, including the program a human can paste back into a box.

class StackSLFitResult(
    /** The wired program that was fitted. */
    val wired: WiredProgram,
    /** The image loss: the renderer, and the MSE whatever [lossKind] the descent ran on. */
    val loss: StackSLImageLoss,
    /** The wire indices that were allowed to move, in the optimiser's own order. */
    val fittedIndices: List<Int>,
    /** The whole wire vector after the fit, indexed by wire index. */
    val fittedWires: DoubleArray,
    /** The descent's provenance: losses, trajectory, gradient norms. */
    val run: OptimizeRun,
    /** [wired]'s source with the fitted literals substituted at their own token positions. */
    val fittedProgram: String,
    /**
     * The jump surrogate's schedule, or `null` for an exact fit.
     *
     * When this is non-null, [OptimizeRun.gradient] is the **surrogate** direction at the final
     * temperature rather than the exact gradient — the loss values are exact either way, because
     * the surrogate only ever changes a partial. Kept so that a reading of a run knows which it is
     * looking at rather than inferring it from the numbers.
     */
    val annealing: AdAnnealing? = null,
    /** What the descent ran on; [seedLoss] and [fittedLoss] are its values. */
    val lossKind: StackSLFitLoss = StackSLFitLoss.MSE
) {

`val seedWires: DoubleArray get() = wired.initial`

The wire vector the fit started from.

`val seedLoss: Double get() = run.initialValue`

The loss before any step.

`val fittedLoss: Double get() = run.value`

The loss at the end.

`val lossRatio: Double get()`

How far the loss fell, as a ratio. 1.0 means the descent achieved nothing.

    val lossRatio: Double get() = if (run.value == 0.0) Double.POSITIVE_INFINITY else run.initialValue / run.value

`val saturatedFraction: Double get() = loss.saturatedFraction(seedWires)`

Fraction of RGB samples clamped at the seed — see StackSLFit's second caveat.

`fun report(): List<String>`

One line per moved literal: what the human wrote, and what the fit made of it.

    fun report(): List<String> = fittedIndices.map { w ->
        val wire = wired.wires[w]
        "wire #$w  '${wire.token}' at ${wire.line}:${wire.column}  " +
            "${wire.initial} -> ${StackSLWires.formatLiteral(fittedWires[w])}"
    }
}