From d2895356e71e1e20c5f231677049c4b1ba2b6e77 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Wed, 12 Jun 2024 16:36:27 -0700 Subject: [PATCH] [compiler] Allow reordering of LoadLocal after their last assignment [ghstack-poisoned] --- .../src/Optimization/InstructionReordering.ts | 104 +++++++++++++++++- 1 file changed, 99 insertions(+), 5 deletions(-) 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 37619b4224787..9a12a0e077bb9 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts @@ -13,14 +13,20 @@ import { HIRFunction, IdentifierId, Instruction, + InstructionId, + MutableRange, + Place, isExpressionBlockKind, + makeInstructionId, markInstructionIds, } from "../HIR"; import { printInstruction } from "../HIR/PrintHIR"; import { + eachInstructionLValue, eachInstructionValueLValue, eachInstructionValueOperand, eachTerminalOperand, + terminalFallthrough, } from "../HIR/visitors"; import { mayAllocate } from "../ReactiveScopes/InferReactiveScopeVariables"; import { getOrInsertWith } from "../Utils/utils"; @@ -69,8 +75,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 +98,76 @@ type Node = { depth: number | null; }; +// Inclusive start and end +type References = { + accessedRanges: AccessedRanges; + lastAssignments: LastAssignments; +}; +type LastAssignments = Map; +type AccessedRanges = Map; +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 = new Map(); @@ -116,7 +189,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 +293,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 +408,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 +423,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; }