[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
Joe Savona committed
Jun 12, 2024 at 15:48 UTC
59e73d90163b16ce820316b1c69710f23096857e
10 files changed
+543
-2
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
+6
@@ -41,6 +41,7 @@ import {
41
deadCodeElimination,
42
pruneMaybeThrows,
43
} from "../Optimization";
44
+import { instructionReordering } from "../Optimization/InstructionReordering";
45
import {
46
CodegenFunction,
47
alignObjectMethodScopes,
@@ -204,6 +205,11 @@ function* runWithEnvironment(
205
deadCodeElimination(hir);
206
yield log({ kind: "hir", name: "DeadCodeElimination", value: hir });
207
208
+ if (env.config.enableInstructionReordering) {
209
+ instructionReordering(hir);
210
+ yield log({ kind: "hir", name: "InstructionReordering", value: hir });
211
+ }
212
+
213
pruneMaybeThrows(hir);
214
yield log({ kind: "hir", name: "PruneMaybeThrows", value: hir });
215
compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
+6
@@ -277,6 +277,12 @@ const EnvironmentConfigSchema = z.object({
277
278
enableEmitHookGuards: ExternalFunctionSchema.nullish(),
279
280
+ /**
281
+ * Enable instruction reordering. See InstructionReordering.ts for the details
282
+ * of the approach.
283
+ */
284
+ enableInstructionReordering: z.boolean().default(false),
285
+
286
/*
287
* Enables instrumentation codegen. This emits a dev-mode only call to an
288
* instrumentation function, for components and hooks that Forget compiles.
compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
+22
@@ -335,6 +335,28 @@ export type HIR = {
335
* statements and not implicit exceptions which may occur.
336
*/
337
export type BlockKind = "block" | "value" | "loop" | "sequence" | "catch";
338
+
339
+/**
340
+ * Returns true for "block" and "catch" block kinds which correspond to statements
341
+ * in the source, including BlockStatement, CatchStatement.
342
+ *
343
+ * Inverse of isExpressionBlockKind()
344
+ */
345
+export function isStatementBlockKind(kind: BlockKind): boolean {
346
+ return kind === "block" || kind === "catch";
347
+}
348
+
349
+/**
350
+ * Returns true for "value", "loop", and "sequence" block kinds which correspond to
351
+ * expressions in the source, such as ConditionalExpression, LogicalExpression, loop
352
+ * initializer/test/updaters, etc
353
+ *
354
+ * Inverse of isStatementBlockKind()
355
+ */
356
+export function isExpressionBlockKind(kind: BlockKind): boolean {
357
+ return !isStatementBlockKind(kind);
358
+}
359
+
360
export type BasicBlock = {
361
kind: BlockKind;
362
id: BlockId;
compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts
new
+353
@@ -0,0 +1,353 @@
1
+/**
2
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
3
+ *
4
+ * This source code is licensed under the MIT license found in the
5
+ * LICENSE file in the root directory of this source tree.
6
+ */
7
+
8
+import { CompilerError } from "..";
9
+import {
10
+ BasicBlock,
11
+ Environment,
12
+ GeneratedSource,
13
+ HIRFunction,
14
+ IdentifierId,
15
+ Instruction,
16
+ isExpressionBlockKind,
17
+ markInstructionIds,
18
+} from "../HIR";
19
+import { printInstruction } from "../HIR/PrintHIR";
20
+import {
21
+ eachInstructionValueLValue,
22
+ eachInstructionValueOperand,
23
+ eachTerminalOperand,
24
+} from "../HIR/visitors";
25
+import { mayAllocate } from "../ReactiveScopes/InferReactiveScopeVariables";
26
+import { getOrInsertWith } from "../Utils/utils";
27
+
28
+/**
29
+ * This pass implements conservative instruction reordering to move instructions closer to
30
+ * to where their produced values are consumed. The goal is to group instructions in a way that
31
+ * is more optimal for future optimizations. Notably, MergeReactiveScopesThatAlwaysInvalidateTogether
32
+ * can only merge two candidate scopes if there are no intervenining instructions that are used by
33
+ * some later code: instruction reordering can move those intervening instructions later in many cases,
34
+ * thereby allowing more scopes to merge together.
35
+ *
36
+ * The high-level approach is to build a dependency graph where nodes correspond either to
37
+ * instructions OR to a particular lvalue assignment of another instruction. So
38
+ * `Destructure [x, y] = z` creates 3 nodes: one for the instruction, and one each for x and y.
39
+ * The lvalue nodes depend on the instruction node that assigns them.
40
+ *
41
+ * Dependency edges are added for all the lvalues and rvalues of each instruction, so for example
42
+ * the node for `t$2 = CallExpression t$0 ( t$1 )` will take dependencies on the nodes for t$0 and t$1.
43
+ *
44
+ * Individual instructions are grouped into two categories:
45
+ * - "Reorderable" instructions include a safe set of instructions that we know are fine to reorder.
46
+ * This includes JSX elements/fragments/text, primitives, template literals, and globals.
47
+ * These instructions are never emitted until they are referenced, and can even be moved across
48
+ * basic blocks until they are used.
49
+ * - All other instructions are non-reorderable, and take an explicit dependency on the last such
50
+ * non-reorderable instruction in their block. This largely ensures that mutations are serialized,
51
+ * since all potentially mutating instructions are in this category.
52
+ *
53
+ * The only remaining mutation not handled by the above is variable reassignment. To ensure that all
54
+ * reads/writes of a variable access the correct version, all references (lvalues and rvalues) to
55
+ * each named variable are serialized. Thus `x = 1; y = x; x = 2; z = x` will establish a chain
56
+ * of dependencies and retain the correct ordering.
57
+ *
58
+ * The algorithm proceeds one basic block at a time, first building up the dependnecy graph and then
59
+ * reordering.
60
+ *
61
+ * The reordering weights nodes according to their transitive dependencies, and whether a particular node
62
+ * needs memoization or not. Larger dependencies go first, followed by smaller dependencies, which in
63
+ * testing seems to allow scopes to merge more effectively. Over time we can likely continue to improve
64
+ * the reordering heuristic.
65
+ *
66
+ * An obvious area for improvement is to allow reordering of LoadLocals that occur after the last write
67
+ * of the named variable. We can add this in a follow-up.
68
+ */
69
+export function instructionReordering(fn: HIRFunction): void {
70
+ // Shared nodes are emitted when they are first used
71
+ const shared: Nodes = new Map();
72
+ for (const [, block] of fn.body.blocks) {
73
+ reorderBlock(fn.env, block, shared);
74
+ }
75
+ CompilerError.invariant(shared.size === 0, {
76
+ reason: `InstructionReordering: expected all reorderable nodes to have been emitted`,
77
+ loc:
78
+ [...shared.values()]
79
+ .map((node) => node.instruction?.loc)
80
+ .filter((loc) => loc != null)[0] ?? GeneratedSource,
81
+ });
82
+ markInstructionIds(fn.body);
83
+}
84
+
85
+const DEBUG = false;
86
+
87
+type Nodes = Map<IdentifierId, Node>;
88
+type Node = {
89
+ instruction: Instruction | null;
90
+ dependencies: Set<IdentifierId>;
91
+ depth: number | null;
92
+};
93
+
94
+function reorderBlock(
95
+ env: Environment,
96
+ block: BasicBlock,
97
+ shared: Nodes
98
+): void {
99
+ const locals: Nodes = new Map();
100
+ const named: Map<string, IdentifierId> = new Map();
101
+ let previous: IdentifierId | null = null;
102
+ for (const instr of block.instructions) {
103
+ const { lvalue, value } = instr;
104
+ // Get or create a node for this lvalue
105
+ const node = getOrInsertWith(
106
+ locals,
107
+ lvalue.identifier.id,
108
+ () =>
109
+ ({
110
+ instruction: instr,
111
+ dependencies: new Set(),
112
+ depth: null,
113
+ }) as Node
114
+ );
115
+ /**
116
+ * Ensure non-reoderable instructions have their order retained by
117
+ * adding explicit dependencies to the previous such instruction.
118
+ */
119
+ if (getReoderability(instr) === Reorderability.Nonreorderable) {
120
+ if (previous !== null) {
121
+ node.dependencies.add(previous);
122
+ }
123
+ previous = lvalue.identifier.id;
124
+ }
125
+ /**
126
+ * Establish dependencies on operands
127
+ */
128
+ for (const operand of eachInstructionValueOperand(value)) {
129
+ const { name, id } = operand.identifier;
130
+ if (name !== null && name.kind === "named") {
131
+ // Serialize all accesses to named variables
132
+ const previous = named.get(name.value);
133
+ if (previous !== undefined) {
134
+ node.dependencies.add(previous);
135
+ }
136
+ named.set(name.value, lvalue.identifier.id);
137
+ } else if (locals.has(id) || shared.has(id)) {
138
+ node.dependencies.add(id);
139
+ }
140
+ }
141
+ /**
142
+ * Establish nodes for lvalues, with dependencies on the node
143
+ * for the instruction itself. This ensures that any consumers
144
+ * of the lvalue will take a dependency through to the original
145
+ * instruction.
146
+ */
147
+ for (const lvalueOperand of eachInstructionValueLValue(value)) {
148
+ const lvalueNode = getOrInsertWith(
149
+ locals,
150
+ lvalueOperand.identifier.id,
151
+ () =>
152
+ ({
153
+ instruction: null,
154
+ dependencies: new Set(),
155
+ depth: null,
156
+ }) as Node
157
+ );
158
+ lvalueNode.dependencies.add(lvalue.identifier.id);
159
+ const name = lvalueOperand.identifier.name;
160
+ if (name !== null && name.kind === "named") {
161
+ const previous = named.get(name.value);
162
+ if (previous !== undefined) {
163
+ node.dependencies.add(previous);
164
+ }
165
+ named.set(name.value, lvalue.identifier.id);
166
+ }
167
+ }
168
+ }
169
+
170
+ const nextInstructions: Array<Instruction> = [];
171
+ const seen = new Set<IdentifierId>();
172
+
173
+ DEBUG && console.log(`bb${block.id}`);
174
+
175
+ // First emit everything that can't be reordered
176
+ if (previous !== null) {
177
+ DEBUG && console.log(`(last non-reorderable instruction)`);
178
+ DEBUG && print(env, locals, shared, seen, previous);
179
+ emit(env, locals, shared, nextInstructions, previous);
180
+ }
181
+ /*
182
+ * For "value" blocks the final instruction represents its value, so we have to be
183
+ * careful to not change the ordering. Emit the last instruction explicitly.
184
+ * Any non-reorderable instructions will get emitted first, and any unused
185
+ * reorderable instructions can be deferred to the shared node list.
186
+ */
187
+ if (isExpressionBlockKind(block.kind) && block.instructions.length !== 0) {
188
+ DEBUG && console.log(`(block value)`);
189
+ DEBUG &&
190
+ print(
191
+ env,
192
+ locals,
193
+ shared,
194
+ seen,
195
+ block.instructions.at(-1)!.lvalue.identifier.id
196
+ );
197
+ emit(
198
+ env,
199
+ locals,
200
+ shared,
201
+ nextInstructions,
202
+ block.instructions.at(-1)!.lvalue.identifier.id
203
+ );
204
+ }
205
+ /*
206
+ * Then emit the dependencies of the terminal operand. In many cases they will have
207
+ * already been emitted in the previous step and this is a no-op.
208
+ * TODO: sort the dependencies based on weight, like we do for other nodes. Not a big
209
+ * deal though since most terminals have a single operand
210
+ */
211
+ for (const operand of eachTerminalOperand(block.terminal)) {
212
+ DEBUG && console.log(`(terminal operand)`);
213
+ DEBUG && print(env, locals, shared, seen, operand.identifier.id);
214
+ emit(env, locals, shared, nextInstructions, operand.identifier.id);
215
+ }
216
+ // Anything not emitted yet is globally reorderable
217
+ for (const [id, node] of locals) {
218
+ if (node.instruction == null) {
219
+ continue;
220
+ }
221
+ CompilerError.invariant(
222
+ node.instruction != null &&
223
+ getReoderability(node.instruction) === Reorderability.Reorderable,
224
+ {
225
+ reason: `Expected all remaining instructions to be reorderable`,
226
+ loc: node.instruction?.loc ?? block.terminal.loc,
227
+ description:
228
+ node.instruction != null
229
+ ? `Instruction [${node.instruction.id}] was not emitted yet but is not reorderable`
230
+ : `Lvalue $${id} was not emitted yet but is not reorderable`,
231
+ }
232
+ );
233
+ DEBUG && console.log(`save shared: $${id}`);
234
+ shared.set(id, node);
235
+ }
236
+
237
+ block.instructions = nextInstructions;
238
+ DEBUG && console.log();
239
+}
240
+
241
+function getDepth(env: Environment, nodes: Nodes, id: IdentifierId): number {
242
+ const node = nodes.get(id)!;
243
+ if (node == null) {
244
+ return 0;
245
+ }
246
+ if (node.depth != null) {
247
+ return node.depth;
248
+ }
249
+ node.depth = 0; // in case of cycles
250
+ let depth =
251
+ node.instruction != null && mayAllocate(env, node.instruction) ? 1 : 0;
252
+ for (const dep of node.dependencies) {
253
+ depth += getDepth(env, nodes, dep);
254
+ }
255
+ node.depth = depth;
256
+ return depth;
257
+}
258
+
259
+function print(
260
+ env: Environment,
261
+ locals: Nodes,
262
+ shared: Nodes,
263
+ seen: Set<IdentifierId>,
264
+ id: IdentifierId,
265
+ depth: number = 0
266
+): void {
267
+ if (seen.has(id)) {
268
+ console.log(`${"| ".repeat(depth)}$${id} <skipped>`);
269
+ return;
270
+ }
271
+ seen.add(id);
272
+ const node = locals.get(id) ?? shared.get(id);
273
+ if (node == null) {
274
+ return;
275
+ }
276
+ const deps = [...node.dependencies];
277
+ deps.sort((a, b) => {
278
+ const aDepth = getDepth(env, locals, a);
279
+ const bDepth = getDepth(env, locals, b);
280
+ return bDepth - aDepth;
281
+ });
282
+ for (const dep of deps) {
283
+ print(env, locals, shared, seen, dep, depth + 1);
284
+ }
285
+ console.log(
286
+ `${"| ".repeat(depth)}$${id} ${printNode(node)} deps=[${deps.map((x) => `$${x}`).join(", ")}]`
287
+ );
288
+}
289
+
290
+function printNode(node: Node): string {
291
+ const { instruction } = node;
292
+ if (instruction === null) {
293
+ return "<lvalue-only>";
294
+ }
295
+ switch (instruction.value.kind) {
296
+ case "FunctionExpression":
297
+ case "ObjectMethod": {
298
+ return `[${instruction.id}] ${instruction.value.kind}`;
299
+ }
300
+ default: {
301
+ return printInstruction(instruction);
302
+ }
303
+ }
304
+}
305
+
306
+function emit(
307
+ env: Environment,
308
+ locals: Nodes,
309
+ shared: Nodes,
310
+ instructions: Array<Instruction>,
311
+ id: IdentifierId
312
+): void {
313
+ const node = locals.get(id) ?? shared.get(id);
314
+ if (node == null) {
315
+ return;
316
+ }
317
+ locals.delete(id);
318
+ shared.delete(id);
319
+ const deps = [...node.dependencies];
320
+ deps.sort((a, b) => {
321
+ const aDepth = getDepth(env, locals, a);
322
+ const bDepth = getDepth(env, locals, b);
323
+ return bDepth - aDepth;
324
+ });
325
+ for (const dep of deps) {
326
+ emit(env, locals, shared, instructions, dep);
327
+ }
328
+ if (node.instruction !== null) {
329
+ instructions.push(node.instruction);
330
+ }
331
+}
332
+
333
+enum Reorderability {
334
+ Reorderable,
335
+ Nonreorderable,
336
+}
337
+function getReoderability(instr: Instruction): Reorderability {
338
+ switch (instr.value.kind) {
339
+ case "JsxExpression":
340
+ case "JsxFragment":
341
+ case "JSXText":
342
+ case "LoadGlobal":
343
+ case "Primitive":
344
+ case "TemplateLiteral":
345
+ case "BinaryExpression":
346
+ case "UnaryExpression": {
347
+ return Reorderability.Reorderable;
348
+ }
349
+ default: {
350
+ return Reorderability.Nonreorderable;
351
+ }
352
+ }
353
+}
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts
+4
-1
@@ -186,7 +186,10 @@ export function isMutable({ id }: Instruction, place: Place): boolean {
186
return id >= range.start && id < range.end;
187
}
188
189
-function mayAllocate(env: Environment, instruction: Instruction): boolean {
189
+export function mayAllocate(
190
+ env: Environment,
191
+ instruction: Instruction
192
+): boolean {
193
const { value } = instruction;
194
switch (value.kind) {
195
case "Destructure": {
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MemoizeFbtAndMacroOperandsInSameScope.ts
+28
@@ -114,6 +114,7 @@ function visit(
114
operand.identifier.mutableRange.start
115
)
116
);
117
+ fbtValues.add(operand.identifier.id);
118
}
119
} else if (
120
isFbtJsxExpression(fbtMacroTags, fbtValues, value) ||
@@ -146,6 +147,33 @@ function visit(
147
*/
148
fbtValues.add(operand.identifier.id);
149
}
150
+ } else if (fbtValues.has(lvalue.identifier.id)) {
151
+ const fbtScope = lvalue.identifier.scope;
152
+ if (fbtScope === null) {
153
+ return;
154
+ }
155
+
156
+ for (const operand of eachReactiveValueOperand(value)) {
157
+ if (
158
+ operand.identifier.name !== null &&
159
+ operand.identifier.name.kind === "named"
160
+ ) {
161
+ /*
162
+ * named identifiers were already locals, we only have to force temporaries
163
+ * into the same scope
164
+ */
165
+ continue;
166
+ }
167
+ operand.identifier.scope = fbtScope;
168
+
169
+ // Expand the jsx element's range to account for its operands
170
+ fbtScope.range.start = makeInstructionId(
171
+ Math.min(
172
+ fbtScope.range.start,
173
+ operand.identifier.mutableRange.start
174
+ )
175
+ );
176
+ }
177
}
178
}
179
}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-reordering.expect.md
new
+100
@@ -0,0 +1,100 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @enableInstructionReordering
6
+import { useState } from "react";
7
+import { Stringify } from "shared-runtime";
8
+
9
+function Component() {
10
+ let [state, setState] = useState(0);
11
+ return (
12
+ <div>
13
+ <Stringify text="Counter" />
14
+ <span>{state}</span>
15
+ <button data-testid="button" onClick={() => setState(state + 1)}>
16
+ increment
17
+ </button>
18
+ </div>
19
+ );
20
+}
21
+
22
+export const FIXTURE_ENTRYPOINT = {
23
+ fn: Component,
24
+ params: [{ value: 42 }],
25
+};
26
+
27
+```
28
+
29
+## Code
30
+
31
+```javascript
32
+import { c as _c } from "react/compiler-runtime"; // @enableInstructionReordering
33
+import { useState } from "react";
34
+import { Stringify } from "shared-runtime";
35
+
36
+function Component() {
37
+ const $ = _c(10);
38
+ const [state, setState] = useState(0);
39
+ let t0;
40
+ if ($[0] !== state) {
41
+ t0 = () => setState(state + 1);
42
+ $[0] = state;
43
+ $[1] = t0;
44
+ } else {
45
+ t0 = $[1];
46
+ }
47
+ let t1;
48
+ if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
49
+ t1 = <Stringify text="Counter" />;
50
+ $[2] = t1;
51
+ } else {
52
+ t1 = $[2];
53
+ }
54
+ let t2;
55
+ if ($[3] !== state) {
56
+ t2 = <span>{state}</span>;
57
+ $[3] = state;
58
+ $[4] = t2;
59
+ } else {
60
+ t2 = $[4];
61
+ }
62
+ let t3;
63
+ if ($[5] !== t0) {
64
+ t3 = (
65
+ <button data-testid="button" onClick={t0}>
66
+ increment
67
+ </button>
68
+ );
69
+ $[5] = t0;
70
+ $[6] = t3;
71
+ } else {
72
+ t3 = $[6];
73
+ }
74
+ let t4;
75
+ if ($[7] !== t2 || $[8] !== t3) {
76
+ t4 = (
77
+ <div>
78
+ {t1}
79
+ {t2}
80
+ {t3}
81
+ </div>
82
+ );
83
+ $[7] = t2;
84
+ $[8] = t3;
85
+ $[9] = t4;
86
+ } else {
87
+ t4 = $[9];
88
+ }
89
+ return t4;
90
+}
91
+
92
+export const FIXTURE_ENTRYPOINT = {
93
+ fn: Component,
94
+ params: [{ value: 42 }],
95
+};
96
+
97
+```
98
+
99
+### Eval output
100
+(kind: ok) <div><div>{"text":"Counter"}</div><span>0</span><button data-testid="button">increment</button></div>
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-reordering.js
new
+21
@@ -0,0 +1,21 @@
1
+// @enableInstructionReordering
2
+import { useState } from "react";
3
+import { Stringify } from "shared-runtime";
4
+
5
+function Component() {
6
+ let [state, setState] = useState(0);
7
+ return (
8
+ <div>
9
+ <Stringify text="Counter" />
10
+ <span>{state}</span>
11
+ <button data-testid="button" onClick={() => setState(state + 1)}>
12
+ increment
13
+ </button>
14
+ </div>
15
+ );
16
+}
17
+
18
+export const FIXTURE_ENTRYPOINT = {
19
+ fn: Component,
20
+ params: [{ value: 42 }],
21
+};
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.expect.md
+2
-1
@@ -2,6 +2,7 @@
2
## Input
3
4
```javascript
5
+// @enableInstructionReordering
6
import { useState } from "react";
7
8
function Component() {
@@ -22,7 +23,7 @@ function Component() {
23
## Code
24
25
```javascript
25
-import { c as _c } from "react/compiler-runtime";
26
+import { c as _c } from "react/compiler-runtime"; // @enableInstructionReordering
27
import { useState } from "react";
28
29
function Component() {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.js
+1
@@ -1,3 +1,4 @@
1
+// @enableInstructionReordering
2
import { useState } from "react";
3
4
function Component() {