diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts index 7b90740417..6708758bab 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts +++ b/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,13 +92,79 @@ type Nodes = Map; type Node = { instruction: Instruction | null; dependencies: Set; + reorderability: Reorderability; depth: number | null; }; +// Inclusive start and end +type References = { + singleUseIdentifiers: SingleUseIdentifiers; + lastAssignments: LastAssignments; +}; +type LastAssignments = Map; +type SingleUseIdentifiers = Set; +enum ReferenceKind { + Read, + Write, +} +function findReferencedRangeOfTemporaries(fn: HIRFunction): References { + const singleUseIdentifiers = 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 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 = new Map(); @@ -102,6 +172,7 @@ 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, @@ -109,6 +180,7 @@ function reorderBlock( ({ instruction: instr, dependencies: new Set(), + reorderability, depth: null, }) as Node ); @@ -116,7 +188,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 (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} `); + DEBUG && console.log(`${"| ".repeat(depth)}$${id} `); 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; } 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 23a0a839ea..2c9e67646b 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts +++ b/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": { 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 index dc49af4401..44eaa5b8b7 100644 --- 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 @@ -34,59 +34,47 @@ import { useState } from "react"; import { Stringify } from "shared-runtime"; function Component() { - const $ = _c(10); + const $ = _c(7); 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 = ( - ); + t1 = {state}; + $[0] = state; + $[1] = t0; + $[2] = t1; + } else { + t0 = $[1]; + t1 = $[2]; + } + let t2; + if ($[3] === Symbol.for("react.memo_cache_sentinel")) { + t2 = ; + $[3] = t2; + } else { + t2 = $[3]; + } + let t3; + if ($[4] !== t1 || $[5] !== t0) { + t3 = ( +
+ {t2} + {t1} + {t0} +
+ ); + $[4] = t1; $[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; + return t3; } export const FIXTURE_ENTRYPOINT = { 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 a2e79ad67f..49558eb0e4 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 @@ -27,7 +27,7 @@ import { c as _c } from "react/compiler-runtime"; // @enableInstructionReorderin import { useState } from "react"; function Component() { - const $ = _c(6); + const $ = _c(4); const [state, setState] = useState(0); let t0; if ($[0] === Symbol.for("react.memo_cache_sentinel")) { @@ -40,34 +40,26 @@ function Component() { } const onClick = t0; let t1; - if ($[1] !== state) { - t1 = Count: {state}; - $[1] = state; - $[2] = t1; + if ($[1] === Symbol.for("react.memo_cache_sentinel")) { + t1 = ; + $[1] = t1; } else { - t1 = $[2]; + t1 = $[1]; } let t2; - if ($[3] === Symbol.for("react.memo_cache_sentinel")) { - t2 = ; + if ($[2] !== state) { + t2 = ( + <> + Count: {state} + {t1} + + ); + $[2] = state; $[3] = t2; } else { t2 = $[3]; } - let t3; - if ($[4] !== t1) { - t3 = ( - <> - {t1} - {t2} - - ); - $[4] = t1; - $[5] = t3; - } else { - t3 = $[5]; - } - return t3; + return t2; } ```