Update base for Update on "[compiler] Optimize emission in normal (non-value) 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, though i still need to double-check all the output changes.

[ghstack-poisoned]

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

@@ -13,11 +13,15 @@ import {
   HIRFunction,
   IdentifierId,
   Instruction,
+  InstructionId,
+  Place,
   isExpressionBlockKind,
+  makeInstructionId,
   markInstructionIds,
 } from "../HIR";
 import { printInstruction } from "../HIR/PrintHIR";
 import {
+  eachInstructionLValue,
   eachInstructionValueLValue,
   eachInstructionValueOperand,
   eachTerminalOperand,
@@ -69,8 +73,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`,
@@ -91,10 +96,76 @@ type Node = {
   depth: number | null;
 };
 
+// Inclusive start and end
+type References = {
+  accessedRanges: AccessedRanges;
+  lastAssignments: LastAssignments;
+};
+type LastAssignments = Map<string, InstructionId>;
+type AccessedRanges = Map<IdentifierId, Range>;
+type Range = { start: InstructionId; end: InstructionId };
+enum ReferenceKind {
+  Read,
+  Write,
+}
+function findReferencedRangeOfTemporaries(fn: HIRFunction): References {
+  const accessedRanges: AccessedRanges = new Map();
+  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 range = getOrInsertWith(
+        accessedRanges,
+        place.identifier.id,
+        () => ({
+          start: instr,
+          end: instr,
+        })
+      );
+      range.end = instr;
+    }
+  }
+  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 { accessedRanges, 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();
@@ -116,7 +187,7 @@ function reorderBlock(
      * Ensure non-reoderable instructions have their order retained by
      * adding explicit dependencies to the previous such instruction.
      */
-    if (getReoderability(instr) === Reorderability.Nonreorderable) {
+    if (getReoderability(instr, references) === Reorderability.Nonreorderable) {
       if (previous !== null) {
         node.dependencies.add(previous);
       }
@@ -220,7 +291,8 @@ function reorderBlock(
     }
     CompilerError.invariant(
       node.instruction != null &&
-        getReoderability(node.instruction) === Reorderability.Reorderable,
+        getReoderability(node.instruction, references) ===
+          Reorderability.Reorderable,
       {
         reason: `Expected all remaining instructions to be reorderable`,
         loc: node.instruction?.loc ?? block.terminal.loc,
@@ -334,7 +406,10 @@ enum Reorderability {
   Reorderable,
   Nonreorderable,
 }
-function getReoderability(instr: Instruction): Reorderability {
+function getReoderability(
+  instr: Instruction,
+  references: References
+): Reorderability {
   switch (instr.value.kind) {
     case "JsxExpression":
     case "JsxFragment":
@@ -346,6 +421,23 @@ 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);
+        const range = references.accessedRanges.get(instr.lvalue.identifier.id);
+        if (
+          lastAssignment !== undefined &&
+          lastAssignment < instr.id &&
+          range !== undefined &&
+          range.end === range.start // this LoadLocal is used exactly once
+        ) {
+          console.log(`reorderable: ${name.value}`);
+          return Reorderability.Reorderable;
+        }
+      }
+      return Reorderability.Nonreorderable;
+    }
     default: {
       return Reorderability.Nonreorderable;
     }