Skip to content

Commit 51f54ec

Browse files
committed
compiler: super early exploration of instruction reordering
See comments in InstructionReordering.ts. This needs substantial iteration before landing in some form, just putting up to share for discussion. ghstack-source-id: 61ef96eb7e545317b47fb055b36a941be6d239ab Pull Request resolved: #29579
1 parent 5fe8c0b commit 51f54ec

File tree

364 files changed

+3005
-2631
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

364 files changed

+3005
-2631
lines changed

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

+4
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ import {
4141
deadCodeElimination,
4242
pruneMaybeThrows,
4343
} from "../Optimization";
44+
import { instructionReordering } from "../Optimization/InstructionReordering";
4445
import {
4546
CodegenFunction,
4647
alignObjectMethodScopes,
@@ -177,6 +178,9 @@ function* runWithEnvironment(
177178
inferTypes(hir);
178179
yield log({ kind: "hir", name: "InferTypes", value: hir });
179180

181+
instructionReordering(hir);
182+
yield log({ kind: "hir", name: "InstructionReordering", value: hir });
183+
180184
if (env.config.validateHooksUsage) {
181185
validateHooksUsage(hir);
182186
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
import {
2+
BasicBlock,
3+
HIRFunction,
4+
IdentifierId,
5+
Instruction,
6+
markInstructionIds,
7+
} from "../HIR";
8+
import { printInstruction } from "../HIR/PrintHIR";
9+
import {
10+
eachInstructionValueLValue,
11+
eachInstructionValueOperand,
12+
eachTerminalOperand,
13+
} from "../HIR/visitors";
14+
import { getOrInsertDefault } from "../Utils/utils";
15+
16+
/**
17+
* WIP early exploration of instruction reordering. This is a fairly aggressive form and has
18+
* some issues. The idea of what's implemented:
19+
*
20+
* The high-level approach is to build a dependency graph where nodes generally correspond
21+
* either to instructions OR to particular lvalue assignments of an expresssion. So
22+
* `Destructure [x, y] = z` creates 3 nodes: one for the instruction, and one each for x and y.
23+
* The lvalue nodes depend on the instruction node that assigns them.
24+
*
25+
* We add dependency edges for all the rvalues/lvalues of each instruction. In addition, we
26+
* implicitly add dependencies btw non-reorderable instructions (more on that criteria) to
27+
* serialize any instruction where order might be observable.
28+
*
29+
* We then distinguish two types of instructions that are reorderable:
30+
* - Primitives, JSXText, JSX elements, and globals can be *globally* reordered, ie across blocks.
31+
* We defer emitting them until they are first used globally.
32+
* - Array and object expressions are reorderable within basic blocks. This could likely be relaxed to be global.
33+
* - StoreLocal, LoadLocal, and Destructure are reorderable within basic blocks. However, we serialize all
34+
* references to each named variable (reads and writes) to ensure that we aren't changing the order of evaluation
35+
* of variable references.
36+
*
37+
* The variable reordering relies on the fact that any variables that could be reassigned via a function expression
38+
* are promoted to "context" variables and use LoadContext/StoreContext, which are not reorderable.
39+
*
40+
* In theory it might even be safe to do this variable reordering globally, but i want to think through that more.
41+
*
42+
* With the above context, the algorithm is approximately:
43+
* - For each basic block:
44+
* - Iterate the instructions to create the dependency graph
45+
* - Re-emit instructions, "pulling" from all the values that are depended upon by the block's terminal.
46+
* - Emit any remaining instructions that cannot be globally reordered, starting from later instructions first.
47+
* - Save any globally-reorderable instructions into a global map that is shared across blocks, so they can be
48+
* emitted by the first block that needs them.
49+
*
50+
* Emitting instructions is currently naive: we just iterate in the order that the dependencies were established.
51+
* If instruction 4 depends on instructions 1, 2, and 3, we'll visit in depth-first order and emit 1, 2, 3, 4.
52+
* That's true even if instruction 1 and 2 are simple instructions (for ex primitives) while instruction 3 has its
53+
* own large dependency tree.
54+
*
55+
* ## Issues/things to explore:
56+
*
57+
* - An obvious improvement is to weight the nodes and emit dependencies based on weight. Alternatively, we could try to
58+
* determine the reactive dependencies of each node, and try to emit nodes that have the same dependencies together.
59+
* - Reordering destructure statements means that we also end up deferring the evaluation of its RHS. So i noticed some
60+
* `const [state, setState] = useState(...)` getting moved around. But i think i might have just messed up the bit that
61+
* ensures non-reorderable instructions (like the useState() call here) are serialized. So this should just be a simple fix,
62+
* if i didn't already fix it (need to go back through the fixture output changes)
63+
* - I also noticed that destructuring being moved meant that some reactive scopes ended up with less precise input, because
64+
* the destructure moved into the reactive scope itself (so the scope depends on the rvalue of the destructure, not the lvalues).
65+
* This is weird, i need to debug.
66+
* - Probably more things.
67+
*/
68+
export function instructionReordering(fn: HIRFunction): void {
69+
const globalDependencies: Dependencies = new Map();
70+
for (const [, block] of fn.body.blocks) {
71+
reorderBlock(block, globalDependencies);
72+
}
73+
markInstructionIds(fn.body);
74+
}
75+
76+
type Dependencies = Map<IdentifierId, Node>;
77+
type Node = {
78+
instruction: Instruction | null;
79+
dependencies: Array<IdentifierId>;
80+
};
81+
82+
function reorderBlock(
83+
block: BasicBlock,
84+
globalDependencies: Dependencies
85+
): void {
86+
const dependencies: Dependencies = new Map();
87+
const locals = new Map<string, IdentifierId>();
88+
let previousIdentifier: IdentifierId | null = null;
89+
for (const instr of block.instructions) {
90+
const node: Node = getOrInsertDefault(
91+
dependencies,
92+
instr.lvalue.identifier.id,
93+
{
94+
instruction: instr,
95+
dependencies: [],
96+
}
97+
);
98+
if (getReorderingLevel(instr) === ReorderingLevel.None) {
99+
if (previousIdentifier !== null) {
100+
node.dependencies.push(previousIdentifier);
101+
}
102+
previousIdentifier = instr.lvalue.identifier.id;
103+
}
104+
for (const operand of eachInstructionValueOperand(instr.value)) {
105+
if (
106+
operand.identifier.name !== null &&
107+
operand.identifier.name.kind === "named"
108+
) {
109+
const previous = locals.get(operand.identifier.name.value);
110+
if (previous !== undefined) {
111+
node.dependencies.push(previous);
112+
} else {
113+
locals.set(operand.identifier.name.value, instr.lvalue.identifier.id);
114+
node.dependencies.push(operand.identifier.id);
115+
}
116+
} else {
117+
if (dependencies.has(operand.identifier.id)) {
118+
node.dependencies.push(operand.identifier.id);
119+
}
120+
}
121+
}
122+
dependencies.set(instr.lvalue.identifier.id, node);
123+
124+
for (const lvalue of eachInstructionValueLValue(instr.value)) {
125+
const lvalueNode = getOrInsertDefault(
126+
dependencies,
127+
lvalue.identifier.id,
128+
{
129+
instruction: null,
130+
dependencies: [],
131+
}
132+
);
133+
lvalueNode.dependencies.push(instr.lvalue.identifier.id);
134+
if (
135+
lvalue.identifier.name !== null &&
136+
lvalue.identifier.name.kind === "named"
137+
) {
138+
const previous = locals.get(lvalue.identifier.name.value);
139+
if (previous !== undefined) {
140+
node.dependencies.push(previous);
141+
}
142+
}
143+
}
144+
}
145+
146+
const instructions: Array<Instruction> = [];
147+
148+
function emit(id: IdentifierId): void {
149+
const node = dependencies.get(id) ?? globalDependencies.get(id);
150+
if (node == null) {
151+
return;
152+
}
153+
dependencies.delete(id);
154+
globalDependencies.delete(id);
155+
for (const dep of node.dependencies) {
156+
emit(dep);
157+
}
158+
if (node.instruction !== null) {
159+
instructions.push(node.instruction);
160+
}
161+
}
162+
163+
for (const operand of eachTerminalOperand(block.terminal)) {
164+
emit(operand.identifier.id);
165+
}
166+
for (const id of Array.from(dependencies.keys()).reverse()) {
167+
const node = dependencies.get(id);
168+
if (node == null) {
169+
continue;
170+
}
171+
if (
172+
node.instruction !== null &&
173+
getReorderingLevel(node.instruction) === ReorderingLevel.Global
174+
) {
175+
globalDependencies.set(id, node);
176+
} else {
177+
emit(id);
178+
}
179+
}
180+
block.instructions = instructions;
181+
}
182+
183+
function printDeps(deps: Dependencies): string {
184+
return (
185+
"[\n" +
186+
Array.from(deps)
187+
.map(
188+
([id, dep]) =>
189+
`$${id} ${
190+
dep.instruction != null ? printInstruction(dep.instruction) : ""
191+
} deps=[${dep.dependencies.map((x) => `$${x}`).join(", ")}]`
192+
)
193+
.join("\n") +
194+
"\n]"
195+
);
196+
}
197+
198+
enum ReorderingLevel {
199+
None = "none",
200+
Local = "local",
201+
Global = "global",
202+
}
203+
function getReorderingLevel(instr: Instruction): ReorderingLevel {
204+
switch (instr.value.kind) {
205+
case "JsxExpression":
206+
case "JsxFragment":
207+
case "JSXText":
208+
case "LoadGlobal":
209+
case "Primitive":
210+
case "TemplateLiteral": {
211+
return ReorderingLevel.Global;
212+
}
213+
case "ArrayExpression":
214+
case "ObjectExpression":
215+
case "LoadLocal":
216+
case "Destructure":
217+
case "StoreLocal": {
218+
return ReorderingLevel.Local;
219+
}
220+
default: {
221+
return ReorderingLevel.None;
222+
}
223+
}
224+
}

compiler/packages/babel-plugin-react-compiler/src/Utils/utils.ts

+3-1
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,9 @@ export function assertExhaustive(_: never, errorMsg: string): never {
3030
throw new Error(errorMsg);
3131
}
3232

33-
// Modifies @param array in place, retaining only the items where the predicate returns true.
33+
/**
34+
* Modifies @param array in place, retaining only the items where the predicate returns true.
35+
*/
3436
export function retainWhere<T>(
3537
array: Array<T>,
3638
predicate: (item: T) => boolean

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-capture-in-method-receiver-and-mutate.expect.md

+2-3
Original file line numberDiff line numberDiff line change
@@ -34,13 +34,12 @@ function Component() {
3434
const $ = _c(1);
3535
let t0;
3636
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
37+
const x = [];
3738
const a = makeObject_Primitives();
3839

39-
const x = [];
40+
t0 = [x, a];
4041
x.push(a);
41-
4242
mutate(x);
43-
t0 = [x, a];
4443
$[0] = t0;
4544
} else {
4645
t0 = $[0];

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-capture-in-method-receiver.expect.md

+7-14
Original file line numberDiff line numberDiff line change
@@ -20,26 +20,19 @@ function Component() {
2020
```javascript
2121
import { c as _c } from "react/compiler-runtime";
2222
function Component() {
23-
const $ = _c(2);
23+
const $ = _c(1);
2424
let t0;
2525
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
26-
t0 = someObj();
27-
$[0] = t0;
28-
} else {
29-
t0 = $[0];
30-
}
31-
const a = t0;
32-
let t1;
33-
if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
3426
const x = [];
35-
x.push(a);
27+
const a = someObj();
3628

37-
t1 = [x, a];
38-
$[1] = t1;
29+
t0 = [x, a];
30+
x.push(a);
31+
$[0] = t0;
3932
} else {
40-
t1 = $[1];
33+
t0 = $[0];
4134
}
42-
return t1;
35+
return t0;
4336
}
4437

4538
```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-computed-load.expect.md

+7-6
Original file line numberDiff line numberDiff line change
@@ -19,19 +19,20 @@ function component(a) {
1919
import { c as _c } from "react/compiler-runtime";
2020
function component(a) {
2121
const $ = _c(2);
22-
let x;
22+
let t0;
2323
if ($[0] !== a) {
24-
x = { a };
25-
const y = {};
24+
const x = { a };
2625

26+
t0 = x;
27+
const y = {};
2728
y.x = x.a;
2829
mutate(y);
2930
$[0] = a;
30-
$[1] = x;
31+
$[1] = t0;
3132
} else {
32-
x = $[1];
33+
t0 = $[1];
3334
}
34-
return x;
35+
return t0;
3536
}
3637

3738
```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-nested-member-path-mutate.expect.md

+8-6
Original file line numberDiff line numberDiff line change
@@ -20,19 +20,21 @@ function component() {
2020
import { c as _c } from "react/compiler-runtime";
2121
function component() {
2222
const $ = _c(1);
23-
let x;
23+
let t0;
2424
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
25-
const z = [];
25+
const x = {};
26+
27+
t0 = x;
2628
const y = {};
29+
const z = [];
2730
y.z = z;
28-
x = {};
2931
x.y = y;
3032
mutate(x.y.z);
31-
$[0] = x;
33+
$[0] = t0;
3234
} else {
33-
x = $[0];
35+
t0 = $[0];
3436
}
35-
return x;
37+
return t0;
3638
}
3739

3840
```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-nested-member-path.expect.md

+8-6
Original file line numberDiff line numberDiff line change
@@ -25,18 +25,20 @@ export const FIXTURE_ENTRYPOINT = {
2525
import { c as _c } from "react/compiler-runtime";
2626
function component() {
2727
const $ = _c(1);
28-
let x;
28+
let t0;
2929
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
30-
const z = [];
30+
const x = {};
31+
32+
t0 = x;
3133
const y = {};
34+
const z = [];
3235
y.z = z;
33-
x = {};
3436
x.y = y;
35-
$[0] = x;
37+
$[0] = t0;
3638
} else {
37-
x = $[0];
39+
t0 = $[0];
3840
}
39-
return x;
41+
return t0;
4042
}
4143

4244
export const FIXTURE_ENTRYPOINT = {

0 commit comments

Comments
 (0)