@woosh/meep-engine 3.15.0 → 3.17.0

package/src/shade/renderer/particles/optimizer/dag/particle_vector_lower.js

@@ -0,0 +1,687 @@
+import { assert } from "../../../../../core/assert.js";
+import { InstructionStream } from "../../isa/InstructionStream.js";
+import { VM_MAX_WIDTH, VM_OP, pack_swizzle, vm_immediate, vm_register } from "../../isa/ParticleVMISA.js";
+import { PV_GROUP } from "./particle_vector_groups.js";
+
+/**
+ * Give the groups an order and a register file, and write the instructions.
+ *
+ * By the time this runs there is nothing left to decide about WHAT the program computes — only when
+ * each instruction runs and which slots it uses. Those are the two halves of the register-count
+ * problem, and the notes' pipeline is exactly this: schedule for the fewest values alive at once,
+ * then colour, which for straight-line code is an interval graph and so optimal for a fixed order.
+ *
+ * ## Ordering
+ *
+ * The value graph left almost nothing to obey. There are no anti-dependencies at all — a value is
+ * not a register until this module makes it one, so the write-after-read and write-after-write
+ * hazards that constrain `../pass_schedule.js` do not exist here. What remains is the true data
+ * dependencies and ONE chain: the RNG draws come off a sequence, so the `RANDOM` groups keep their
+ * order among themselves. Stores and kills are free to go anywhere after the values they read, which
+ * is the practical difference between the two optimizers' schedules — over there a store is a
+ * barrier nothing crosses.
+ *
+ * The order is chosen by the same backward greedy `../pass_schedule.js` uses: place last the
+ * instruction that leaves fewest slots alive above it. It is a greedy, and a greedy can be fooled;
+ * an exact minimum-register schedule is a constraint-programming problem, which is where the notes
+ * point next and where a solver would be needed.
+ *
+ * ## Colouring
+ *
+ * A group's run is live from where it is written to the last instruction that reads it, so the live
+ * ranges are intervals and the lowest-base first fit is the classic optimal left-edge colouring —
+ * for SCALARS. A run of two or more is a shape rather than a point and first fit is then a
+ * heuristic; the same caveat `../pass_allocate_registers.js` carries, and the reason it is worth
+ * running that pass over the result afterwards.
+ *
+ * One exception to "sources are read before the destination is written", which is otherwise true of
+ * every instruction in this ISA and is what lets a destination reuse a slot that dies at the same
+ * instruction: a GATHER is several `MOV`s, and the first of them would clobber a slot a later one
+ * still has to read. So a gather's run is allocated before its operands are released and never
+ * overlaps them.
+ */
+
+/**
+ * @typedef {object} ParticleLoweredSection
+ * @property {InstructionStream} stream
+ * @property {number} peak the most slots alive at once, which is what the file has to hold
+ * @property {number[]} order the group indices, in the order they run
+ */
+
+/**
+ * @param {import("./particle_vector_groups.js").ParticleValuePacking} packing
+ * @param {object} params
+ * @param {number} params.register_slots
+ * @param {function(number[]): number} params.intern_constant adds a run to the shared pool and
+ *   returns its first word
+ * @returns {ParticleLoweredSection|null} `null` when the program does not fit the register file
+ */
+export function lower_particle_values(packing, { register_slots, intern_constant }) {
+    const order = schedule_groups(packing.groups);
+    const bases = allocate_groups(packing.groups, order, register_slots, packing.coalitions);
+
+    if (bases === null) {
+        return null;
+    }
+
+    const stream = new InstructionStream();
+
+    for (const index of order) {
+        emit_group(stream, packing.groups[index], bases, intern_constant);
+    }
+
+    return { stream, peak: stream.register_count, order };
+}
+
+/**
+ * Per group, the groups it must come after.
+ *
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup[]} groups
+ * @returns {Set<number>[]}
+ */
+export function build_group_dependencies(groups) {
+    const predecessors = groups.map(() => new Set());
+    const loads = groups.filter(group => group.op === VM_OP.LOAD_ATTR);
+
+    for (const group of groups) {
+        for (const operand of group.operands) {
+            for (const read of operand.reads) {
+                predecessors[group.index].add(read.group);
+            }
+        }
+
+        for (const source of group.gather) {
+            if (source.group >= 0) {
+                predecessors[group.index].add(source.group);
+            }
+        }
+
+        // The one ordering the value graph did not already discharge. Lifting forwarded every load of
+        // a word that had been stored, so what is left is loads of the word AS THE SECTION FOUND IT —
+        // and those are only true until something stores over that word. The value is a fact about
+        // the record before the section ran; the INSTRUCTION that fetches it is not.
+        if (group.op === VM_OP.STORE_ATTR) {
+            for (const load of loads) {
+                if (words_overlap(load.raw[0], load.raw[1], group.raw[1], group.raw[2])) {
+                    predecessors[group.index].add(load.index);
+                }
+            }
+        }
+    }
+
+    // The draws are a sequence: the state after `n` of them is what the section is compared on, so
+    // they may move anywhere as a block but never past one another. Which one comes first is what
+    // `draw` says and not where the group happens to sit in the list.
+    const draws = groups.filter(group => group.op === VM_OP.RANDOM).sort((a, b) => a.draw - b.draw);
+
+    for (let i = 1; i < draws.length; i++) {
+        predecessors[draws[i].index].add(draws[i - 1].index);
+    }
+
+    return predecessors;
+}
+
+/**
+ * @param {number} base_a
+ * @param {number} span_a
+ * @param {number} base_b
+ * @param {number} span_b
+ * @returns {boolean} whether the two record-word ranges share a word
+ */
+function words_overlap(base_a, span_a, base_b, span_b) {
+    return base_a < base_b + span_b && base_b < base_a + span_a;
+}
+
+/**
+ * The order the groups run in: whichever of two candidate schedules holds fewer slots at once.
+ *
+ * The first is the section's own order, rebuilt — a forward topological walk that always takes the
+ * ready group whose work came earliest. It is not clever and it is not meant to be; it is the
+ * schedule the program arrived with, which is a schedule somebody's frontend already thought about,
+ * and it is what stops this stage from ever being the reason a program stopped fitting.
+ *
+ * The second is the backward greedy `../pass_schedule.js` uses: place last the group that leaves
+ * fewest slots alive above it. It wins on programs whose values were computed a long way from where
+ * they are used — which is most hand-lowered code — and loses on programs made of many independent
+ * chains, where the liveness measure is flat and it interleaves twelve noise samples that should
+ * have been done one at a time.
+ *
+ * Neither is the minimum. An exact minimum-register schedule of a DAG is NP-complete and the answer
+ * at this size is a constraint solver, which is the next thing the notes point at; measuring two
+ * cheap orders and keeping the better is what is here instead, and it is at least never worse than
+ * the program's own.
+ *
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup[]} groups
+ * @returns {number[]} group indices, in the order they run
+ */
+export function schedule_groups(groups) {
+    const predecessors = build_group_dependencies(groups);
+    const sequence = groups.map(group => group.sequence);
+    const original = schedule_in_priority_order(groups, predecessors, sequence);
+    // Three orders, and each of them wins on programs the other two lose on. Over 1,600 generated
+    // programs the section's own order won 1,376, the greedy 214 and the sink 10 — the greedy is
+    // worth its quadratic ready set on the small ones and hopeless on the large, where it interleaves
+    // every independent chain there is; the sink is worth having for the shapes where a value was
+    // computed a long way from where it is read. Iterating the sink over POSITIONS rather than the
+    // section's numbering was tried and won nothing at all, on any program measured.
+    const candidates = [
+        original,
+        schedule_in_priority_order(
+            groups,
+            predecessors,
+            as_late_as_possible(groups, predecessors, original, sequence),
+            sequence
+        ),
+        schedule_by_liveness(groups, predecessors, sequence),
+    ];
+
+    let best = candidates[0];
+    let best_peak = peak_slots(groups, best);
+
+    for (const candidate of candidates.slice(1)) {
+        const peak = peak_slots(groups, candidate);
+
+        if (peak < best_peak) {
+            best = candidate;
+            best_peak = peak;
+        }
+    }
+
+    return best;
+}
+
+/**
+ * The latest place in the section's own order each group could have gone: just before whatever reads
+ * it first.
+ *
+ * A group that is placed where it was WRITTEN holds its run from there until its reader; a group
+ * placed just before its reader holds it for an instruction. That is the whole of the difference
+ * between a peak of a hundred and a peak of thirty on a program made of a dozen noise samples, where
+ * grouping lanes into vectors has pulled a value's readers away from where the value was computed.
+ *
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup[]} groups
+ * @param {Set<number>[]} predecessors
+ * @param {number[]} order any valid order, walked backwards so a group is pulled down only once
+ *   every reader of it already has been
+ * @param {number[]} base per group, the time it would run at if nothing pulled it
+ * @returns {number[]} per group, its priority — smaller runs earlier
+ */
+function as_late_as_possible(groups, predecessors, order, base) {
+    // A group with no reader has nowhere later to go, so it stays where it was; everything else is
+    // pulled down to just above the first thing that reads it.
+    const latest = base.slice();
+
+    for (let i = order.length - 1; i >= 0; i--) {
+        for (const predecessor of predecessors[order[i]]) {
+            latest[predecessor] = Math.min(latest[predecessor], latest[order[i]]);
+        }
+    }
+
+    return latest;
+}
+
+/**
+ * A forward topological walk that always takes the ready group with the smallest priority.
+ *
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup[]} groups
+ * @param {Set<number>[]} predecessors
+ * @param {number[]} priority per group
+ * @param {number[]} [tiebreak] per group, consulted where the priority is equal — which it very often
+ *   is once groups have been pulled down onto the same reader
+ * @returns {number[]}
+ */
+function schedule_in_priority_order(groups, predecessors, priority, tiebreak = priority) {
+    const pending = predecessors.map(set => set.size);
+    /** @type {Set<number>} */
+    const ready = new Set();
+    /** @type {number[][]} per group, what becomes ready once it is placed */
+    const successors = groups.map(() => []);
+
+    for (let i = 0; i < groups.length; i++) {
+        for (const predecessor of predecessors[i]) {
+            successors[predecessor].push(i);
+        }
+
+        if (pending[i] === 0) {
+            ready.add(i);
+        }
+    }
+
+    /** @type {number[]} */
+    const order = [];
+
+    while (order.length < groups.length) {
+        let best = -1;
+
+        for (const candidate of ready) {
+            if (best < 0 || earlier(priority, tiebreak, candidate, best)) {
+                best = candidate;
+            }
+        }
+
+        assert.ok(best >= 0, "particle value lowering: the group graph has a cycle");
+
+        order.push(best);
+        ready.delete(best);
+
+        for (const successor of successors[best]) {
+            if (--pending[successor] === 0) {
+                ready.add(successor);
+            }
+        }
+    }
+
+    return order;
+}
+
+/**
+ * @param {number[]} priority
+ * @param {number[]} tiebreak
+ * @param {number} candidate
+ * @param {number} best
+ * @returns {boolean}
+ */
+function earlier(priority, tiebreak, candidate, best) {
+    if (priority[candidate] !== priority[best]) {
+        return priority[candidate] < priority[best];
+    }
+
+    return tiebreak[candidate] !== tiebreak[best] ? tiebreak[candidate] < tiebreak[best] : candidate < best;
+}
+
+/**
+ * The most slots an order holds at once — what the register file has to be.
+ *
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup[]} groups
+ * @param {number[]} order
+ * @returns {number}
+ */
+export function peak_slots(groups, order) {
+    const position = new Int32Array(groups.length).fill(-1);
+
+    for (let i = 0; i < order.length; i++) {
+        position[order[i]] = i;
+    }
+
+    const last_use = new Int32Array(groups.length).fill(-1);
+
+    for (const group of groups) {
+        for (const source of sources_of(group)) {
+            last_use[source] = Math.max(last_use[source], position[group.index]);
+        }
+    }
+
+    /** @type {number[][]} */
+    const dying = order.map(() => []);
+
+    let live = 0;
+    let peak = 0;
+
+    for (let i = 0; i < order.length; i++) {
+        const group = groups[order[i]];
+
+        live += group.span;
+        peak = Math.max(peak, live);
+        dying[Math.max(i, last_use[group.index])].push(group.span);
+
+        for (const span of dying[i]) {
+            live -= span;
+        }
+    }
+
+    return peak;
+}
+
+/**
+ * Backwards, one group at a time, from those whose consumers are all already placed: the one that
+ * leaves fewest slots alive above it.
+ *
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup[]} groups
+ * @param {Set<number>[]} predecessors
+ * @param {number[]} priority per group; ties go to the one that should run later
+ * @returns {number[]} group indices, in the order they run
+ */
+function schedule_by_liveness(groups, predecessors, priority) {
+    const pending = new Int32Array(groups.length);
+
+    for (const set of predecessors) {
+        for (const predecessor of set) {
+            pending[predecessor]++;
+        }
+    }
+
+    /** @type {Set<number>} groups every consumer of which is placed */
+    const ready = new Set();
+
+    for (let i = 0; i < groups.length; i++) {
+        if (pending[i] === 0) {
+            ready.add(i);
+        }
+    }
+
+    /** @type {Set<number>} groups whose runs are alive just above the ones placed so far */
+    const live = new Set();
+    /** @type {number[]} */
+    const reversed = [];
+
+    while (reversed.length < groups.length) {
+        let best = -1;
+        let best_live = Infinity;
+
+        for (const candidate of ready) {
+            const alive = live_above(groups, candidate, live);
+
+            // Ties go to the group whose work came LATEST in the section as it was written, so that
+            // a stretch this measure is indifferent about comes back in the order it arrived.
+            if (alive < best_live || (alive === best_live && later(priority, candidate, best))) {
+                best = candidate;
+                best_live = alive;
+            }
+        }
+
+        assert.ok(best >= 0, "particle value lowering: the group graph has a cycle");
+
+        live.delete(best);
+
+        for (const source of sources_of(groups[best])) {
+            live.add(source);
+        }
+
+        reversed.push(best);
+        ready.delete(best);
+
+        for (const predecessor of predecessors[best]) {
+            if (--pending[predecessor] === 0) {
+                ready.add(predecessor);
+            }
+        }
+    }
+
+    return reversed.reverse();
+}
+
+/**
+ * @param {number[]} priority
+ * @param {number} candidate
+ * @param {number} best the best so far, or `-1`
+ * @returns {boolean}
+ */
+function later(priority, candidate, best) {
+    if (best < 0) {
+        return true;
+    }
+
+    return priority[candidate] !== priority[best] ? priority[candidate] > priority[best] : candidate > best;
+}
+
+/**
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup} group
+ * @returns {number[]} the groups it reads, with repeats
+ */
+function sources_of(group) {
+    const sources = [];
+
+    for (const operand of group.operands) {
+        for (const read of operand.reads) {
+            sources.push(read.group);
+        }
+    }
+
+    for (const source of group.gather) {
+        if (source.group >= 0) {
+            sources.push(source.group);
+        }
+    }
+
+    return sources;
+}
+
+/**
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup[]} groups
+ * @param {number} candidate
+ * @param {Set<number>} below the groups alive just under it
+ * @returns {number} slots alive just above it, were it placed here
+ */
+function live_above(groups, candidate, below) {
+    const above = new Set(below);
+
+    above.delete(candidate);
+
+    for (const source of sources_of(groups[candidate])) {
+        above.add(source);
+    }
+
+    let slots = 0;
+
+    for (const group of above) {
+        slots += groups[group].span;
+    }
+
+    return slots;
+}
+
+/**
+ * Give each group's run the lowest base that is free for as long as it lives.
+ *
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup[]} groups
+ * @param {number[]} order
+ * @param {number} register_slots
+ * @param {number[][]} [coalitions] runs that have to end up end to end, in that order
+ * @returns {Int32Array|null} per group, the slot its run starts at; `null` when the file is too small
+ */
+export function allocate_groups(groups, order, register_slots, coalitions = []) {
+    const position = positions_of(groups, order);
+    const { last, gathered } = lane_uses(groups, position);
+    const bases = new Int32Array(groups.length).fill(-1);
+    const free = new Uint8Array(register_slots).fill(1);
+    /** @type {number[][]} per position, the slots that come free once it has run */
+    const release = order.map(() => []);
+
+    /**
+     * Hand each slot of a run back at its own time.
+     *
+     * @param {import("./particle_vector_groups.js").ParticleValueGroup} group
+     * @param {number} base
+     */
+    const release_lanes = (group, base) => {
+        for (let lane = 0; lane < group.span; lane++) {
+            const key = `${group.index}:${lane}`;
+            // A slot is free where its last reader runs, since an instruction reads every source
+            // before it writes — so the reader may have it. One nothing reads comes free one position
+            // later, this being where it is written; and so does one a GATHER reads, since a gather
+            // fills its run one `MOV` at a time and must not clobber a lane it has yet to read.
+            const end = Math.max(
+                position[group.index] + 1,
+                (last.get(key) ?? -1) + (gathered.has(key) ? 1 : 0)
+            );
+
+            if (end < order.length) {
+                release[end].push(base + lane);
+            }
+        }
+    };
+
+    for (let i = 0; i < order.length; i++) {
+        const group = groups[order[i]];
+
+        for (const slot of release[i]) {
+            free[slot] = 1;
+        }
+
+        // A coalition member gets its slot when the coalition is claimed, which is before it runs.
+        if (group.span === 0 || bases[group.index] >= 0) {
+            continue;
+        }
+
+        // A coalition is claimed whole, by whichever of its runs is written first: they have to end
+        // up adjacent, and the only moment that can be arranged is before any of them has a slot.
+        const members = group.coalition >= 0 ? coalitions[group.coalition] : [group.index];
+        const width = members.reduce((slots, member) => slots + groups[member].span, 0);
+        const base = first_fit(free, width);
+
+        if (base < 0) {
+            return null;
+        }
+
+        free.fill(0, base, base + width);
+
+        let at = base;
+
+        for (const member of members) {
+            bases[member] = at;
+            release_lanes(groups[member], at);
+            at += groups[member].span;
+        }
+    }
+
+    return bases;
+}
+
+/**
+ * Per SLOT of each group's run, the last position that reads it — and which of those slots a gather
+ * reads.
+ *
+ * Per slot and not per run, which is the whole point: a `RANDOM.4` whose fourth draw nothing wants,
+ * a builtin loaded three lanes wide for its third, a `CROSS` two lanes of which are read — all of
+ * those hold slots that died the moment they were written, and freeing the run as a unit would keep
+ * them to the end of the last lane anybody wanted.
+ *
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup[]} groups
+ * @param {number[]} position
+ * @returns {{last: Map<string, number>, gathered: Set<string>}} keyed `group:lane`
+ */
+function lane_uses(groups, position) {
+    /** @type {Map<string, number>} */
+    const last = new Map();
+    /** @type {Set<string>} */
+    const gathered = new Set();
+
+    const touch = (group, lane, at) => {
+        const key = `${group}:${lane}`;
+
+        last.set(key, Math.max(last.get(key) ?? -1, at));
+
+        return key;
+    };
+
+    for (const group of groups) {
+        for (const operand of group.operands) {
+            // `reads` and not `lanes`: a coalition operand's lanes are slots of the whole coalition,
+            // and what is being kept alive is a lane of one of the runs inside it.
+            for (const read of operand.reads) {
+                touch(read.group, read.lane, position[group.index]);
+            }
+        }
+
+        for (const source of group.gather) {
+            if (source.group >= 0) {
+                gathered.add(touch(source.group, source.lane, position[group.index]));
+            }
+        }
+    }
+
+    return { last, gathered };
+}
+
+/**
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup[]} groups
+ * @param {number[]} order
+ * @returns {number[]} per group, where the order put it
+ */
+function positions_of(groups, order) {
+    const positions = new Array(groups.length).fill(-1);
+
+    for (let i = 0; i < order.length; i++) {
+        positions[order[i]] = i;
+    }
+
+    return positions;
+}
+
+/**
+ * @param {Uint8Array} free
+ * @param {number} span
+ * @returns {number} the lowest base with `span` free slots above it, or `-1`
+ */
+function first_fit(free, span) {
+    let run = 0;
+
+    for (let slot = 0; slot < free.length; slot++) {
+        run = free[slot] === 1 ? run + 1 : 0;
+
+        if (run === span) {
+            return slot - span + 1;
+        }
+    }
+
+    return -1;
+}
+
+/**
+ * @param {InstructionStream} stream
+ * @param {import("./particle_vector_groups.js").ParticleValueGroup} group
+ * @param {Int32Array} bases
+ * @param {function(number[]): number} intern_constant
+ */
+
+function emit_group(stream, group, bases, intern_constant) {
+    const base = bases[group.index];
+
+    switch (group.kind) {
+        case PV_GROUP.GATHER:
+            for (let lane = 0; lane < group.gather.length; lane++) {
+                const source = group.gather[lane];
+                const operand = source.group >= 0
+                    ? vm_register(bases[source.group] + source.lane, pack_swizzle(0, 1, 2, 3))
+                    : vm_immediate(source.value);
+
+                stream.emit(VM_OP.MOV, base + lane, operand, 0, 0, 1);
+            }
+
+            return;
+
+        case PV_GROUP.CONSTANT:
+            stream.emit(VM_OP.LOAD_CONST, base, intern_constant(group.constants), 0, 0, group.width);
+
+            return;
+
+        default: {
+            const sources = [0, 0, 0];
+
+            for (let slot = 0; slot < 3; slot++) {
+                if (group.raw[slot] >= 0) {
+                    sources[slot] = group.raw[slot];
+                }
+            }
+
+            for (const operand of group.operands) {
+                sources[operand.slot] = operand.group < 0
+                    ? vm_immediate(operand.value)
+                    : register_operand(bases[operand.group], operand.lanes);
+            }
+
+            stream.emit(group.op, group.span === 0 ? 0 : base, sources[0], sources[1], sources[2], group.width);
+        }
+    }
+}
+
+/**
+ * A register operand reading given lanes of a run: the base is the lowest of them and the selectors
+ * say the rest. Every lane of an operand comes from ONE group, whose run is at most four slots, so
+ * the selectors always fit the two bits they have.
+ *
+ * @param {number} group_base
+ * @param {number[]} lanes which lane of the run each lane of the operand reads
+ * @returns {number} a register operand word
+ */
+function register_operand(group_base, lanes) {
+    const lowest = Math.min(...lanes);
+    const selector = lane => (lane < lanes.length ? lanes[lane] - lowest : lane);
+
+    assert.ok(
+        lanes.every(lane => lane - lowest < VM_MAX_WIDTH),
+        `particle value lowering: lanes ${lanes.join(",")} do not fit one operand`
+    );
+
+    return vm_register(group_base + lowest, pack_swizzle(selector(0), selector(1), selector(2), selector(3)));
+}