@woosh/meep-engine 3.15.0 → 3.16.0

package/src/shade/renderer/particles/optimizer/optimize_particle_program.js

@@ -0,0 +1,478 @@
+import { assert } from "../../../../core/assert.js";
+import { PARTICLE_VM_INSTRUCTION_WORDS, PARTICLE_VM_REGISTER_SLOTS } from "../ParticleConstants.js";
+import { InstructionStream } from "../isa/InstructionStream.js";
+import { ParticleAssembler } from "../isa/ParticleAssembler.js";
+import { ParticleProgram } from "../isa/ParticleProgram.js";
+import { VM_OP, vm_decode_op, vm_decode_width } from "../isa/ParticleVMISA.js";
+import { PARTICLE_PASS_DEAD_CODE, pass_eliminate_dead_code } from "./pass_eliminate_dead_code.js";
+import {
+    PARTICLE_PASS_DEDUP_CONSTANTS,
+    pass_dedup_constants,
+    visit_particle_constant_reads,
+} from "./pass_dedup_constants.js";
+import { PARTICLE_PASS_FOLD_CONSTANTS, pass_fold_constants } from "./pass_fold_constants.js";
+import { PARTICLE_PASS_NARROW_WIDTHS, pass_narrow_widths } from "./pass_narrow_widths.js";
+import { PARTICLE_PASS_PACK_REGISTERS, pass_pack_registers } from "./pass_pack_registers.js";
+import { PARTICLE_PASS_PEEPHOLE, pass_peephole } from "./pass_peephole.js";
+import { PARTICLE_PASS_PROPAGATE_COPIES, pass_propagate_copies } from "./pass_propagate_copies.js";
+import { PARTICLE_PASS_SHARE_SUBEXPRESSIONS, pass_share_subexpressions } from "./pass_share_subexpressions.js";
+
+/**
+ * Binary-level optimizer for particle VM programs: instructions in, instructions out.
+ *
+ * It runs over a compiled {@link InstructionStream} rather than inside any one frontend, because
+ * there is more than one frontend — the node-graph compiler, the assembly text pane in the editor,
+ * and whatever comes next — and the same waste appears in all of them. The node compiler lowers a
+ * value at a time and has no way to know that the broadcast it just emitted will be read one lane
+ * wide; a pass over the finished stream can simply look.
+ *
+ * ## Shape
+ *
+ * Passes are ordinary functions from an instruction array to a new one, with `null` for "changed
+ * nothing". They never reorder: every pass rewrites an instruction in place or drops it, which is
+ * what makes the ordering constraints — `RANDOM`'s draw sequence, `STORE_ATTR`/`LOAD_ATTR`
+ * aliasing — hold structurally rather than by each pass remembering them.
+ *
+ * The pipeline sweeps the enabled passes in order until a whole sweep changes nothing or
+ * `max_iterations` sweeps have run, then packs the register file once. Passes that pay to repeat on
+ * their own take their own `max_iterations`, so effort is the caller's to spend.
+ *
+ * ## What must not change
+ *
+ * A program's meaning is the record it leaves, the kill flag it sets, and the RNG state it ends on.
+ * Every pass here preserves all three to the bit, with one deliberate exception: folding `x * 0` to
+ * `+0` gives up the sign of that zero, and gives up `NaN` where something upstream overflowed. It is
+ * on by default — see the argument in `pass_fold_constants.js` — and
+ * `{ passes: { fold_constants: { fold_multiply_by_zero: false } } }` takes the optimizer back to
+ * bit-identical. `optimize_particle_program.spec.js` and the differential harness beside it hold
+ * both settings to their respective promises, over the real distribution compositions, the
+ * hand-written spec programs and randomly generated ones.
+ */
+
+/**
+ * @callback ParticleOptimizerPass
+ * @param {number[][]} instructions five-word tuples, in program order
+ * @param {ParticleOptimizerContext} context
+ * @returns {number[][]|null} the rewritten stream, or `null` when the pass changed nothing
+ */
+
+/**
+ * @typedef {object} ParticleOptimizerContext
+ * @property {ArrayLike<number>} constants the shared constant pool, one float per word
+ * @property {((values: number[]) => number)|null} intern_constant interns a run of constants and
+ *   returns its first word index — {@link ../isa/ParticleAssembler.js}'s `constant`. Without it a
+ *   pass can read the pool but not add to it, which only costs it the folds that need a new literal.
+ * @property {number} register_slots size of the register file
+ * @property {number} max_iterations how many times this pass may repeat itself
+ */
+
+/**
+ * The passes, in the order a sweep runs them. `final` passes run once, after the sweeps settle.
+ *
+ * @type {Array<{name: string, run: ParticleOptimizerPass, max_iterations: number, final?: boolean}>}
+ */
+export const PARTICLE_OPTIMIZER_PASSES = [
+    { name: PARTICLE_PASS_DEAD_CODE, run: pass_eliminate_dead_code, max_iterations: 1 },
+    { name: PARTICLE_PASS_NARROW_WIDTHS, run: pass_narrow_widths, max_iterations: 1 },
+    { name: PARTICLE_PASS_FOLD_CONSTANTS, run: pass_fold_constants, max_iterations: 1 },
+    // Before sharing, and after the folds that intern new literals: two loads of one value are the
+    // same subexpression only once they name the same pool word.
+    { name: PARTICLE_PASS_DEDUP_CONSTANTS, run: pass_dedup_constants, max_iterations: 1 },
+    { name: PARTICLE_PASS_SHARE_SUBEXPRESSIONS, run: pass_share_subexpressions, max_iterations: 1 },
+    { name: PARTICLE_PASS_PROPAGATE_COPIES, run: pass_propagate_copies, max_iterations: 3 },
+    { name: PARTICLE_PASS_PEEPHOLE, run: pass_peephole, max_iterations: 1 },
+    { name: PARTICLE_PASS_PACK_REGISTERS, run: pass_pack_registers, max_iterations: 1, final: true },
+];
+
+/**
+ * Passes enabled unless the caller says otherwise. The switches are for bisecting and for measuring
+ * what each pass is worth; the two rewrites that change the arithmetic rather than the program are
+ * options of their passes rather than passes, so that the list stays a list of what to run and not a
+ * list of what to trust. `fold_constants.fold_multiply_by_zero` is on, `peephole.fuse_multiply_add`
+ * is off, and both are argued where they live.
+ *
+ * @type {Readonly<Object<string, boolean>>}
+ */
+export const PARTICLE_OPTIMIZER_DEFAULT_PASSES = Object.freeze(
+    Object.fromEntries(PARTICLE_OPTIMIZER_PASSES.map(pass => [pass.name, true]))
+);
+
+/**
+ * How many times the whole pass list may sweep before the optimizer gives up on reaching a fixed
+ * point. Passes feed each other — folding a `SELECT` makes a broadcast dead, deleting it lets the
+ * value above narrow — so one sweep is never enough. Two settle every program in this tree and
+ * three settle the worst of a hundred generated ones; the budget is well clear of that, because a
+ * program that ran out of sweeps would be a program whose optimization depended on the budget.
+ *
+ * @type {number}
+ */
+export const PARTICLE_OPTIMIZER_DEFAULT_ITERATIONS = 8;
+
+/**
+ * Optimize one section of a particle program.
+ *
+ * @param {InstructionStream} stream
+ * @param {object} [options]
+ * @param {ArrayLike<number>} [options.constants] the shared constant pool; without it the passes
+ *   that reason about literal values do nothing
+ * @param {function(number[]): number} [options.intern_constant] adds a run to that pool
+ * @param {Object<string, boolean|object>} [options.passes] per-pass switches, overriding
+ *   {@link PARTICLE_OPTIMIZER_DEFAULT_PASSES}; an object value also carries that pass's options
+ * @param {number} [options.max_iterations] sweeps of the whole pass list
+ * @param {boolean} [options.speculate] whether to try lane peeling both ways and keep the better
+ *   program; on unless told otherwise. Off makes the pipeline a straight run of its passes, which is
+ *   what a caller measuring one pass in isolation wants.
+ * @param {number} [options.register_slots]
+ * @returns {InstructionStream} a fresh stream, whose `register_count` is re-derived from what the
+ *   optimized instructions actually touch
+ */
+export function optimize_particle_program(stream, options = {}) {
+    assert.ok(stream instanceof InstructionStream, "stream must be an InstructionStream");
+
+    const register_slots = options.register_slots ?? PARTICLE_VM_REGISTER_SLOTS;
+    const max_iterations = options.max_iterations ?? PARTICLE_OPTIMIZER_DEFAULT_ITERATIONS;
+
+    assert.isNonNegativeInteger(max_iterations, "max_iterations");
+
+    /** @type {number[][]} */
+    const instructions = stream.instructions.map(words => words.slice());
+
+    const base_context = {
+        constants: options.constants ?? [],
+        intern_constant: options.intern_constant ?? null,
+        register_slots,
+    };
+
+    const speculative = to_stream(run_pipeline(instructions, base_context, options, max_iterations));
+
+    if (options.speculate === false || !peeling_could_apply(options)) {
+        return speculative;
+    }
+
+    // Peeling a constant tail off an instruction always costs one instruction where it fires, and
+    // pays only through what dies above it — a mask's load, the broadcast feeding it, the arithmetic
+    // that built them. Whether that cascade happens is not a local question, and a heuristic for it
+    // would be wrong on somebody's effect. So the pipeline is run both ways and the better program
+    // kept; these are at most a few hundred instructions and settling one twice costs nothing worth
+    // measuring.
+    const conservative = to_stream(
+        run_pipeline(instructions, base_context, without_peeling(options), max_iterations)
+    );
+
+    return better_of(conservative, speculative);
+}
+
+/**
+ * Sweep the enabled passes to a fixed point, then run the closing ones once.
+ *
+ * @param {number[][]} instructions
+ * @param {object} base_context
+ * @param {object} options
+ * @param {number} max_iterations
+ * @returns {number[][]}
+ */
+function run_pipeline(instructions, base_context, options, max_iterations) {
+    const sweep = PARTICLE_OPTIMIZER_PASSES.filter(pass => pass.final !== true);
+    const closing = PARTICLE_OPTIMIZER_PASSES.filter(pass => pass.final === true);
+
+    let current = instructions;
+
+    for (let iteration = 0; iteration < max_iterations; iteration++) {
+        let changed = false;
+
+        for (const pass of sweep) {
+            const result = run_pass(pass, current, base_context, options);
+
+            if (result !== null) {
+                current = result;
+                changed = true;
+            }
+        }
+
+        if (!changed) {
+            break;
+        }
+    }
+
+    for (const pass of closing) {
+        const result = run_pass(pass, current, base_context, options);
+
+        if (result !== null) {
+            current = result;
+        }
+    }
+
+    return current;
+}
+
+/**
+ * Fewer instructions wins; a tie goes to the smaller register file, and a tie there to the
+ * speculative one, whose instructions are the narrower of the two.
+ *
+ * @param {InstructionStream} conservative
+ * @param {InstructionStream} speculative
+ * @returns {InstructionStream}
+ */
+function better_of(conservative, speculative) {
+    if (conservative.instruction_count !== speculative.instruction_count) {
+        return conservative.instruction_count < speculative.instruction_count ? conservative : speculative;
+    }
+
+    return conservative.register_count < speculative.register_count ? conservative : speculative;
+}
+
+/** @returns {boolean} whether the run just made could have peeled anything */
+function peeling_could_apply(options) {
+    const setting = options.passes?.[PARTICLE_PASS_FOLD_CONSTANTS];
+
+    if (setting === false) {
+        return false;
+    }
+
+    return !(typeof setting === "object" && setting !== null && setting.peel_constant_lanes === false);
+}
+
+/** @returns {object} the same options with lane peeling turned off */
+function without_peeling(options) {
+    const setting = options.passes?.[PARTICLE_PASS_FOLD_CONSTANTS];
+    const fold = typeof setting === "object" && setting !== null ? setting : {};
+
+    return {
+        ...options,
+        passes: {
+            ...options.passes,
+            [PARTICLE_PASS_FOLD_CONSTANTS]: { ...fold, peel_constant_lanes: false },
+        },
+    };
+}
+
+/**
+ * @param {{name: string, run: ParticleOptimizerPass, max_iterations: number}} pass
+ * @param {number[][]} instructions
+ * @param {object} base_context
+ * @param {object} options
+ * @returns {number[][]|null}
+ */
+function run_pass(pass, instructions, base_context, options) {
+    const setting = options.passes?.[pass.name];
+
+    if (setting === false) {
+        return null;
+    }
+
+    const overrides = typeof setting === "object" && setting !== null ? setting : {};
+
+    return pass.run(instructions, {
+        ...base_context,
+        max_iterations: pass.max_iterations,
+        ...overrides,
+    });
+}
+
+/**
+ * A `passes` setting enabling only the named passes.
+ *
+ * The plain object form is additive — anything it does not mention keeps its default — which is what
+ * a caller turning one pass off wants and the opposite of what a caller measuring one pass wants.
+ *
+ * @param {...string} names
+ * @returns {Object<string, boolean>}
+ */
+export function only_particle_passes(...names) {
+    const switches = {};
+
+    for (const pass of PARTICLE_OPTIMIZER_PASSES) {
+        switches[pass.name] = names.includes(pass.name);
+    }
+
+    return switches;
+}
+
+/**
+ * Optimize both sections of a compiled program against their shared constant pool.
+ *
+ * The pool is shared and may grow — a folded expression needs somewhere to put its literal — so it
+ * is interned through a {@link ParticleAssembler}, which is the code that already knows how to reuse
+ * a run that is anywhere in the pool rather than only one on a boundary.
+ *
+ * @param {ParticleProgram} program
+ * @param {object} [options] as {@link optimize_particle_program}, minus the pool it supplies itself
+ * @returns {ParticleProgram}
+ */
+export function optimize_compiled_particle_program(program, options = {}) {
+    const assembler = new ParticleAssembler();
+
+    assembler.constants = Array.from(program.constants);
+
+    const shared = {
+        ...options,
+        constants: assembler.constants,
+        intern_constant: (values) => assembler.constant(...values),
+    };
+
+    const init = optimize_particle_program(stream_from_words(program.init), shared);
+    const update = optimize_particle_program(stream_from_words(program.update), shared);
+
+    const init_words = init.toWords();
+    const update_words = update.toWords();
+
+    // Last, and here rather than in a pass: which pool words are dead is a question about the whole
+    // program, and a pass sees one section. Running it before the sections settled would be wrong
+    // twice over — the folds are still interning, and the speculation below has not yet thrown away
+    // the run whose literals are the dead ones.
+    const constants = sweep_particle_constants(
+        [init_words, update_words],
+        assembler.constants,
+        sweeping_is_enabled(options)
+    );
+
+    return new ParticleProgram({
+        init: init_words,
+        update: update_words,
+        constants,
+        register_count: Math.max(init.register_count, update.register_count),
+    });
+}
+
+/**
+ * Whether the caller left the constant-pool passes on. See the note in
+ * {@link ./pass_dedup_constants.js} on why the compaction follows the rewrite's switch.
+ *
+ * @param {object} options
+ * @returns {boolean}
+ */
+function sweeping_is_enabled(options) {
+    const setting = options.passes?.[PARTICLE_PASS_DEDUP_CONSTANTS];
+
+    if (setting === false) {
+        return false;
+    }
+
+    return !(typeof setting === "object" && setting !== null && setting.sweep === false);
+}
+
+/**
+ * Drop the pool words no section loads any more, and move the loads onto what is left.
+ *
+ * Marking is per word rather than per run, and that is what makes compaction safe for runs that
+ * overlap — a scalar sharing the leading word of a `vec4` is something the assembler deliberately
+ * produces. Every word of a live run is marked, and compaction keeps marked words in order, so a run
+ * whose words were consecutive still is; only its base moves.
+ *
+ * @param {Uint32Array[]} sections instruction words, rewritten in place
+ * @param {ArrayLike<number>} constants the shared pool
+ * @param {boolean} enabled
+ * @returns {Float32Array} the pool the sections now address
+ */
+function sweep_particle_constants(sections, constants, enabled) {
+    if (!enabled) {
+        return new Float32Array(constants);
+    }
+
+    const live = new Uint8Array(constants.length);
+
+    let well_formed = true;
+
+    for (const words of sections) {
+        visit_particle_constant_reads(instruction_tuples(words), (index, base, span) => {
+            if (base < 0 || base + span > constants.length) {
+                // Out of the pool already. Compacting under it would turn a program that fails
+                // validation into one that fails it somewhere else, so nothing is touched.
+                well_formed = false;
+                return;
+            }
+
+            for (let i = 0; i < span; i++) {
+                live[base + i] = 1;
+            }
+        });
+    }
+
+    if (!well_formed) {
+        return new Float32Array(constants);
+    }
+
+    /** @type {number[]} */
+    const kept = [];
+
+    /** @type {Int32Array} old word index -> new one; only meaningful where `live` */
+    const remap = new Int32Array(constants.length);
+
+    for (let i = 0; i < constants.length; i++) {
+        remap[i] = kept.length;
+
+        if (live[i] === 1) {
+            kept.push(constants[i]);
+        }
+    }
+
+    if (kept.length !== constants.length) {
+        for (const words of sections) {
+            for (let base = 0; base < words.length; base += PARTICLE_VM_INSTRUCTION_WORDS) {
+                if (vm_decode_op(words[base]) !== VM_OP.LOAD_CONST) {
+                    continue;
+                }
+
+                words[base + 2] = remap[words[base + 2]];
+            }
+        }
+    }
+
+    return new Float32Array(kept);
+}
+
+/**
+ * A packed section as the five-word tuples the pass helpers read.
+ *
+ * @param {Uint32Array} words
+ * @returns {Uint32Array[]} views, so nothing is copied
+ */
+function instruction_tuples(words) {
+    const out = [];
+
+    for (let base = 0; base < words.length; base += PARTICLE_VM_INSTRUCTION_WORDS) {
+        out.push(words.subarray(base, base + PARTICLE_VM_INSTRUCTION_WORDS));
+    }
+
+    return out;
+}
+
+/**
+ * @param {Uint32Array} words a packed instruction section
+ * @returns {InstructionStream}
+ */
+export function stream_from_words(words) {
+    const out = new InstructionStream();
+
+    for (let base = 0; base < words.length; base += PARTICLE_VM_INSTRUCTION_WORDS) {
+        out.emit(
+            vm_decode_op(words[base]),
+            words[base + 1],
+            words[base + 2],
+            words[base + 3],
+            words[base + 4],
+            vm_decode_width(words[base])
+        );
+    }
+
+    return out;
+}
+
+/**
+ * Re-emit into a fresh stream, so the register high-water mark is the optimized program's rather
+ * than the one it was built from. A file that shrank is itself part of the win — it is the
+ * interpreter's dominant register cost (see the occupancy note in `vm/chunk_particle_vm.js`) — and
+ * carrying the old count forward would hide it.
+ *
+ * @param {number[][]} instructions
+ * @returns {InstructionStream}
+ */
+function to_stream(instructions) {
+    const out = new InstructionStream();
+
+    for (const words of instructions) {
+        out.emit(vm_decode_op(words[0]), words[1], words[2], words[3], words[4], vm_decode_width(words[0]));
+    }
+
+    return out;
+}