[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-source-id: 98824627b66f7a43aeaf141c21efddc60c3cc0b3
Pull Request resolved: #29883

Author: josephsavona

compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts

@@ -281,7 +281,7 @@ const EnvironmentConfigSchema = z.object({
    * Enable instruction reordering. See InstructionReordering.ts for the details
    * of the approach.
    */
-  enableInstructionReordering: z.boolean().default(false),
+  enableInstructionReordering: z.boolean().default(true),
 
   /*
    * Enables instrumentation codegen. This emits a dev-mode only call to an

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

@@ -26,7 +26,6 @@ import {
   eachInstructionValueOperand,
   eachTerminalOperand,
 } from "../HIR/visitors";
-import { mayAllocate } from "../ReactiveScopes/InferReactiveScopeVariables";
 import { getOrInsertWith } from "../Utils/utils";
 
 /**
@@ -93,6 +92,7 @@ type Nodes = Map<IdentifierId, Node>;
 type Node = {
   instruction: Instruction | null;
   dependencies: Set<IdentifierId>;
+  reorderability: Reorderability;
   depth: number | null;
 };
 
@@ -173,21 +173,23 @@ function reorderBlock(
   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, references) === Reorderability.Nonreorderable) {
+    if (reorderability === Reorderability.Nonreorderable) {
       if (previous !== null) {
         node.dependencies.add(previous);
       }
@@ -243,67 +245,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, references) ===
-          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 don't just mean moving
+     * them later: we can also move then _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;
@@ -319,8 +379,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);
   }
@@ -355,7 +414,7 @@ function print(
     print(env, locals, shared, seen, dep, depth + 1);
   }
   console.log(
-    `${"|   ".repeat(depth)}$${id} ${printNode(node)} deps=[${deps.map((x) => `$${x}`).join(", ")}]`
+    `${"|   ".repeat(depth)}$${id} ${printNode(node)} deps=[${deps.map((x) => `$${x}`).join(", ")}] depth=${node.depth}`
   );
 }
 
@@ -406,7 +465,7 @@ enum Reorderability {
   Reorderable,
   Nonreorderable,
 }
-function getReoderability(
+function getReorderability(
   instr: Instruction,
   references: References
 ): Reorderability {
@@ -432,7 +491,6 @@ function getReoderability(
           range !== undefined &&
           range.end === range.start // this LoadLocal is used exactly once
         ) {
-          console.log(`reorderable: ${name.value}`);
           return Reorderability.Reorderable;
         }
       }

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

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-computed-load.expect.md

@@ -19,19 +19,20 @@ function component(a) {
 import { c as _c } from "react/compiler-runtime";
 function component(a) {
   const $ = _c(2);
-  let x;
+  let t0;
   if ($[0] !== a) {
-    x = { a };
+    const x = { a };
     const y = {};
 
+    t0 = x;
     y.x = x.a;
     mutate(y);
     $[0] = a;
-    $[1] = x;
+    $[1] = t0;
   } else {
-    x = $[1];
+    t0 = $[1];
   }
-  return x;
+  return t0;
 }
 
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-nested-member-path-mutate.expect.md

@@ -20,19 +20,21 @@ function component() {
 import { c as _c } from "react/compiler-runtime";
 function component() {
   const $ = _c(1);
-  let x;
+  let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     const z = [];
     const y = {};
     y.z = z;
-    x = {};
+    const x = {};
     x.y = y;
+
+    t0 = x;
     mutate(x.y.z);
-    $[0] = x;
+    $[0] = t0;
   } else {
-    x = $[0];
+    t0 = $[0];
   }
-  return x;
+  return t0;
 }
 
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-nested-member-path.expect.md

@@ -25,18 +25,20 @@ export const FIXTURE_ENTRYPOINT = {
 import { c as _c } from "react/compiler-runtime";
 function component() {
   const $ = _c(1);
-  let x;
+  let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     const z = [];
     const y = {};
     y.z = z;
-    x = {};
+    const x = {};
+
+    t0 = x;
     x.y = y;
-    $[0] = x;
+    $[0] = t0;
   } else {
-    x = $[0];
+    t0 = $[0];
   }
-  return x;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/align-scopes-within-nested-valueblock-in-array.expect.md

@@ -52,11 +52,12 @@ import { Stringify, identity, makeArray, mutate } from "shared-runtime";
  * handles this correctly.
  */
 function Foo(t0) {
-  const $ = _c(4);
+  const $ = _c(3);
   const { cond1, cond2 } = t0;
-  const arr = makeArray({ a: 2 }, 2, []);
   let t1;
-  if ($[0] !== cond1 || $[1] !== cond2 || $[2] !== arr) {
+  if ($[0] !== cond1 || $[1] !== cond2) {
+    const arr = makeArray({ a: 2 }, 2, []);
+
     t1 = cond1 ? (
       <>
         <div>{identity("foo")}</div>
@@ -65,10 +66,9 @@ function Foo(t0) {
     ) : null;
     $[0] = cond1;
     $[1] = cond2;
-    $[2] = arr;
-    $[3] = t1;
+    $[2] = t1;
   } else {
-    t1 = $[3];
+    t1 = $[2];
   }
   return t1;
 }