[compiler] Instruction reordering

Adds a pass just after DCE to reorder safely reorderable instructions (jsx, primitives, globals) closer to where they are used, to allow other optimization passes to be more effective. Notably, the reordering allows scope merging to be more effective, since that pass relies on two scopes not having intervening instructions — in many cases we can now reorder such instructions out of the way and unlock merging, as demonstrated in the changed fixtures.

The algorithm itself is described in the docblock.

note: This is a cleaned up version of #29579 that is ready for review.

ghstack-source-id: c54a806cad7aefba4ac1876c9fd9b25f9177e95a
Pull Request resolved: https://github.com/facebook/react/pull/29863

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -41,6 +41,7 @@ import {
   deadCodeElimination,
   pruneMaybeThrows,
 } from "../Optimization";
+import { instructionReordering } from "../Optimization/InstructionReordering";
 import {
   CodegenFunction,
   alignObjectMethodScopes,
@@ -204,6 +205,11 @@ function* runWithEnvironment(
   deadCodeElimination(hir);
   yield log({ kind: "hir", name: "DeadCodeElimination", value: hir });
 
+  if (env.config.enableInstructionReordering) {
+    instructionReordering(hir);
+    yield log({ kind: "hir", name: "InstructionReordering", value: hir });
+  }
+
   pruneMaybeThrows(hir);
   yield log({ kind: "hir", name: "PruneMaybeThrows", value: hir });
 

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

@@ -277,6 +277,12 @@ const EnvironmentConfigSchema = z.object({
 
   enableEmitHookGuards: ExternalFunctionSchema.nullish(),
 
+  /**
+   * Enable instruction reordering. See InstructionReordering.ts for the details
+   * of the approach.
+   */
+  enableInstructionReordering: z.boolean().default(false),
+
   /*
    * Enables instrumentation codegen. This emits a dev-mode only call to an
    * instrumentation function, for components and hooks that Forget compiles.

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

@@ -335,6 +335,28 @@ export type HIR = {
  * statements and not implicit exceptions which may occur.
  */
 export type BlockKind = "block" | "value" | "loop" | "sequence" | "catch";
+
+/**
+ * Returns true for "block" and "catch" block kinds which correspond to statements
+ * in the source, including BlockStatement, CatchStatement.
+ *
+ * Inverse of isExpressionBlockKind()
+ */
+export function isStatementBlockKind(kind: BlockKind): boolean {
+  return kind === "block" || kind === "catch";
+}
+
+/**
+ * Returns true for "value", "loop", and "sequence" block kinds which correspond to
+ * expressions in the source, such as ConditionalExpression, LogicalExpression, loop
+ * initializer/test/updaters, etc
+ *
+ * Inverse of isStatementBlockKind()
+ */
+export function isExpressionBlockKind(kind: BlockKind): boolean {
+  return !isStatementBlockKind(kind);
+}
+
 export type BasicBlock = {
   kind: BlockKind;
   id: BlockId;

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

@@ -0,0 +1,353 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { CompilerError } from "..";
+import {
+  BasicBlock,
+  Environment,
+  GeneratedSource,
+  HIRFunction,
+  IdentifierId,
+  Instruction,
+  isExpressionBlockKind,
+  markInstructionIds,
+} from "../HIR";
+import { printInstruction } from "../HIR/PrintHIR";
+import {
+  eachInstructionValueLValue,
+  eachInstructionValueOperand,
+  eachTerminalOperand,
+} from "../HIR/visitors";
+import { mayAllocate } from "../ReactiveScopes/InferReactiveScopeVariables";
+import { getOrInsertWith } from "../Utils/utils";
+
+/**
+ * This pass implements conservative instruction reordering to move instructions closer to
+ * to where their produced values are consumed. The goal is to group instructions in a way that
+ * is more optimal for future optimizations. Notably, MergeReactiveScopesThatAlwaysInvalidateTogether
+ * can only merge two candidate scopes if there are no intervenining instructions that are used by
+ * some later code: instruction reordering can move those intervening instructions later in many cases,
+ * thereby allowing more scopes to merge together.
+ *
+ * The high-level approach is to build a dependency graph where nodes correspond either to
+ * instructions OR to a particular lvalue assignment of another instruction. So
+ * `Destructure [x, y] = z` creates 3 nodes: one for the instruction, and one each for x and y.
+ * The lvalue nodes depend on the instruction node that assigns them.
+ *
+ * Dependency edges are added for all the lvalues and rvalues of each instruction, so for example
+ * the node for `t$2 = CallExpression t$0 ( t$1 )` will take dependencies on the nodes for t$0 and t$1.
+ *
+ * Individual instructions are grouped into two categories:
+ * - "Reorderable" instructions include a safe set of instructions that we know are fine to reorder.
+ *   This includes JSX elements/fragments/text, primitives, template literals, and globals.
+ *   These instructions are never emitted until they are referenced, and can even be moved across
+ *   basic blocks until they are used.
+ * - All other instructions are non-reorderable, and take an explicit dependency on the last such
+ *   non-reorderable instruction in their block. This largely ensures that mutations are serialized,
+ *   since all potentially mutating instructions are in this category.
+ *
+ * The only remaining mutation not handled by the above is variable reassignment. To ensure that all
+ * reads/writes of a variable access the correct version, all references (lvalues and rvalues) to
+ * each named variable are serialized. Thus `x = 1; y = x; x = 2; z = x` will establish a chain
+ * of dependencies and retain the correct ordering.
+ *
+ * The algorithm proceeds one basic block at a time, first building up the dependnecy graph and then
+ * reordering.
+ *
+ * The reordering weights nodes according to their transitive dependencies, and whether a particular node
+ * needs memoization or not. Larger dependencies go first, followed by smaller dependencies, which in
+ * testing seems to allow scopes to merge more effectively. Over time we can likely continue to improve
+ * the reordering heuristic.
+ *
+ * An obvious area for improvement is to allow reordering of LoadLocals that occur after the last write
+ * of the named variable. We can add this in a follow-up.
+ */
+export function instructionReordering(fn: HIRFunction): void {
+  // Shared nodes are emitted when they are first used
+  const shared: Nodes = new Map();
+  for (const [, block] of fn.body.blocks) {
+    reorderBlock(fn.env, block, shared);
+  }
+  CompilerError.invariant(shared.size === 0, {
+    reason: `InstructionReordering: expected all reorderable nodes to have been emitted`,
+    loc:
+      [...shared.values()]
+        .map((node) => node.instruction?.loc)
+        .filter((loc) => loc != null)[0] ?? GeneratedSource,
+  });
+  markInstructionIds(fn.body);
+}
+
+const DEBUG = false;
+
+type Nodes = Map<IdentifierId, Node>;
+type Node = {
+  instruction: Instruction | null;
+  dependencies: Set<IdentifierId>;
+  depth: number | null;
+};
+
+function reorderBlock(
+  env: Environment,
+  block: BasicBlock,
+  shared: Nodes
+): 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 node = getOrInsertWith(
+      locals,
+      lvalue.identifier.id,
+      () =>
+        ({
+          instruction: instr,
+          dependencies: new Set(),
+          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 (previous !== null) {
+        node.dependencies.add(previous);
+      }
+      previous = lvalue.identifier.id;
+    }
+    /**
+     * Establish dependencies on operands
+     */
+    for (const operand of eachInstructionValueOperand(value)) {
+      const { name, id } = operand.identifier;
+      if (name !== null && name.kind === "named") {
+        // Serialize all accesses to named variables
+        const previous = named.get(name.value);
+        if (previous !== undefined) {
+          node.dependencies.add(previous);
+        }
+        named.set(name.value, lvalue.identifier.id);
+      } else if (locals.has(id) || shared.has(id)) {
+        node.dependencies.add(id);
+      }
+    }
+    /**
+     * Establish nodes for lvalues, with dependencies on the node
+     * for the instruction itself. This ensures that any consumers
+     * of the lvalue will take a dependency through to the original
+     * instruction.
+     */
+    for (const lvalueOperand of eachInstructionValueLValue(value)) {
+      const lvalueNode = getOrInsertWith(
+        locals,
+        lvalueOperand.identifier.id,
+        () =>
+          ({
+            instruction: null,
+            dependencies: new Set(),
+            depth: null,
+          }) as Node
+      );
+      lvalueNode.dependencies.add(lvalue.identifier.id);
+      const name = lvalueOperand.identifier.name;
+      if (name !== null && name.kind === "named") {
+        const previous = named.get(name.value);
+        if (previous !== undefined) {
+          node.dependencies.add(previous);
+        }
+        named.set(name.value, lvalue.identifier.id);
+      }
+    }
+  }
+
+  const nextInstructions: Array<Instruction> = [];
+  const seen = new Set<IdentifierId>();
+
+  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.
+   */
+  if (isExpressionBlockKind(block.kind) && 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,
+      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`,
+      }
+    );
+    DEBUG && console.log(`save shared: $${id}`);
+    shared.set(id, node);
+  }
+
+  block.instructions = nextInstructions;
+  DEBUG && console.log();
+}
+
+function getDepth(env: Environment, nodes: Nodes, id: IdentifierId): number {
+  const node = nodes.get(id)!;
+  if (node == null) {
+    return 0;
+  }
+  if (node.depth != null) {
+    return node.depth;
+  }
+  node.depth = 0; // in case of cycles
+  let depth =
+    node.instruction != null && mayAllocate(env, node.instruction) ? 1 : 0;
+  for (const dep of node.dependencies) {
+    depth += getDepth(env, nodes, dep);
+  }
+  node.depth = depth;
+  return depth;
+}
+
+function print(
+  env: Environment,
+  locals: Nodes,
+  shared: Nodes,
+  seen: Set<IdentifierId>,
+  id: IdentifierId,
+  depth: number = 0
+): void {
+  if (seen.has(id)) {
+    console.log(`${"|   ".repeat(depth)}$${id} <skipped>`);
+    return;
+  }
+  seen.add(id);
+  const node = locals.get(id) ?? shared.get(id);
+  if (node == null) {
+    return;
+  }
+  const deps = [...node.dependencies];
+  deps.sort((a, b) => {
+    const aDepth = getDepth(env, locals, a);
+    const bDepth = getDepth(env, locals, b);
+    return bDepth - aDepth;
+  });
+  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(", ")}]`
+  );
+}
+
+function printNode(node: Node): string {
+  const { instruction } = node;
+  if (instruction === null) {
+    return "<lvalue-only>";
+  }
+  switch (instruction.value.kind) {
+    case "FunctionExpression":
+    case "ObjectMethod": {
+      return `[${instruction.id}] ${instruction.value.kind}`;
+    }
+    default: {
+      return printInstruction(instruction);
+    }
+  }
+}
+
+function emit(
+  env: Environment,
+  locals: Nodes,
+  shared: Nodes,
+  instructions: Array<Instruction>,
+  id: IdentifierId
+): void {
+  const node = locals.get(id) ?? shared.get(id);
+  if (node == null) {
+    return;
+  }
+  locals.delete(id);
+  shared.delete(id);
+  const deps = [...node.dependencies];
+  deps.sort((a, b) => {
+    const aDepth = getDepth(env, locals, a);
+    const bDepth = getDepth(env, locals, b);
+    return bDepth - aDepth;
+  });
+  for (const dep of deps) {
+    emit(env, locals, shared, instructions, dep);
+  }
+  if (node.instruction !== null) {
+    instructions.push(node.instruction);
+  }
+}
+
+enum Reorderability {
+  Reorderable,
+  Nonreorderable,
+}
+function getReoderability(instr: Instruction): Reorderability {
+  switch (instr.value.kind) {
+    case "JsxExpression":
+    case "JsxFragment":
+    case "JSXText":
+    case "LoadGlobal":
+    case "Primitive":
+    case "TemplateLiteral":
+    case "BinaryExpression":
+    case "UnaryExpression": {
+      return Reorderability.Reorderable;
+    }
+    default: {
+      return Reorderability.Nonreorderable;
+    }
+  }
+}

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

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

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

@@ -114,6 +114,7 @@ function visit(
               operand.identifier.mutableRange.start
             )
           );
+          fbtValues.add(operand.identifier.id);
         }
       } else if (
         isFbtJsxExpression(fbtMacroTags, fbtValues, value) ||
@@ -146,6 +147,33 @@ function visit(
            */
           fbtValues.add(operand.identifier.id);
         }
+      } else if (fbtValues.has(lvalue.identifier.id)) {
+        const fbtScope = lvalue.identifier.scope;
+        if (fbtScope === null) {
+          return;
+        }
+
+        for (const operand of eachReactiveValueOperand(value)) {
+          if (
+            operand.identifier.name !== null &&
+            operand.identifier.name.kind === "named"
+          ) {
+            /*
+             * named identifiers were already locals, we only have to force temporaries
+             * into the same scope
+             */
+            continue;
+          }
+          operand.identifier.scope = fbtScope;
+
+          // Expand the jsx element's range to account for its operands
+          fbtScope.range.start = makeInstructionId(
+            Math.min(
+              fbtScope.range.start,
+              operand.identifier.mutableRange.start
+            )
+          );
+        }
       }
     }
   }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-reordering.expect.md

@@ -0,0 +1,100 @@
+
+## Input
+
+```javascript
+// @enableInstructionReordering
+import { useState } from "react";
+import { Stringify } from "shared-runtime";
+
+function Component() {
+  let [state, setState] = useState(0);
+  return (
+    <div>
+      <Stringify text="Counter" />
+      <span>{state}</span>
+      <button data-testid="button" onClick={() => setState(state + 1)}>
+        increment
+      </button>
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableInstructionReordering
+import { useState } from "react";
+import { Stringify } from "shared-runtime";
+
+function Component() {
+  const $ = _c(10);
+  const [state, setState] = useState(0);
+  let t0;
+  if ($[0] !== state) {
+    t0 = () => setState(state + 1);
+    $[0] = state;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  let t1;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = <Stringify text="Counter" />;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  let t2;
+  if ($[3] !== state) {
+    t2 = <span>{state}</span>;
+    $[3] = state;
+    $[4] = t2;
+  } else {
+    t2 = $[4];
+  }
+  let t3;
+  if ($[5] !== t0) {
+    t3 = (
+      <button data-testid="button" onClick={t0}>
+        increment
+      </button>
+    );
+    $[5] = t0;
+    $[6] = t3;
+  } else {
+    t3 = $[6];
+  }
+  let t4;
+  if ($[7] !== t2 || $[8] !== t3) {
+    t4 = (
+      <div>
+        {t1}
+        {t2}
+        {t3}
+      </div>
+    );
+    $[7] = t2;
+    $[8] = t3;
+    $[9] = t4;
+  } else {
+    t4 = $[9];
+  }
+  return t4;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div><div>{"text":"Counter"}</div><span>0</span><button data-testid="button">increment</button></div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-reordering.js

@@ -0,0 +1,21 @@
+// @enableInstructionReordering
+import { useState } from "react";
+import { Stringify } from "shared-runtime";
+
+function Component() {
+  let [state, setState] = useState(0);
+  return (
+    <div>
+      <Stringify text="Counter" />
+      <span>{state}</span>
+      <button data-testid="button" onClick={() => setState(state + 1)}>
+        increment
+      </button>
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @enableInstructionReordering
 import { useState } from "react";
 
 function Component() {
@@ -22,7 +23,7 @@ function Component() {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @enableInstructionReordering
 import { useState } from "react";
 
 function Component() {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.js

@@ -1,3 +1,4 @@
+// @enableInstructionReordering
 import { useState } from "react";
 
 function Component() {