diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts index 6ba3b5dcc3e60..a93c7e256d5a5 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/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 }); diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts index 2a94eec79bdc0..45344e72662db 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -265,6 +265,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. diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts index d544269869297..5d1a14711a02d 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts +++ b/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; diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts new file mode 100644 index 0000000000000..37619b4224787 --- /dev/null +++ b/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; +type Node = { + instruction: Instruction | null; + dependencies: Set; + depth: number | null; +}; + +function reorderBlock( + env: Environment, + block: BasicBlock, + shared: Nodes +): void { + const locals: Nodes = new Map(); + const named: Map = 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 = []; + const seen = new Set(); + + 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, + id: IdentifierId, + depth: number = 0 +): void { + if (seen.has(id)) { + console.log(`${"| ".repeat(depth)}$${id} `); + 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 ""; + } + 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, + 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; + } + } +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts index 2c9e67646b155..23a0a839ea4b1 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts +++ b/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": { diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeFbtOperandsInSameScope.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeFbtOperandsInSameScope.ts index daab49d22b267..7bca5a5ce463d 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeFbtOperandsInSameScope.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeFbtOperandsInSameScope.ts @@ -94,6 +94,7 @@ function visit(fn: HIRFunction, fbtValues: Set): void { operand.identifier.mutableRange.start ) ); + fbtValues.add(operand.identifier.id); } } else if ( isFbtJsxExpression(fbtValues, value) || @@ -126,6 +127,33 @@ function visit(fn: HIRFunction, fbtValues: Set): void { */ 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 + ) + ); + } } } } diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-reordering.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-reordering.expect.md new file mode 100644 index 0000000000000..dc49af4401613 --- /dev/null +++ b/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 ( +
+ + {state} + +
+ ); +} + +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 = ; + $[2] = t1; + } else { + t1 = $[2]; + } + let t2; + if ($[3] !== state) { + t2 = {state}; + $[3] = state; + $[4] = t2; + } else { + t2 = $[4]; + } + let t3; + if ($[5] !== t0) { + t3 = ( + + ); + $[5] = t0; + $[6] = t3; + } else { + t3 = $[6]; + } + let t4; + if ($[7] !== t2 || $[8] !== t3) { + t4 = ( +
+ {t1} + {t2} + {t3} +
+ ); + $[7] = t2; + $[8] = t3; + $[9] = t4; + } else { + t4 = $[9]; + } + return t4; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: 42 }], +}; + +``` + +### Eval output +(kind: ok)
{"text":"Counter"}
0
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-reordering.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-reordering.js new file mode 100644 index 0000000000000..ad566a062c773 --- /dev/null +++ b/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 ( +
+ + {state} + +
+ ); +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: 42 }], +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.expect.md index eacff50f8892c..a2e79ad67f095 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.expect.md +++ b/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() { diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.js index 2775c83acedd5..89c3e2ddb6da6 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.js +++ b/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() {