@woosh/meep-engine 3.16.0 → 3.17.0

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

@@ -1,271 +1,392 @@
-import { VM_MAX_WIDTH, VM_OP, vm_encode_word0 } from "../isa/ParticleVMISA.js";
-import {
-    ParticleValueNumbering,
-    decode_particle_instructions,
-    particle_instruction_is_well_formed,
-    particle_run_is_dead_after,
-    particle_run_is_touched_between,
-} from "./particle_program_analysis.js";
-
-/** @type {string} */
-export const PARTICLE_PASS_PEEPHOLE = "peephole";
-
-/**
- * Local rewrites over a couple of neighbouring instructions — the ones that need no analysis beyond
- * "these two are really one".
- *
- * ## Constant loads
- *
- * A frontend emits one load per literal, because it lowers one value at a time. Where two of them
- * land in adjacent registers from adjacent pool words, they are one load: the pool is addressed by
- * word, so a two-word run is as loadable as a one-word one. The graph compiler's `mad` node — a
- * multiplier and an addend, interned side by side — is exactly this shape.
- *
- * ## Multiplying by a broadcast
- *
- * `MUL` against a value whose every lane is the same is `SCALE`, which reads one lane and takes the
- * broadcast out of the program. Widening a scalar to meet a wider operation is the compiler's
- * standard move, and `SCALE` is the opcode that exists so it does not have to.
- *
- * ## Fusing a multiply and an add — off by default
- *
- * `MUL` then `ADD` is `MAD`, one instruction rather than two. It is also, unavoidably, a DIFFERENT
- * NUMBER: the pair rounds the product to `f32` before adding, and `MAD` does not — and on hardware
- * that contracts it into an FMA, does not round it at all. That is the ordinary compiler trade and
- * it is a perfectly reasonable one to take for a particle effect, but it is not the bit-identity the
- * rest of this optimizer holds itself to, so it is a switch rather than a default:
- * `{ passes: { peephole: { fuse_multiply_add: true } } }`.
- *
- * @type {import("./optimize_particle_program.js").ParticleOptimizerPass}
- */
-export function pass_peephole(instructions, context) {
-    const options = {
-        merge_constant_loads: true,
-        multiply_by_broadcast: true,
-        fuse_multiply_add: false,
-        ...context,
-    };
-
-    let current = instructions;
-    let changed = false;
-
-    if (options.multiply_by_broadcast) {
-        const result = scale_by_broadcast(current, context.register_slots);
-
-        if (result !== null) {
-            current = result;
-            changed = true;
-        }
-    }
-
-    for (const rewrite of [
-        options.merge_constant_loads ? merge_one_constant_load : null,
-        options.fuse_multiply_add ? fuse_one_multiply_add : null,
-    ]) {
-        if (rewrite === null) {
-            continue;
-        }
-
-        // Applied one at a time and re-derived: each rewrite deletes an instruction, and re-deriving
-        // a program of this size is cheaper than keeping indices and liveness honest across a splice.
-        for (let attempt = 0; attempt < instructions.length; attempt++) {
-            const result = rewrite(current);
-
-            if (result === null) {
-                break;
-            }
-
-            current = result;
-            changed = true;
-        }
-    }
-
-    return changed ? current : null;
-}
-
-/**
- * `MUL` by a value whose lanes all hold the same number is `SCALE` of the register that number came
- * from — which usually leaves the broadcast that widened it dead.
- *
- * @param {number[][]} instructions
- * @param {number} slots
- * @returns {number[][]|null}
- */
-function scale_by_broadcast(instructions, slots) {
-    const decoded = decode_particle_instructions(instructions);
-    const values = new ParticleValueNumbering(slots);
-
-    /** @type {number[][]} */
-    const out = [];
-
-    let changed = false;
-
-    for (const instruction of decoded) {
-        let words = instruction.words;
-
-        if (instruction.op === VM_OP.MUL && instruction.operation_width > 1) {
-            // Either operand may be the broadcast one; `SCALE` keeps the vector in slot 0.
-            const candidate = as_scale(instruction, instruction.reads[1], instruction.reads[0], values)
-                ?? as_scale(instruction, instruction.reads[0], instruction.reads[1], values);
-
-            if (candidate !== null) {
-                words = candidate;
-                changed = true;
-            }
-        }
-
-        out.push(words);
-        values.publish(words);
-    }
-
-    return changed ? out : null;
-}
-
-/**
- * @param {import("./particle_program_analysis.js").ParticleDecodedInstruction} instruction
- * @param {import("./particle_program_analysis.js").ParticleOperandRun} broadcast the operand hoped
- *   to hold one value across every lane
- * @param {import("./particle_program_analysis.js").ParticleOperandRun} vector the other one
- * @param {ParticleValueNumbering} values
- * @returns {number[]|null}
- */
-function as_scale(instruction, broadcast, vector, values) {
-    const ids = values.ids_of(broadcast.base, broadcast.span);
-
-    if (ids === null || ids.some(id => id !== ids[0])) {
-        return null;
-    }
-
-    // The lane the value was computed in, which is the single slot SCALE will read.
-    const scalar = values.origin_of(broadcast.base, 1);
-
-    if (scalar < 0) {
-        return null;
-    }
-
-    const words = [
-        vm_encode_word0(VM_OP.SCALE, instruction.width),
-        instruction.write.base,
-        vector.base,
-        scalar,
-        0,
-    ];
-
-    return particle_instruction_is_well_formed(words) ? words : null;
-}
-
-/**
- * Two constant loads that fill adjacent registers from adjacent pool words are one load.
- *
- * @param {number[][]} instructions
- * @returns {number[][]|null}
- */
-function merge_one_constant_load(instructions) {
-    const decoded = decode_particle_instructions(instructions);
-
-    for (let first = 0; first < decoded.length; first++) {
-        const a = decoded[first];
-
-        if (a.op !== VM_OP.LOAD_CONST) {
-            continue;
-        }
-
-        for (let second = first + 1; second < decoded.length; second++) {
-            const b = decoded[second];
-
-            if (b.op === VM_OP.HALT) {
-                break;
-            }
-
-            if (b.op !== VM_OP.LOAD_CONST
-                || b.words[1] !== a.words[1] + a.operation_width
-                || b.words[2] !== a.words[2] + a.operation_width
-                || a.operation_width + b.operation_width > VM_MAX_WIDTH) {
-                continue;
-            }
-
-            const base = a.words[1];
-            const span = a.operation_width + b.operation_width;
-
-            // The second load moves back to the first one's position, so nothing in between may
-            // read or write either register run.
-            if (particle_run_is_touched_between(decoded, first + 1, second, base, span)) {
-                continue;
-            }
-
-            const words = [vm_encode_word0(VM_OP.LOAD_CONST, span), base, a.words[2], 0, 0];
-
-            if (!particle_instruction_is_well_formed(words)) {
-                continue;
-            }
-
-            const out = instructions.slice();
-
-            out[first] = words;
-            out.splice(second, 1);
-
-            return out;
-        }
-    }
-
-    return null;
-}
-
-/**
- * `MUL` into a register that the very next instruction adds to something is `MAD`.
- *
- * Adjacent only: the pair is what a frontend emits for `a * b + c`, and looking further would buy
- * little for a rewrite that is off by default anyway.
- *
- * @param {number[][]} instructions
- * @returns {number[][]|null}
- */
-function fuse_one_multiply_add(instructions) {
-    const decoded = decode_particle_instructions(instructions);
-
-    for (let i = 0; i + 1 < decoded.length; i++) {
-        const multiply = decoded[i];
-        const add = decoded[i + 1];
-
-        if (multiply.op !== VM_OP.MUL || add.op !== VM_OP.ADD || multiply.width !== add.width) {
-            continue;
-        }
-
-        const product = multiply.write.base;
-        const span = multiply.write.span;
-
-        // The addend is whichever operand of the ADD is not the product.
-        const addend = add.reads.find(read => read.base !== product);
-        const uses_product = add.reads.some(read => read.base === product);
-
-        if (!uses_product || addend === undefined) {
-            continue;
-        }
-
-        // The product must be nobody else's: the MAD does not leave it in a register.
-        if (product !== add.write.base && !particle_run_is_dead_after(decoded, i + 1, product, span)) {
-            continue;
-        }
-
-        const words = [
-            vm_encode_word0(VM_OP.MAD, multiply.width),
-            add.write.base,
-            multiply.reads[0].base,
-            multiply.reads[1].base,
-            addend.base,
-        ];
-
-        if (!particle_instruction_is_well_formed(words)) {
-            continue;
-        }
-
-        const out = instructions.slice();
-
-        out[i] = words;
-        out.splice(i + 1, 1);
-
-        return out;
-    }
-
-    return null;
-}
+import { VM_KIND, VM_MAX_WIDTH, VM_OP, vm_f32_value, vm_immediate, vm_pack_instruction, vm_register } from "../isa/ParticleVMISA.js";
+import { PARTICLE_EXACT_INTEGER_LIMIT, ParticleIntegerBounds } from "./particle_integer_bounds.js";
+import {
+    decode_particle_instructions,
+    particle_instruction_is_well_formed,
+    particle_run_is_dead_after,
+    particle_run_is_touched_between,
+    particle_run_is_written_between,
+    particle_runs_overlap,
+} from "./particle_program_analysis.js";
+
+/** @type {string} */
+export const PARTICLE_PASS_PEEPHOLE = "peephole";
+
+/**
+ * Local rewrites over a couple of neighbouring instructions — the ones that need no analysis beyond
+ * "these two are really one".
+ *
+ * ## Constant loads
+ *
+ * A frontend emits one load per literal, because it lowers one value at a time. Where two of them
+ * land in adjacent registers from adjacent pool words, they are one load: the pool is addressed by
+ * word, so a two-word run is as loadable as a one-word one. The graph compiler's `mad` node — a
+ * multiplier and an addend, interned side by side — is exactly this shape.
+ *
+ * ## Copies
+ *
+ * The same rewrite over `MOV`, and for the same reason: a `vec3` node is a copy per lane, because
+ * writing one lane on its own is the only thing a lane-at-a-time constructor can do, and three
+ * copies of consecutive slots into consecutive slots are one copy of the run, in whatever order the
+ * lanes were filled. Two copies of the same literal into adjacent slots are one copy too. It matters beyond the instructions it saves — a wide
+ * copy is a whole value moving, which is the shape {@link ./pass_propagate_copies.js} can then
+ * forward into its producer and delete outright, and three narrow ones are not.
+ *
+ * ## Fusing a multiply and an add
+ *
+ * `MUL` then `ADD` (or `SUB`) is `MAD`, one instruction rather than two. It is also, in general, a
+ * DIFFERENT NUMBER: the pair rounds the product to `f32` before adding, and `MAD` does not — and on
+ * hardware that contracts it into an FMA, does not round it at all. That is the ordinary compiler
+ * trade and it is a perfectly reasonable one to take for a particle effect, but it is not the
+ * bit-identity the rest of this optimizer holds itself to, so it is a switch rather than a default:
+ * `{ fuse_multiply_add: true }`.
+ *
+ * The exception, on by default, is a product that provably needs no rounding: an integer within
+ * 2^24, which {@link ParticleIntegerBounds} proves for the masks effects are made of —
+ * `cmp * 2 - 1`, `sign * step + p`, a lane count times a flag. Those pairs are fused with nothing
+ * given up, and `fuse_exact_multiply_add: false` switches even that off.
+ *
+ * @type {import("./optimize_particle_program.js").ParticleOptimizerPass}
+ */
+export function pass_peephole(instructions, context) {
+    const options = {
+        merge_constant_loads: true,
+        merge_moves: true,
+        fuse_exact_multiply_add: true,
+        fuse_multiply_add: false,
+        ...context,
+    };
+
+    let current = instructions;
+    let changed = false;
+
+    for (const rewrite of [
+        options.merge_constant_loads ? merge_one_constant_load : null,
+        options.merge_moves ? merge_one_move : null,
+        options.fuse_exact_multiply_add ? (stream) => fuse_one_multiply_add(stream, context, true) : null,
+        options.fuse_multiply_add ? (stream) => fuse_one_multiply_add(stream, context, false) : null,
+    ]) {
+        if (rewrite === null) {
+            continue;
+        }
+
+        // Applied one at a time and re-derived: each rewrite deletes an instruction, and re-deriving
+        // a program of this size is cheaper than keeping indices and liveness honest across a splice.
+        for (let attempt = 0; attempt < instructions.length; attempt++) {
+            const result = rewrite(current);
+
+            if (result === null) {
+                break;
+            }
+
+            current = result;
+            changed = true;
+        }
+    }
+
+    return changed ? current : null;
+}
+
+/**
+ * Two constant loads that fill adjacent registers from adjacent pool words are one load. The pool
+ * cannot change under them, so there is nothing to ask about the source.
+ *
+ * @param {number[][]} instructions
+ * @returns {number[][]|null}
+ */
+function merge_one_constant_load(instructions) {
+    return merge_one_pair(
+        instructions,
+        VM_OP.LOAD_CONST,
+        (a, b) => (b.words[1] === a.words[1] + a.operation_width ? a.words[1] : null),
+        () => true
+    );
+}
+
+/**
+ * Two copies that fill adjacent registers from adjacent registers — or from one literal — are one
+ * copy.
+ *
+ * The second one's register source has to survive the move: it is read at the first one's position
+ * now, so anything that OVERWRITES it between the two — the first copy itself included, which is
+ * what `MOV r1, r0; MOV r2, r1` is — makes the merge a different program. Reads in between are
+ * nobody's business, and a literal cannot change at all.
+ *
+ * @param {number[][]} instructions
+ * @returns {number[][]|null}
+ */
+function merge_one_move(instructions) {
+    return merge_one_pair(
+        instructions,
+        VM_OP.MOV,
+        (a, b) => {
+            const [from_a] = a.sources;
+            const [from_b] = b.sources;
+
+            if (from_a.kind !== from_b.kind) {
+                return null;
+            }
+
+            if (from_a.kind === VM_KIND.IMM) {
+                return from_a.word === from_b.word ? from_a.word : null;
+            }
+
+            // Plain runs only: a swizzled copy is not a run of its source.
+            const plain = from_a.word === vm_register(from_a.base) && from_b.word === vm_register(from_b.base);
+
+            return plain && from_b.base === from_a.base + a.operation_width ? from_a.word : null;
+        },
+        (decoded, first, second, b) => b.sources[0].kind === VM_KIND.IMM
+            || !particle_run_is_written_between(decoded, first, second, b.sources[0].base, b.operation_width)
+    );
+}
+
+/**
+ * Two instructions of one opcode that fill adjacent destination runs from sources that combine, with
+ * nothing in between to stop them being one instruction of the combined width. Which of them fills
+ * the lower lanes does not matter; the combined instruction starts where the lower one did.
+ *
+ * The second one moves back to the first one's position, so nothing between them may read or write
+ * the run the two of them now fill together — and whatever else that opcode needs of the gap is the
+ * caller's to state.
+ *
+ * @param {number[][]} instructions
+ * @param {number} op the opcode both instructions carry: one that writes its width from its first
+ *   source alone, `MOV` or `LOAD_CONST`
+ * @param {function(import("./particle_program_analysis.js").ParticleDecodedInstruction, import("./particle_program_analysis.js").ParticleDecodedInstruction): number|null} combined_source
+ *   given the instruction filling the lower lanes and the one filling the upper, the first source
+ *   word of the merged instruction, or `null` when the two sources do not combine
+ * @param {function(import("./particle_program_analysis.js").ParticleDecodedInstruction[], number, number, import("./particle_program_analysis.js").ParticleDecodedInstruction): boolean} source_survives
+ *   whether the second instruction's source still holds what it did, read at the first's position
+ * @returns {number[][]|null}
+ */
+function merge_one_pair(instructions, op, combined_source, source_survives) {
+    const decoded = decode_particle_instructions(instructions);
+
+    for (let first = 0; first < decoded.length; first++) {
+        const a = decoded[first];
+
+        if (a.op !== op) {
+            continue;
+        }
+
+        for (let second = first + 1; second < decoded.length; second++) {
+            const b = decoded[second];
+
+            if (b.op === VM_OP.HALT) {
+                break;
+            }
+
+            if (b.op !== op || a.operation_width + b.operation_width > VM_MAX_WIDTH) {
+                continue;
+            }
+
+            // Either may be the one filling the lower lanes: a run built from the top down is the
+            // same run, and a lane-at-a-time lowering fills whichever lane it reaches first.
+            const above = b.dst === a.dst + a.operation_width;
+            const below = b.dst + b.operation_width === a.dst;
+
+            if (!above && !below) {
+                continue;
+            }
+
+            const [lower, upper] = above ? [a, b] : [b, a];
+            const source = combined_source(lower, upper);
+
+            if (source === null) {
+                continue;
+            }
+
+            const base = lower.dst;
+            const span = a.operation_width + b.operation_width;
+
+            if (particle_run_is_touched_between(decoded, first + 1, second, base, span)
+                || !source_survives(decoded, first, second, b)) {
+                continue;
+            }
+
+            const kind = a.sources.length > 0 ? a.sources[0].kind : VM_KIND.NONE;
+            const words = vm_pack_instruction(op, span, base, [kind, VM_KIND.NONE, VM_KIND.NONE], [source, 0, 0]);
+
+            if (!particle_instruction_is_well_formed(words)) {
+                continue;
+            }
+
+            const out = instructions.slice();
+
+            out[first] = words;
+            out.splice(second, 1);
+
+            return out;
+        }
+    }
+
+    return null;
+}
+
+/**
+ * `MUL` into a register that an `ADD` or a `SUB` below consumes is `MAD`.
+ *
+ * The consumer may be anywhere below the multiply, as long as the product is nobody else's on the
+ * way — nothing between them reads or writes it — and the two factors still hold what they did, so
+ * the fused instruction reads at the consumer's position what the multiply read at its own. The
+ * consumer has to read the product as it lies: a swizzled read reorders lanes the multiply computed
+ * in order. Afterwards the product must be dead, since the `MAD` does not leave it anywhere.
+ *
+ * With `exact`, only a product {@link ParticleIntegerBounds} proves to be an integer within
+ * {@link PARTICLE_EXACT_INTEGER_LIMIT} is fused. Such a product needs no rounding, so the fused
+ * form is the same number as the pair on every executor, FMA or not. Without it any pair is fused,
+ * which is the `fuse_multiply_add` switch and a different number.
+ *
+ * A `SUB` is fused by negating a literal, which is exact: `p - c` is `MAD(a, b, -c)` when `c` is a
+ * literal, and `c - p` is `MAD(a, -b, c)` when one factor is. A register cannot be negated in place,
+ * so the other shapes are left.
+ *
+ * @param {number[][]} instructions
+ * @param {import("./optimize_particle_program.js").ParticleOptimizerContext} context
+ * @param {boolean} exact
+ * @returns {number[][]|null}
+ */
+function fuse_one_multiply_add(instructions, context, exact) {
+    const decoded = decode_particle_instructions(instructions);
+    const bounds = exact ? new ParticleIntegerBounds(context.constants ?? [], context.register_slots) : null;
+
+    for (let i = 0; i < decoded.length; i++) {
+        const multiply = decoded[i];
+
+        if (multiply.op === VM_OP.MUL && (bounds === null || product_is_exact(multiply, bounds))) {
+            const fused = fuse_multiply_into_consumer(decoded, i);
+
+            if (fused !== null) {
+                const out = instructions.slice();
+
+                out[fused.consumer] = fused.words;
+                out.splice(i, 1);
+
+                return out;
+            }
+        }
+
+        // Published after the search, so the factors are bounded as the multiply read them.
+        bounds?.publish(multiply.words);
+    }
+
+    return null;
+}
+
+/**
+ * @param {import("./particle_program_analysis.js").ParticleDecodedInstruction} multiply
+ * @param {ParticleIntegerBounds} bounds as they stand before the multiply
+ * @returns {boolean} whether its product is an integer small enough to need no rounding
+ */
+function product_is_exact(multiply, bounds) {
+    const [a, b] = multiply.sources;
+    const width = multiply.operation_width;
+
+    // `<=` rather than `!(>)`: an unknown factor times a zero bound is `NaN`, and not exact.
+    return bounds.bound_of_source(a, width) * bounds.bound_of_source(b, width) <= PARTICLE_EXACT_INTEGER_LIMIT;
+}
+
+/**
+ * The consumer of the product of the `MUL` at `index`, and the `MAD` that replaces it there.
+ *
+ * @param {import("./particle_program_analysis.js").ParticleDecodedInstruction[]} decoded
+ * @param {number} index
+ * @returns {{consumer: number, words: number[]}|null} `null` when the product is read by something
+ *   other than one fusable consumer, or the factors do not survive to it
+ */
+function fuse_multiply_into_consumer(decoded, index) {
+    const multiply = decoded[index];
+    const product = multiply.write.base;
+    const span = multiply.write.span;
+    const is_product = (source) => source.kind === VM_KIND.REG && source.word === vm_register(product);
+    const touches_product = (source) => source.kind === VM_KIND.REG
+        && particle_runs_overlap(source.base, source.span, product, span);
+
+    for (let j = index + 1; j < decoded.length; j++) {
+        const consumer = decoded[j];
+
+        if (consumer.op === VM_OP.HALT) {
+            return null;
+        }
+
+        const position = (consumer.op === VM_OP.ADD || consumer.op === VM_OP.SUB) && consumer.width === multiply.width
+            ? consumer.sources.findIndex(is_product)
+            : -1;
+
+        if (position < 0) {
+            // Anything else that touches the product wants it as a value of its own.
+            if ((consumer.write !== null && particle_runs_overlap(consumer.write.base, consumer.write.span, product, span))
+                || consumer.reads.some(touches_product)) {
+                return null;
+            }
+
+            continue;
+        }
+
+        const other = consumer.sources[1 - position];
+
+        if (touches_product(other)) {
+            return null;
+        }
+
+        const factors_survive = multiply.reads.every(
+            factor => !particle_run_is_written_between(decoded, index + 1, j, factor.base, factor.span)
+        );
+        const covered = consumer.write.base <= product && consumer.write.base + consumer.write.span >= product + span;
+
+        if (!factors_survive || (!covered && !particle_run_is_dead_after(decoded, j, product, span))) {
+            return null;
+        }
+
+        const words = fused_words(multiply, consumer, position, other);
+
+        return words !== null && particle_instruction_is_well_formed(words) ? { consumer: j, words } : null;
+    }
+
+    return null;
+}
+
+/**
+ * @param {import("./particle_program_analysis.js").ParticleDecodedInstruction} multiply
+ * @param {import("./particle_program_analysis.js").ParticleDecodedInstruction} consumer an `ADD` or `SUB`
+ * @param {number} position which of the consumer's operands is the product
+ * @param {import("./particle_program_analysis.js").ParticleOperandRun} other its other operand
+ * @returns {number[]|null} the `MAD`, or `null` when the shape needs a register negated
+ */
+function fused_words(multiply, consumer, position, other) {
+    let [a, b] = multiply.sources;
+    let addend = other;
+
+    if (consumer.op === VM_OP.SUB) {
+        if (position === 0) {
+            // p - c: negate the addend.
+            if (addend.kind !== VM_KIND.IMM) {
+                return null;
+            }
+
+            addend = negated(addend);
+        } else if (b.kind === VM_KIND.IMM) {
+            // c - p: negate a factor.
+            b = negated(b);
+        } else if (a.kind === VM_KIND.IMM) {
+            a = negated(a);
+        } else {
+            return null;
+        }
+    }
+
+    return vm_pack_instruction(
+        VM_OP.MAD, multiply.width, consumer.write.base,
+        [a.kind, b.kind, addend.kind],
+        [a.word, b.word, addend.word]
+    );
+}
+
+/**
+ * @param {import("./particle_program_analysis.js").ParticleOperandRun} literal an `IMM` operand
+ * @returns {{kind: number, word: number}} the same literal, negated
+ */
+function negated(literal) {
+    return vm_immediate(-vm_f32_value(literal.word));
+}