onejs-sl 0.1.0

package/src/encode.ts

@@ -0,0 +1,462 @@
+/**
+ * Turns a recorded program into the buffer the VM shader evaluates.
+ *
+ * Phase 2 of `Specs/SHADER_LANG.md`, and the first consumer of the Phase 1 IR.
+ * Three passes, all here, all pure TypeScript:
+ *
+ *   1. reachable  drop anything the result does not depend on
+ *   2. live ranges  the last instruction that reads each value
+ *   3. linear scan  assign a register, free it when its range ends
+ *
+ * WHY EIGHT REGISTERS. Phase 0 measured a dynamically indexed register file in a
+ * fragment shader on three stacks. On ANGLE D3D11 and on NVIDIA's GL driver,
+ * shrinking the file from 16 entries to 8 costs 3.0x to 4.0x less, which is what
+ * a spill to per thread local memory looks like. Eight is not a guess, it is the
+ * number the hardware asked for.
+ *
+ * WHY AN INDEXED STORE AND NOT AN UNROLLED CHAIN. At 8 registers the two Windows
+ * backends want opposite things: D3D11 is 1.38x to 1.44x FASTER with the indexed
+ * store, GL is about 0.84x faster with the unrolled chain, and both reach the
+ * same best case on their preferred one. Branching on the renderer would take
+ * the best of both and ship a code path that can only ever be tested on one of
+ * the two machines that disagree. One shape everywhere has the better worst
+ * case, 0.199 against 0.237, and no untestable branch.
+ */
+
+import {
+    MAX_TEXTURES, SLError, TYPE, type NodeRef, type Program, type SLNode, type SLType,
+} from "./ir"
+import { INPUT_ID, SLOP } from "./ops"
+import { emitShader } from "./hlsl"
+import { emitGLSL, emitWGSL } from "./web"
+import { uniformDefaults } from "./sl"
+
+// Lives in ops.ts with the other wire constants; re-exported so nothing that
+// reached it through the encoder has to move.
+export { INPUT_ID }
+
+/**
+ * Registers in the VM's file. Phase 0's answer, not a preference.
+ *
+ * A program that needs more is refused rather than spilled. A spill slot is a
+ * second dynamic index into a second array, and any program that needs one is
+ * already past the budget this design is honest about.
+ */
+export const REGISTERS = 8
+
+/** Instructions in one program. The VM's loop is bounded and must be. */
+export const MAX_INSTRUCTIONS = 256
+
+/**
+ * Two texels per instruction, fixed width, so instruction `i` sits at texels
+ * `2i` and `2i+1` and the shader needs no cursor to advance.
+ *
+ * Fixed width wastes a texel on instructions with no immediate. That is the
+ * right trade: variable width means decoding a length and advancing a cursor,
+ * which is a dependent read per instruction and defeats the point of storing the
+ * program in a texture at all.
+ *
+ *   texel 2i     [ op, dst, a, b ]
+ *   texel 2i+1   depends on the op:
+ *                  SWIZZLE   [ c0, c1, c2, c3 ], unused channels are -1
+ *                  COMPOSE   [ c, d, argCount, 0 ], the third and fourth args
+ *                  SAMPLE    [ slot, 0, 0, 0 ]
+ *                  FBM       [ octaves, 0, 0, 0 ]
+ *                  CONST     [ x, y, z, w ]
+ *                  otherwise unused, and zero
+ *
+ * Every value is a float. Integers ride in floats exactly up to 2^24, which is
+ * far above any field here, and the encoder asserts each one is in range rather
+ * than trusting it.
+ */
+export const TEXELS_PER_INSTRUCTION = 2
+
+export interface Encoded {
+    /** Flat RGBA float data, `TEXELS_PER_INSTRUCTION * 4` numbers per instruction. */
+    data: Float32Array
+    instructions: number
+    /** Register holding the result when the program halts. */
+    resultRegister: number
+    /** Registers actually used, for the shader to size nothing and for reporting. */
+    registersUsed: number
+    /**
+     * Uniform names in SLOT ORDER, so a host can set one by name.
+     *
+     * The VM addresses uniforms by slot and knows nothing about names, and the
+     * encoded instructions carry the slot only. Without this the host had a
+     * name from the game and no way to turn it into a slot, so it set a
+     * material property instead and every uniform stayed at zero.
+     */
+    uniforms: string[]
+    /**
+     * Declared uniform defaults, four floats per slot in slot order.
+     *
+     * The generated shader writes these into its Properties block, so a
+     * compiled material starts at them, while the VM's uniform array starts at
+     * zero. Without carrying them across, an unset uniform was its default
+     * after an eject and 0 in the browser: one program, two pictures, nothing
+     * to see in either. The host seeds these right after uploading and lets
+     * the caller's own `uniforms` write over them.
+     */
+    defaults: number[]
+    /**
+     * Texture names in SLOT order, so a host can bind one by name.
+     *
+     * The same reason the uniform names are here: the VM addresses a texture by
+     * slot and the generated shader declares `_Tex0`, so a host handed the name
+     * an author wrote had no way to reach either. Setting a material property
+     * called `grain` bound nothing, in both backends, silently.
+     */
+    textures: string[]
+    hash: string
+    /**
+     * The lowest VM encoding that can run this buffer. A VM older than this
+     * refuses the program; see `SL_WIRE_VERSION`.
+     */
+    wire: number
+    /**
+     * The program as HLSL, for a host that can compile it.
+     *
+     * Lazy, and absent from enumeration: in Play nothing ever reads it, so the
+     * emitter never runs there. In an editor the host asks for it once per
+     * program it has no compiled shader for, records it, and generates the
+     * shader, which is how an ejected game ends up compiled without anybody
+     * writing a manifest.
+     */
+    readonly hlsl: string
+    /**
+     * The program as WGSL and as GLSL ES 3.00, for a browser to compile on
+     * Unity's own device (WebGPU and WebGL2 respectively). Lazy and absent
+     * from enumeration like `hlsl`: a `.sl` import carries them as plain
+     * strings from the build instead, and a host reads only the one its
+     * backend needs.
+     */
+    readonly wgsl: string
+    readonly glsl: string
+}
+
+interface Instr {
+    op: number
+    dst: number
+    a: number
+    b: number
+    imm: [number, number, number, number]
+}
+
+/**
+ * Nodes the result actually depends on, in order.
+ *
+ * Dead nodes are dropped rather than encoded. An author can produce them easily
+ * by computing something and not using it, and the hash already ignores them, so
+ * encoding them would make the buffer disagree with its own hash about what the
+ * program is.
+ */
+export function reachable(nodes: SLNode[], result: NodeRef): NodeRef[] {
+    const keep = new Set<NodeRef>()
+    const stack = [result]
+    while (stack.length > 0) {
+        const ref = stack.pop()!
+        if (keep.has(ref)) continue
+        keep.add(ref)
+        const n = nodes[ref]
+        if (n.k === "swizzle") stack.push(n.src)
+        else if (n.k === "call") for (const a of n.args) stack.push(a)
+    }
+    // Ascending, which is still topological because a node only refers backwards.
+    return [...keep].sort((x, y) => x - y)
+}
+
+/**
+ * The last position in `order` that reads each node.
+ *
+ * A register is free the instant its value's last reader has run, which is what
+ * lets an eight entry file hold a program with far more than eight values in it.
+ */
+export function liveRanges(nodes: SLNode[], order: NodeRef[], result: NodeRef): Map<NodeRef, number> {
+    const pos = new Map<NodeRef, number>()
+    order.forEach((ref, i) => pos.set(ref, i))
+    const last = new Map<NodeRef, number>()
+    order.forEach((ref, i) => {
+        const n = nodes[ref]
+        const reads: NodeRef[] = n.k === "swizzle" ? [n.src] : n.k === "call" ? n.args : []
+        for (const r of reads) last.set(r, i)
+    })
+    // The result outlives every instruction, or it could be freed and overwritten
+    // before anybody reads it.
+    last.set(result, order.length)
+    for (const ref of order) if (!last.has(ref)) last.set(ref, pos.get(ref)!)
+    return last
+}
+
+/**
+ * Names the loops behind an over-long program. A loop unrolls, so a count that
+ * reads as one line of source is emitted once per iteration, and without this
+ * the error blames "instructions" the author never wrote. A loop is named only
+ * when it fills a quarter of the budget or more: a two iteration loop in a
+ * program made long by other code is not the thing to lower.
+ *
+ * Worded for both surfaces: the span comes from `sl.repeat` in the EDSL and
+ * from `for` in a `.sl` file, and an author who wrote one should not be told
+ * about the other.
+ */
+function blameLoops(program: Program, order: NodeRef[]): string {
+    const parts: string[] = []
+    for (const loop of program.loops ?? []) {
+        const ops = order.filter((r) => r >= loop.start && r < loop.end).length
+        if (ops * 4 >= MAX_INSTRUCTIONS) parts.push(`a loop of ${loop.count} accounts for ${ops}`)
+    }
+    if (parts.length === 0) return ""
+    return ` A loop unrolls, so every iteration is emitted in full: ${parts.join(", ")}. ` +
+        `Lower the count or slim the body.`
+}
+
+/**
+ * The program as the VM runs it.
+ *
+ * The IR carries a shape's whole parameter list, up to six. One instruction
+ * holds a shape id and four immediates, so a shape given a fifth or sixth
+ * becomes SDF_WIDE: its first two parameters stay immediates and the other
+ * four come from a constant register. Only the VM sees this; every compiled
+ * backend reads the IR's SDF directly. The hash is the program's own, since
+ * nothing about what it computes changed.
+ */
+export function forVm(program: Program): Program {
+    const wide = (n: SLNode) => n.k === "call" && n.op === SLOP.SDF && (n.imm?.length ?? 0) > 5
+    if (!program.nodes.some(wide)) return program
+    const nodes: SLNode[] = []
+    const map: NodeRef[] = []
+    program.nodes.forEach((n, i) => {
+        if (n.k === "call" && wide(n)) {
+            const im = n.imm!
+            const rest = nodes.push({ k: "const", type: TYPE.VEC4, v: [im[3] ?? 0, im[4] ?? 0, im[5] ?? 0, im[6] ?? 0] }) - 1
+            map[i] = nodes.push({
+                k: "call", type: n.type, op: SLOP.SDF_WIDE, args: [map[n.args[0]!]!, rest], imm: [im[0] ?? 0, im[1] ?? 0, im[2] ?? 0, 0],
+            }) - 1
+            return
+        }
+        if (n.k === "swizzle") map[i] = nodes.push({ ...n, src: map[n.src]! }) - 1
+        else if (n.k === "call") map[i] = nodes.push({ ...n, args: n.args.map((a) => map[a]!) }) - 1
+        else map[i] = nodes.push(n) - 1
+    })
+    const at = (ref: NodeRef) => ref < map.length ? map[ref]! : nodes.length
+    return { ...program, nodes, result: map[program.result]!, loops: program.loops.map((l) => ({ ...l, start: at(l.start), end: at(l.end) })) }
+}
+
+/**
+ * The limits `encode` refuses a program for that are not about registers: its
+ * length and its sampler slots. Shared with `vmFit`, so the two cannot
+ * disagree about what fits.
+ */
+export function checkBudget(program: Program, order: NodeRef[]): void {
+    if (order.length > MAX_INSTRUCTIONS) {
+        throw new SLError(
+            `this program is ${order.length} operations and the VM runs at most ${MAX_INSTRUCTIONS}. ` +
+            `The loop in the shader is bounded and has to be.` + blameLoops(program, order),
+        )
+    }
+    for (const t of program.textures) {
+        if (t.slot >= MAX_TEXTURES) throw new SLError(`texture "${t.name}" is past the sampler budget`)
+    }
+}
+
+/**
+ * A register for every node in `order`, by linear scan, and the most in use
+ * at once. Throws when the file runs out. Shared with `vmFit`.
+ */
+export function allocate(program: Program, order: NodeRef[]): { reg: Map<NodeRef, number>, peak: number } {
+    const last = liveRanges(program.nodes, order, program.result)
+    const reg = new Map<NodeRef, number>()
+    const free: number[] = []
+    for (let r = REGISTERS - 1; r >= 0; r--) free.push(r)
+    /** Which node currently owns each register, for freeing. */
+    const owner = new Array<NodeRef | null>(REGISTERS).fill(null)
+    let peak = 0
+
+    order.forEach((ref, i) => {
+        // Free first, so a value whose last reader is THIS instruction can hand
+        // its register straight to this instruction's result. Without that an
+        // eight register file would behave like a seven register one on every
+        // chained expression, which is most of them.
+        for (let r = 0; r < REGISTERS; r++) {
+            const o = owner[r]
+            if (o !== null && last.get(o)! < i) { owner[r] = null; free.push(r) }
+        }
+
+        const args = readsOf(program.nodes[ref]).map((r) => {
+            const rr = reg.get(r)
+            if (rr === undefined) throw new SLError(`internal: node ${r} has no register when ${ref} needs it`)
+            return rr
+        })
+
+        // Operands are read before the destination is written, so a register
+        // whose last use is this instruction can also be the destination.
+        for (let r = 0; r < REGISTERS; r++) {
+            const o = owner[r]
+            if (o !== null && last.get(o)! === i && !reg.has(ref)) {
+                // Only reclaim a register this instruction is actually reading;
+                // anything else is still live for a later instruction.
+                if (args.includes(r)) { owner[r] = null; free.push(r) }
+            }
+        }
+
+        const dst = free.pop()
+        if (dst === undefined) {
+            throw new SLError(
+                `this program needs more than ${REGISTERS} registers at once. The VM has ${REGISTERS} ` +
+                `because that is what the hardware measured fastest, and spilling is deliberately not ` +
+                `implemented: a spill slot is a second dynamic index into a second array. Split the ` +
+                `program, or reuse fewer intermediate values at the same time.`,
+            )
+        }
+        owner[dst] = ref
+        reg.set(ref, dst)
+        peak = Math.max(peak, REGISTERS - free.length)
+    })
+    return { reg, peak }
+}
+
+function readsOf(n: SLNode): NodeRef[] {
+    return n.k === "swizzle" ? [n.src] : n.k === "call" ? n.args : []
+}
+
+/**
+ * The lowest VM that can run `vm`, the result of `forVm(source)`: 2 only where
+ * a wide shape made SDF_WIDE, so everything else still runs on a wire 1 VM.
+ */
+export function wireOf(source: Program, vm: Program): number {
+    return vm === source ? 1 : 2
+}
+
+export function encode(source: Program): Encoded {
+    const program = forVm(source)
+    const order = reachable(program.nodes, program.result)
+    checkBudget(program, order)
+    const { reg, peak } = allocate(program, order)
+
+    const out = order.map((ref) => {
+        const n = program.nodes[ref]
+        const reads = readsOf(n)
+        // The width of the first operand, for the ops whose result narrows.
+        const srcWidth = reads.length > 0 ? program.nodes[reads[0]].type : undefined
+        return emit(n, reg.get(ref)!, reads.map((r) => reg.get(r)!), srcWidth)
+    })
+
+    const data = new Float32Array(out.length * TEXELS_PER_INSTRUCTION * 4)
+    out.forEach((ins, i) => {
+        const o = i * 8
+        data[o] = ins.op; data[o + 1] = ins.dst; data[o + 2] = ins.a; data[o + 3] = ins.b
+        data[o + 4] = ins.imm[0]; data[o + 5] = ins.imm[1]
+        data[o + 6] = ins.imm[2]; data[o + 7] = ins.imm[3]
+    })
+
+    const encoded = {
+        data,
+        instructions: out.length,
+        resultRegister: reg.get(program.result)!,
+        registersUsed: peak,
+        // Slot order, which is declaration order: Builder.uniform pushes and
+        // uses the resulting index as the slot.
+        uniforms: program.uniforms.map((u) => u.name),
+        defaults: uniformDefaults(program),
+        textures: program.textures.map((t) => t.name),
+        hash: program.hash,
+        wire: wireOf(source, program),
+    } as Encoded
+    let hlsl: string | undefined
+    let wgsl: string | undefined
+    let glsl: string | undefined
+    Object.defineProperty(encoded, "hlsl", {
+        enumerable: false,
+        get: () => (hlsl ??= emitShader(source)),
+    })
+    Object.defineProperty(encoded, "wgsl", {
+        enumerable: false,
+        get: () => (wgsl ??= emitWGSL(source)),
+    })
+    Object.defineProperty(encoded, "glsl", {
+        enumerable: false,
+        get: () => (glsl ??= emitGLSL(source)),
+    })
+    return encoded
+}
+
+const NONE: [number, number, number, number] = [0, 0, 0, 0]
+
+/**
+ * Ops whose result is narrower than their input, so the shader cannot infer how
+ * many components to read from the destination type.
+ */
+const WIDTH_SENSITIVE = new Set<number>([
+    SLOP.LENGTH, SLOP.DISTANCE, SLOP.DOT, SLOP.NORMALIZE, SLOP.LUMINANCE,
+])
+
+function emit(n: SLNode, dst: number, args: number[], srcWidth?: SLType): Instr {
+    switch (n.k) {
+        case "const": {
+            const v = n.v
+            return { op: SLOP.CONST, dst, a: 0, b: n.type, imm: [v[0] ?? 0, v[1] ?? 0, v[2] ?? 0, v[3] ?? 0] }
+        }
+        case "input":
+            return { op: SLOP.INPUT, dst, a: INPUT_ID[n.name], b: n.type, imm: [...NONE] }
+        case "uniform":
+            return { op: SLOP.UNIFORM, dst, a: n.slot, b: n.type, imm: [...NONE] }
+        case "swizzle": {
+            const c = n.chans
+            return {
+                op: SLOP.SWIZZLE, dst, a: args[0], b: n.type,
+                imm: [c[0] ?? -1, c[1] ?? -1, c[2] ?? -1, c[3] ?? -1],
+            }
+        }
+        case "call": {
+            if (n.op === SLOP.COMPOSE) {
+                if (args.length > 4) throw new SLError(`compose takes at most 4 parts, got ${args.length}`)
+                return {
+                    op: SLOP.COMPOSE, dst, a: args[0] ?? 0, b: args[1] ?? 0,
+                    imm: [args[2] ?? 0, args[3] ?? 0, args.length, n.type],
+                }
+            }
+            // Ops whose RESULT is narrower than their input need the source
+            // width, because a register is a float4 whatever the value in it is
+            // and length(vec2) is not length of four components. Found by
+            // writing the shader, not by writing the encoder: nothing in the IR
+            // is wrong, the information simply was not reaching the consumer.
+            if (WIDTH_SENSITIVE.has(n.op)) {
+                const src = srcWidth
+                if (src === undefined) throw new SLError(`internal: ${n.op} has no source width`)
+                return { op: n.op, dst, a: args[0] ?? 0, b: args[1] ?? 0, imm: [src, 0, 0, 0] }
+            }
+            // SDF carries a shape id plus four parameters, which is one more
+            // number than the immediate holds. The id goes in the second
+            // operand slot, free because SDF reads one register.
+            if (n.op === SLOP.SDF) {
+                const im = n.imm ?? []
+                return {
+                    op: SLOP.SDF, dst, a: args[0] ?? 0, b: im[0] ?? 0,
+                    imm: [im[1] ?? 0, im[2] ?? 0, im[3] ?? 0, im[4] ?? 0],
+                }
+            }
+            if (args.length > 3) {
+                throw new SLError(`internal: ${n.op} has ${args.length} operands and the encoding holds 3`)
+            }
+            // A third operand rides in the immediate, which is free for every op
+            // that has one (clamp, mix, smoothstep, select) because none of them
+            // carries a literal.
+            const imm: [number, number, number, number] = n.imm !== undefined
+                ? [n.imm[0] ?? 0, n.imm[1] ?? 0, n.imm[2] ?? 0, n.imm[3] ?? 0]
+                : [args[2] ?? 0, 0, 0, 0]
+            return { op: n.op, dst, a: args[0] ?? 0, b: args[1] ?? 0, imm }
+        }
+    }
+}
+
+/**
+ * Every field must survive a float32 round trip exactly, or an instruction
+ * decodes as a different instruction. Cheap to check and impossible to debug
+ * from the wrong picture it would otherwise produce.
+ */
+export function checkEncodable(v: number, what: string): void {
+    if (!Number.isFinite(v)) throw new SLError(`${what} is ${v}, which cannot be encoded`)
+    if (Number.isInteger(v) && Math.abs(v) > 0x1000000) {
+        throw new SLError(`${what} is ${v}, past the 2^24 an integer survives in a float`)
+    }
+}