[compiler] Optimize instruction reordering

compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts

@@ -13,16 +13,19 @@ import {
   HIRFunction,
   IdentifierId,
   Instruction,
+  InstructionId,
+  Place,
   isExpressionBlockKind,
+  makeInstructionId,
   markInstructionIds,
 } from "../HIR";
 import { printInstruction } from "../HIR/PrintHIR";
 import {
+  eachInstructionLValue,
   eachInstructionValueLValue,
   eachInstructionValueOperand,
   eachTerminalOperand,
 } from "../HIR/visitors";
-import { mayAllocate } from "../ReactiveScopes/InferReactiveScopeVariables";
 import { getOrInsertWith } from "../Utils/utils";
 
 /**
@@ -69,8 +72,9 @@ import { getOrInsertWith } from "../Utils/utils";
 export function instructionReordering(fn: HIRFunction): void {
   // Shared nodes are emitted when they are first used
   const shared: Nodes = new Map();
+  const references = findReferencedRangeOfTemporaries(fn);
   for (const [, block] of fn.body.blocks) {
-    reorderBlock(fn.env, block, shared);
+    reorderBlock(fn.env, block, shared, references);
   }
   CompilerError.invariant(shared.size === 0, {
     reason: `InstructionReordering: expected all reorderable nodes to have been emitted`,
@@ -88,35 +92,103 @@ type Nodes = Map<IdentifierId, Node>;
 type Node = {
   instruction: Instruction | null;
   dependencies: Set<IdentifierId>;
+  reorderability: Reorderability;
   depth: number | null;
 };
 
+// Inclusive start and end
+type References = {
+  singleUseIdentifiers: SingleUseIdentifiers;
+  lastAssignments: LastAssignments;
+};
+type LastAssignments = Map<string, InstructionId>;
+type SingleUseIdentifiers = Set<IdentifierId>;
+enum ReferenceKind {
+  Read,
+  Write,
+}
+function findReferencedRangeOfTemporaries(fn: HIRFunction): References {
+  const singleUseIdentifiers = new Map<IdentifierId, number>();
+  const lastAssignments: LastAssignments = new Map();
+  function reference(
+    instr: InstructionId,
+    place: Place,
+    kind: ReferenceKind
+  ): void {
+    if (
+      place.identifier.name !== null &&
+      place.identifier.name.kind === "named"
+    ) {
+      if (kind === ReferenceKind.Write) {
+        const name = place.identifier.name.value;
+        const previous = lastAssignments.get(name);
+        if (previous === undefined) {
+          lastAssignments.set(name, instr);
+        } else {
+          lastAssignments.set(
+            name,
+            makeInstructionId(Math.max(previous, instr))
+          );
+        }
+      }
+      return;
+    } else if (kind === ReferenceKind.Read) {
+      const previousCount = singleUseIdentifiers.get(place.identifier.id) ?? 0;
+      singleUseIdentifiers.set(place.identifier.id, previousCount + 1);
+    }
+  }
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      for (const operand of eachInstructionValueLValue(instr.value)) {
+        reference(instr.id, operand, ReferenceKind.Read);
+      }
+      for (const lvalue of eachInstructionLValue(instr)) {
+        reference(instr.id, lvalue, ReferenceKind.Write);
+      }
+    }
+    for (const operand of eachTerminalOperand(block.terminal)) {
+      reference(block.terminal.id, operand, ReferenceKind.Read);
+    }
+  }
+  return {
+    singleUseIdentifiers: new Set(
+      [...singleUseIdentifiers]
+        .filter(([, count]) => count === 1)
+        .map(([id]) => id)
+    ),
+    lastAssignments,
+  };
+}
+
 function reorderBlock(
   env: Environment,
   block: BasicBlock,
-  shared: Nodes
+  shared: Nodes,
+  references: References
 ): void {
   const locals: Nodes = new Map();
   const named: Map<string, IdentifierId> = new Map();
   let previous: IdentifierId | null = null;
   for (const instr of block.instructions) {
     const { lvalue, value } = instr;
     // Get or create a node for this lvalue
+    const reorderability = getReorderability(instr, references);
     const node = getOrInsertWith(
       locals,
       lvalue.identifier.id,
       () =>
         ({
           instruction: instr,
           dependencies: new Set(),
+          reorderability,
           depth: null,
         }) as Node
     );
     /**
      * Ensure non-reoderable instructions have their order retained by
      * adding explicit dependencies to the previous such instruction.
      */
-    if (getReoderability(instr) === Reorderability.Nonreorderable) {
+    if (reorderability === Reorderability.Nonreorderable) {
       if (previous !== null) {
         node.dependencies.add(previous);
       }
@@ -172,66 +244,125 @@ function reorderBlock(
 
   DEBUG && console.log(`bb${block.id}`);
 
-  // First emit everything that can't be reordered
-  if (previous !== null) {
-    DEBUG && console.log(`(last non-reorderable instruction)`);
-    DEBUG && print(env, locals, shared, seen, previous);
-    emit(env, locals, shared, nextInstructions, previous);
-  }
-  /*
-   * For "value" blocks the final instruction represents its value, so we have to be
-   * careful to not change the ordering. Emit the last instruction explicitly.
-   * Any non-reorderable instructions will get emitted first, and any unused
-   * reorderable instructions can be deferred to the shared node list.
+  /**
+   * The ideal order for emitting instructions may change the final instruction,
+   * but value blocks have special semantics for the final instruction of a block -
+   * that's the expression's value!. So we choose between a less optimal strategy
+   * for value blocks which preserves the final instruction order OR a more optimal
+   * ordering for statement-y blocks.
    */
-  if (isExpressionBlockKind(block.kind) && block.instructions.length !== 0) {
-    DEBUG && console.log(`(block value)`);
-    DEBUG &&
-      print(
+  if (isExpressionBlockKind(block.kind)) {
+    // First emit everything that can't be reordered
+    if (previous !== null) {
+      DEBUG && console.log(`(last non-reorderable instruction)`);
+      DEBUG && print(env, locals, shared, seen, previous);
+      emit(env, locals, shared, nextInstructions, previous);
+    }
+    /*
+     * For "value" blocks the final instruction represents its value, so we have to be
+     * careful to not change the ordering. Emit the last instruction explicitly.
+     * Any non-reorderable instructions will get emitted first, and any unused
+     * reorderable instructions can be deferred to the shared node list.
+     */
+    if (block.instructions.length !== 0) {
+      DEBUG && console.log(`(block value)`);
+      DEBUG &&
+        print(
+          env,
+          locals,
+          shared,
+          seen,
+          block.instructions.at(-1)!.lvalue.identifier.id
+        );
+      emit(
         env,
         locals,
         shared,
-        seen,
+        nextInstructions,
         block.instructions.at(-1)!.lvalue.identifier.id
       );
-    emit(
-      env,
-      locals,
-      shared,
-      nextInstructions,
-      block.instructions.at(-1)!.lvalue.identifier.id
-    );
-  }
-  /*
-   * Then emit the dependencies of the terminal operand. In many cases they will have
-   * already been emitted in the previous step and this is a no-op.
-   * TODO: sort the dependencies based on weight, like we do for other nodes. Not a big
-   * deal though since most terminals have a single operand
-   */
-  for (const operand of eachTerminalOperand(block.terminal)) {
-    DEBUG && console.log(`(terminal operand)`);
-    DEBUG && print(env, locals, shared, seen, operand.identifier.id);
-    emit(env, locals, shared, nextInstructions, operand.identifier.id);
-  }
-  // Anything not emitted yet is globally reorderable
-  for (const [id, node] of locals) {
-    if (node.instruction == null) {
-      continue;
     }
-    CompilerError.invariant(
-      node.instruction != null &&
-        getReoderability(node.instruction) === Reorderability.Reorderable,
-      {
-        reason: `Expected all remaining instructions to be reorderable`,
-        loc: node.instruction?.loc ?? block.terminal.loc,
-        description:
-          node.instruction != null
-            ? `Instruction [${node.instruction.id}] was not emitted yet but is not reorderable`
-            : `Lvalue $${id} was not emitted yet but is not reorderable`,
+    /*
+     * Then emit the dependencies of the terminal operand. In many cases they will have
+     * already been emitted in the previous step and this is a no-op.
+     * TODO: sort the dependencies based on weight, like we do for other nodes. Not a big
+     * deal though since most terminals have a single operand
+     */
+    for (const operand of eachTerminalOperand(block.terminal)) {
+      DEBUG && console.log(`(terminal operand)`);
+      DEBUG && print(env, locals, shared, seen, operand.identifier.id);
+      emit(env, locals, shared, nextInstructions, operand.identifier.id);
+    }
+    // Anything not emitted yet is globally reorderable
+    for (const [id, node] of locals) {
+      if (node.instruction == null) {
+        continue;
       }
-    );
-    DEBUG && console.log(`save shared: $${id}`);
-    shared.set(id, node);
+      CompilerError.invariant(
+        node.reorderability === Reorderability.Reorderable,
+        {
+          reason: `Expected all remaining instructions to be reorderable`,
+          loc: node.instruction?.loc ?? block.terminal.loc,
+          description:
+            node.instruction != null
+              ? `Instruction [${node.instruction.id}] was not emitted yet but is not reorderable`
+              : `Lvalue $${id} was not emitted yet but is not reorderable`,
+        }
+      );
+
+      DEBUG && console.log(`save shared: $${id}`);
+      shared.set(id, node);
+    }
+  } else {
+    /**
+     * If this is not a value block, then the order within the block doesn't matter
+     * and we can optimize more. The observation is that blocks often have instructions
+     * such as:
+     *
+     * ```
+     * t$0 = nonreorderable
+     * t$1 = nonreorderable <-- this gets in the way of merging t$0 and t$2
+     * t$2 = reorderable deps[ t$0 ]
+     * return t$2
+     * ```
+     *
+     * Ie where there is some pair of nonreorderable+reorderable values, with some intervening
+     * also non-reorderable instruction. If we emit all non-reorderable instructions first,
+     * then we'll keep the original order. But reordering instructions doesn't just mean moving
+     * them later: we can also move them _earlier_. By starting from terminal operands we
+     * end up emitting:
+     *
+     * ```
+     * t$0 = nonreorderable // dep of t$2
+     * t$2 = reorderable deps[ t$0 ]
+     * t$1 = nonreorderable <-- not in the way of merging anymore!
+     * return t$2
+     * ```
+     *
+     * Ie all nonreorderable transitive deps of the terminal operands will get emitted first,
+     * but we'll be able to intersperse the depending reorderable instructions in between
+     * them in a way that works better with scope merging.
+     */
+    for (const operand of eachTerminalOperand(block.terminal)) {
+      DEBUG && console.log(`(terminal operand)`);
+      DEBUG && print(env, locals, shared, seen, operand.identifier.id);
+      emit(env, locals, shared, nextInstructions, operand.identifier.id);
+    }
+    // Anything not emitted yet is globally reorderable
+    for (const id of Array.from(locals.keys()).reverse()) {
+      const node = locals.get(id);
+      if (node === undefined) {
+        continue;
+      }
+      if (node.reorderability === Reorderability.Reorderable) {
+        DEBUG && console.log(`save shared: $${id}`);
+        shared.set(id, node);
+      } else {
+        DEBUG && console.log("leftover");
+        DEBUG && print(env, locals, shared, seen, id);
+        emit(env, locals, shared, nextInstructions, id);
+      }
+    }
   }
 
   block.instructions = nextInstructions;
@@ -247,8 +378,7 @@ function getDepth(env: Environment, nodes: Nodes, id: IdentifierId): number {
     return node.depth;
   }
   node.depth = 0; // in case of cycles
-  let depth =
-    node.instruction != null && mayAllocate(env, node.instruction) ? 1 : 0;
+  let depth = node.reorderability === Reorderability.Reorderable ? 1 : 10;
   for (const dep of node.dependencies) {
     depth += getDepth(env, nodes, dep);
   }
@@ -265,7 +395,7 @@ function print(
   depth: number = 0
 ): void {
   if (seen.has(id)) {
-    console.log(`${"|   ".repeat(depth)}$${id} <skipped>`);
+    DEBUG && console.log(`${"|   ".repeat(depth)}$${id} <skipped>`);
     return;
   }
   seen.add(id);
@@ -282,11 +412,12 @@ function print(
   for (const dep of deps) {
     print(env, locals, shared, seen, dep, depth + 1);
   }
-  console.log(
-    `${"|   ".repeat(depth)}$${id} ${printNode(node)} deps=[${deps
-      .map((x) => `$${x}`)
-      .join(", ")}]`
-  );
+  DEBUG &&
+    console.log(
+      `${"|   ".repeat(depth)}$${id} ${printNode(node)} deps=[${deps
+        .map((x) => `$${x}`)
+        .join(", ")}] depth=${node.depth}`
+    );
 }
 
 function printNode(node: Node): string {
@@ -336,7 +467,10 @@ enum Reorderability {
   Reorderable,
   Nonreorderable,
 }
-function getReoderability(instr: Instruction): Reorderability {
+function getReorderability(
+  instr: Instruction,
+  references: References
+): Reorderability {
   switch (instr.value.kind) {
     case "JsxExpression":
     case "JsxFragment":
@@ -348,6 +482,20 @@ function getReoderability(instr: Instruction): Reorderability {
     case "UnaryExpression": {
       return Reorderability.Reorderable;
     }
+    case "LoadLocal": {
+      const name = instr.value.place.identifier.name;
+      if (name !== null && name.kind === "named") {
+        const lastAssignment = references.lastAssignments.get(name.value);
+        if (
+          lastAssignment !== undefined &&
+          lastAssignment < instr.id &&
+          references.singleUseIdentifiers.has(instr.lvalue.identifier.id)
+        ) {
+          return Reorderability.Reorderable;
+        }
+      }
+      return Reorderability.Nonreorderable;
+    }
     default: {
       return Reorderability.Nonreorderable;
     }

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts

@@ -186,10 +186,7 @@ export function isMutable({ id }: Instruction, place: Place): boolean {
   return id >= range.start && id < range.end;
 }
 
-export function mayAllocate(
-  env: Environment,
-  instruction: Instruction
-): boolean {
+function mayAllocate(env: Environment, instruction: Instruction): boolean {
   const { value } = instruction;
   switch (value.kind) {
     case "Destructure": {