From de873ab1da41099524b16c0e07ddaaf69d0766df Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Fri, 3 Jan 2025 16:47:36 -0800 Subject: [PATCH] [compiler] Early sketch of ReactiveIR Early sketch of a new representation for the middle-phase of the compiler. This is a sea-of-nodes approach, representing instructions and terminals as nodes with direct and control dependencies. Something like this: ```js let array = []; if (cond) { array.push(value); } return array; ``` Would correspond roughly to a graph as follows. Note that the actual representation uses `Instruction` as-is, and maps dependencies into local `Place`s, but for ease of reading i've mapped operands as the node they come from (n0 etc): ``` n0 = ArrayExpression [] n1 = StoreLocal 'array' = n0 // nodes from the consequent n2 = LoadLocal 'array' n3 = PropertyLoad n3 . 'push' n4 = LoadLocal 'value' n5 = MethodCall n2 . n3 ( n4 ) // if terminal n6 = LoadLocal 'cond' n7 = If test=n6 consequent=n5 alternate=(not pictured) // return terminal n8 = LoadLocal 'array'; n9 = Return value=n8 exit=n9 ``` Note that even if there are multiple returns, there will always be a single "exit" node, which is just the node corresponding to the last statement of the outer block in the original program. What's implemented so far is very basic translation of HIR to ReactiveGraph. The only dependencies created are direct data dependencies and basic control flow, which means that e.g. code inside an if consequent will stay "inside" that consequent, but it's possible for mutations or reassignments to get reordered. The reordering happens because we reverse-postorder the graph after construction, which automatically orders strictly based on dependencies and moves things around otherwise. There is lots left to do here, including cleaning up the node dependencies and how they map to local Places, establishing other forms of dependencies, etc. ghstack-source-id: fc633994d81aa93389a2dbbbe83358eea7eeccb0 Pull Request resolved: https://github.com/facebook/react/pull/31974 --- .../src/Entrypoint/Pipeline.ts | 11 + .../src/HIR/Environment.ts | 6 + .../src/ReactiveIR/BuildReactiveGraph.ts | 228 ++++++++++++ .../src/ReactiveIR/ReactiveIR.ts | 352 ++++++++++++++++++ .../fixtures/compiler/reactive-graph.js | 8 + 5 files changed, 605 insertions(+) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/BuildReactiveGraph.ts create mode 100644 compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/ReactiveIR.ts create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reactive-graph.js 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 ca6abc0748eb9..f5052cb1bcebc 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts @@ -99,6 +99,8 @@ import {propagateScopeDependenciesHIR} from '../HIR/PropagateScopeDependenciesHI import {outlineJSX} from '../Optimization/OutlineJsx'; import {optimizePropsMethodCalls} from '../Optimization/OptimizePropsMethodCalls'; import {transformFire} from '../Transform'; +import {buildReactiveGraph} from '../ReactiveIR/BuildReactiveGraph'; +import {printReactiveGraph} from '../ReactiveIR/ReactiveIR'; export type CompilerPipelineValue = | {kind: 'ast'; name: string; value: CodegenFunction} @@ -373,6 +375,15 @@ function runWithEnvironment( }); } + if (env.config.enableReactiveGraph) { + const reactiveGraph = buildReactiveGraph(hir); + log({ + kind: 'debug', + name: 'BuildReactiveGraph', + value: printReactiveGraph(reactiveGraph), + }); + } + const reactiveFunction = buildReactiveFunction(hir); log({ kind: 'reactive', 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 f3f426df56e44..f6543e9337371 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -395,6 +395,12 @@ const EnvironmentConfigSchema = z.object({ */ enableInstructionReordering: z.boolean().default(false), + /** + * Enables ReactiveGraph-based optimizations including reordering across terminal + * boundaries + */ + enableReactiveGraph: z.boolean().default(false), + /** * Enables function outlinining, where anonymous functions that do not close over * local variables can be extracted into top-level helper functions. diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/BuildReactiveGraph.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/BuildReactiveGraph.ts new file mode 100644 index 0000000000000..c9e5e98b0e1ef --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/BuildReactiveGraph.ts @@ -0,0 +1,228 @@ +/** + * 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, SourceLocation} from '..'; +import {BlockId, HIRFunction, Identifier, IdentifierId, Place} from '../HIR'; +import {printIdentifier, printPlace} from '../HIR/PrintHIR'; +import { + eachInstructionLValue, + eachInstructionValueOperand, + terminalFallthrough, +} from '../HIR/visitors'; +import { + IfNode, + InstructionNode, + LoadArgumentNode, + makeReactiveId, + NodeDependencies, + NodeReference, + populateReactiveGraphNodeOutputs, + printReactiveNodes, + ReactiveGraph, + ReactiveId, + ReactiveNode, + ReturnNode, + reversePostorderReactiveGraph, + ScopeNode, +} from './ReactiveIR'; + +export function buildReactiveGraph(fn: HIRFunction): ReactiveGraph { + const builder = new Builder(); + for (const param of fn.params) { + const place = param.kind === 'Identifier' ? param : param.place; + const node: LoadArgumentNode = { + kind: 'LoadArgument', + id: builder.nextReactiveId, + loc: place.loc, + outputs: [], + place: {...place}, + }; + builder.nodes.set(node.id, node); + builder.declare(node.id, place); + } + + const exitNode = buildBlockScope(fn, builder, fn.body.entry); + + const graph: ReactiveGraph = { + async: fn.async, + directives: fn.directives, + env: fn.env, + exit: exitNode, + fnType: fn.fnType, + generator: fn.generator, + id: fn.id, + loc: fn.loc, + nextNodeId: builder._nextNodeId, + nodes: builder.nodes, + params: fn.params, + }; + populateReactiveGraphNodeOutputs(graph); + reversePostorderReactiveGraph(graph); + return graph; +} + +class Builder { + _nextNodeId: number = 0; + #environment: Map = new Map(); + nodes: Map = new Map(); + args: Set = new Set(); + + get nextReactiveId(): ReactiveId { + return makeReactiveId(this._nextNodeId++); + } + + declare(node: ReactiveId, place: Place): void { + this.#environment.set(place.identifier.id, {node, from: place}); + } + + lookup( + identifier: Identifier, + loc: SourceLocation, + ): {node: ReactiveId; from: Place} { + const dep = this.#environment.get(identifier.id); + if (dep == null) { + console.log(printReactiveNodes(this.nodes)); + for (const [id, dep] of this.#environment) { + console.log(`t#${id} => £${dep.node} . ${printPlace(dep.from)}`); + } + + console.log(); + console.log(`could not find ${printIdentifier(identifier)}`); + } + CompilerError.invariant(dep != null, { + reason: `No source node for identifier ${printIdentifier(identifier)}`, + loc, + }); + return dep; + } +} + +function buildBlockScope( + fn: HIRFunction, + builder: Builder, + entry: BlockId, +): ReactiveId { + let block = fn.body.blocks.get(entry)!; + let lastNode: ReactiveNode = { + kind: 'Empty', + id: builder.nextReactiveId, + loc: block.terminal.loc, + outputs: [], + }; + builder.nodes.set(lastNode.id, lastNode); + while (true) { + // iterate instructions of the block + for (const instr of block.instructions) { + const dependencies: NodeDependencies = new Map(); + for (const operand of eachInstructionValueOperand(instr.value)) { + const dep = builder.lookup(operand.identifier, operand.loc); + dependencies.set(dep.node, { + from: {...dep.from}, + as: {...operand}, + }); + } + const node: InstructionNode = { + kind: 'Value', + controlDependency: null, + dependencies, + id: builder.nextReactiveId, + loc: instr.loc, + outputs: [], + value: instr, + }; + builder.nodes.set(node.id, node); + lastNode = node; + for (const lvalue of eachInstructionLValue(instr)) { + builder.declare(node.id, lvalue); + } + } + + // handle the terminal + const terminal = block.terminal; + switch (terminal.kind) { + case 'if': { + const testDep = builder.lookup( + terminal.test.identifier, + terminal.test.loc, + ); + const test: NodeReference = { + node: testDep.node, + from: testDep.from, + as: {...terminal.test}, + }; + const consequent = buildBlockScope(fn, builder, terminal.consequent); + const alternate = buildBlockScope(fn, builder, terminal.alternate); + const ifNode: IfNode = { + kind: 'If', + alternate, + consequent, + id: builder.nextReactiveId, + loc: terminal.loc, + outputs: [], + test, + }; + builder.nodes.set(ifNode.id, ifNode); + lastNode = ifNode; + break; + } + case 'return': { + const valueDep = builder.lookup( + terminal.value.identifier, + terminal.value.loc, + ); + const value: NodeReference = { + node: valueDep.node, + from: valueDep.from, + as: {...terminal.value}, + }; + const returnNode: ReturnNode = { + kind: 'Return', + id: builder.nextReactiveId, + loc: terminal.loc, + outputs: [], + value, + }; + builder.nodes.set(returnNode.id, returnNode); + lastNode = returnNode; + break; + } + case 'scope': { + const body = buildBlockScope(fn, builder, terminal.block); + const scopeNode: ScopeNode = { + kind: 'Scope', + body, + dependencies: new Map(), + id: builder.nextReactiveId, + loc: terminal.scope.loc, + outputs: [], + scope: terminal.scope, + }; + builder.nodes.set(scopeNode.id, scopeNode); + lastNode = scopeNode; + break; + } + case 'goto': { + break; + } + default: { + CompilerError.throwTodo({ + reason: `Support ${terminal.kind} nodes`, + loc: terminal.loc, + }); + } + } + + // Continue iteration in the fallthrough + const fallthrough = terminalFallthrough(terminal); + if (fallthrough != null) { + block = fn.body.blocks.get(fallthrough)!; + } else { + break; + } + } + return lastNode.id; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/ReactiveIR.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/ReactiveIR.ts new file mode 100644 index 0000000000000..ae7b14a3455b5 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/ReactiveIR.ts @@ -0,0 +1,352 @@ +/** + * 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 { + Environment, + Instruction, + Place, + ReactiveScope, + SourceLocation, + SpreadPattern, +} from '../HIR'; +import {ReactFunctionType} from '../HIR/Environment'; +import {printInstruction, printPlace} from '../HIR/PrintHIR'; +import {assertExhaustive} from '../Utils/utils'; + +export type ReactiveGraph = { + nodes: Map; + nextNodeId: number; + exit: ReactiveId; + loc: SourceLocation; + id: string | null; + params: Array; + generator: boolean; + async: boolean; + env: Environment; + directives: Array; + fnType: ReactFunctionType; +}; + +/* + * Simulated opaque type for Reactive IDs to prevent using normal numbers as ids + * accidentally. + */ +const opaqueReactiveId = Symbol(); +export type ReactiveId = number & {[opaqueReactiveId]: 'ReactiveId'}; + +export function makeReactiveId(id: number): ReactiveId { + CompilerError.invariant(id >= 0 && Number.isInteger(id), { + reason: 'Expected reactive node id to be a non-negative integer', + description: null, + loc: null, + suggestions: null, + }); + return id as ReactiveId; +} + +export type ReactiveNode = + | LoadArgumentNode + | InstructionNode + | IfNode + | EmptyNode + | ReturnNode + | ScopeNode; + +export type NodeReference = { + node: ReactiveId; + from: Place; + as: Place; +}; + +export type NodeDependencies = Map; +export type NodeDependency = {from: Place; as: Place}; + +export type LoadArgumentNode = { + kind: 'LoadArgument'; + id: ReactiveId; + loc: SourceLocation; + outputs: Array; + place: Place; +}; + +// An individual instruction +export type InstructionNode = { + kind: 'Value'; + id: ReactiveId; + loc: SourceLocation; + outputs: Array; + dependencies: NodeDependencies; + controlDependency: ReactiveId | null; + value: Instruction; +}; + +export type ReturnNode = { + kind: 'Return'; + id: ReactiveId; + loc: SourceLocation; + value: NodeReference; + outputs: Array; +}; + +export type IfNode = { + kind: 'If'; + id: ReactiveId; + loc: SourceLocation; + outputs: Array; + test: NodeReference; + consequent: ReactiveId; + alternate: ReactiveId; +}; + +export type EmptyNode = { + kind: 'Empty'; + id: ReactiveId; + loc: SourceLocation; + outputs: Array; +}; + +export type ScopeNode = { + kind: 'Scope'; + id: ReactiveId; + loc: SourceLocation; + outputs: Array; + scope: ReactiveScope; + /** + * The hoisted dependencies of the scope. Instructions "within" the scope + * (ie, the declarations or their deps) will also depend on these same values + * but we explicitly describe them here to ensure that all deps come before the scope + */ + dependencies: NodeDependencies; + /** + * The nodes that produce the values declared by the scope + */ + // declarations: NodeDependencies; + body: ReactiveId; +}; + +function _staticInvariantReactiveNodeHasIdLocationAndOutputs( + node: ReactiveNode, +): [ReactiveId, SourceLocation, Array] { + // If this fails, it is because a variant of ReactiveNode is missing a .id and/or .loc - add it! + return [node.id, node.loc, node.outputs]; +} + +/** + * Populates the outputs of each node in the graph + */ +export function populateReactiveGraphNodeOutputs(graph: ReactiveGraph): void { + // Populate node outputs + for (const [, node] of graph.nodes) { + node.outputs.length = 0; + } + for (const [, node] of graph.nodes) { + for (const dep of eachNodeDependency(node)) { + const sourceNode = graph.nodes.get(dep); + CompilerError.invariant(sourceNode != null, { + reason: `Expected source dependency ${dep} to exist`, + loc: node.loc, + }); + sourceNode.outputs.push(node.id); + } + } + const exitNode = graph.nodes.get(graph.exit)!; + exitNode.outputs.push(graph.exit); +} + +/** + * Puts the nodes of the graph into reverse postorder, such that nodes + * appear before any of their "successors" (consumers/dependents). + */ +export function reversePostorderReactiveGraph(graph: ReactiveGraph): void { + const nodes: Map = new Map(); + function visit(id: ReactiveId): void { + if (nodes.has(id)) { + return; + } + const node = graph.nodes.get(id); + CompilerError.invariant(node != null, { + reason: `Missing definition for ID ${id}`, + loc: null, + }); + for (const dep of eachNodeDependency(node)) { + visit(dep); + } + nodes.set(id, node); + } + for (const [_id, node] of graph.nodes) { + if (node.outputs.length === 0 && node.kind !== 'Empty') { + visit(node.id); + } + } + visit(graph.exit); + graph.nodes = nodes; +} + +export function* eachNodeDependency(node: ReactiveNode): Iterable { + switch (node.kind) { + case 'LoadArgument': + case 'Empty': { + break; + } + case 'If': { + yield node.test.node; + yield node.consequent; + yield node.alternate; + break; + } + case 'Return': { + yield node.value.node; + break; + } + case 'Value': { + yield* [...node.dependencies.keys()]; + if (node.controlDependency != null) { + yield node.controlDependency; + } + break; + } + case 'Scope': { + yield* [...node.dependencies.keys()]; + // yield* [...node.declarations.keys()]; + yield node.body; + break; + } + default: { + assertExhaustive(node, `Unexpected node kind '${(node as any).kind}'`); + } + } +} + +export function* eachNodeReference( + node: ReactiveNode, +): Iterable { + switch (node.kind) { + case 'LoadArgument': + case 'Empty': { + break; + } + case 'Return': { + yield node.value; + break; + } + case 'If': { + yield node.test; + break; + } + case 'Value': { + yield* [...node.dependencies].map(([node, dep]) => ({ + node, + from: dep.from, + as: dep.as, + })); + break; + } + case 'Scope': { + yield* [...node.dependencies].map(([node, dep]) => ({ + node, + from: dep.from, + as: dep.as, + })); + // yield* [...node.declarations].map(([node, dep]) => ({ + // node, + // from: dep.from, + // as: dep.as, + // })); + break; + } + default: { + assertExhaustive(node, `Unexpected node kind '${(node as any).kind}'`); + } + } +} + +function printNodeReference({node, from, as}: NodeReference): string { + return `£${node}.${printPlace(from)} => ${printPlace(as)}`; +} + +export function printNodeDependencies(deps: NodeDependencies): string { + const buffer: Array = []; + for (const [id, dep] of deps) { + buffer.push(printNodeReference({node: id, from: dep.from, as: dep.as})); + } + return buffer.join(', '); +} + +export function printReactiveGraph(graph: ReactiveGraph): string { + const buffer: Array = []; + buffer.push( + `${graph.fnType} ${graph.id ?? ''}(` + + graph.params + .map(param => { + if (param.kind === 'Identifier') { + return printPlace(param); + } else { + return `...${printPlace(param.place)}`; + } + }) + .join(', ') + + ')', + ); + writeReactiveNodes(buffer, graph.nodes); + buffer.push(`Exit £${graph.exit}`); + return buffer.join('\n'); +} + +export function printReactiveNodes( + nodes: Map, +): string { + const buffer: Array = []; + writeReactiveNodes(buffer, nodes); + return buffer.join('\n'); +} + +function writeReactiveNodes( + buffer: Array, + nodes: Map, +): void { + for (const [id, node] of nodes) { + const deps = [...eachNodeReference(node)] + .map(id => printNodeReference(id)) + .join(' '); + switch (node.kind) { + case 'LoadArgument': { + buffer.push(`£${id} LoadArgument ${printPlace(node.place)}`); + break; + } + case 'Empty': { + buffer.push(`£${id} Empty deps=[${deps}]`); + break; + } + case 'Return': { + buffer.push(`£${id} Return ${printNodeReference(node.value)}`); + break; + } + case 'If': { + buffer.push( + `£${id} If test=${printNodeReference(node.test)} consequent=£${node.consequent} alternate=£${node.alternate}`, + ); + break; + } + case 'Value': { + buffer.push(`£${id} Intermediate deps=[${deps}]`); + buffer.push(' ' + printInstruction(node.value)); + break; + } + case 'Scope': { + buffer.push( + // `£${id} Scope @${node.scope.id} deps=[${printNodeDependencies(node.dependencies)}] declarations=[${printNodeDependencies(node.declarations)}]`, + `£${id} Scope @${node.scope.id} deps=[${printNodeDependencies(node.dependencies)}] body=£${node.body}`, + ); + break; + } + default: { + assertExhaustive(node, `Unexpected node kind ${(node as any).kind}`); + } + } + } +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reactive-graph.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reactive-graph.js new file mode 100644 index 0000000000000..ac76010e16d75 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reactive-graph.js @@ -0,0 +1,8 @@ +// @enableReactiveGraph +function Component(props) { + const elements = []; + if (props.value) { + elements.push(
{props.value}
); + } + return elements; +}