[compiler] Optimize instruction reordering

Note: due to a bad rebase i included #29883 here. Both were stamped so i'm not gonna bother splitting it back up aain.

This PR includes two changes:
* First, allow `LoadLocal` to be reordered if a) the load occurs after the last write to a variable and b) the LoadLocal lvalue is used exactly once
* Uses a more optimal reordering for statement blocks, while keeping the existing approach for expression blocks.

In #29863 I tried to find a clean way to share code for emitting instructions between value blocks and regular blocks. The catch is that value blocks have special meaning for their final instruction — that's the value of the block — so reordering can't change the last instruction. However, in finding a clean way to share code for these two categories of code, i also inadvertently reduced the effectiveness of the optimization.

This PR updates to use different strategies for these two kinds of blocks: value blocks use the code from #29863 where we first emit all non-reorderable instructions in their original order, then try to emit reorderable values. The reason this is suboptimal, though, is that we want to move instructions closer to their dependencies so that they can invalidate (merge) together. Emitting the reorderable values last prevents this.

So for normal blocks, we now emit terminal operands first. This will invariably cause some of the non-reorderable instructions to be emitted, but it will intersperse reoderable instructions in between, right after their dependencies. This maximizes our ability to merge scopes.

I think the complexity cost of two strategies is worth the benefit, as evidenced by the reduced memo slots in the fixtures.

ghstack-source-id: ad3e516fa474235ced8c5d56f4541d2a7c413608
Pull Request resolved: #29882

Author: josephsavona

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": {