@samitouri / QOS-React / commits / 2caaa05c08

[compiler] Optimize instruction reordering

Note: due to a bad rebase i included #29883 here. Both were stamped so i'm not gonna bother splitting it back up aain. This PR includes two changes: * First, allow `LoadLocal` to be reordered if a) the load occurs after the last write to a variable and b) the LoadLocal lvalue is used exactly once * Uses a more optimal reordering for statement blocks, while keeping the existing approach for expression blocks. In #29863 I tried to find a clean way to share code for emitting instructions between value blocks and regular blocks. The catch is that value blocks have special meaning for their final instruction — that's the value of the block — so reordering can't change the last instruction. However, in finding a clean way to share code for these two categories of code, i also inadvertently reduced the effectiveness of the optimization. This PR updates to use different strategies for these two kinds of blocks: value blocks use the code from #29863 where we first emit all non-reorderable instructions in their original order, then try to emit reorderable values. The reason this is suboptimal, though, is that we want to move instructions closer to their dependencies so that they can invalidate (merge) together. Emitting the reorderable values last prevents this. So for normal blocks, we now emit terminal operands first. This will invariably cause some of the non-reorderable instructions to be emitted, but it will intersperse reoderable instructions in between, right after their dependencies. This maximizes our ability to merge scopes. I think the complexity cost of two strategies is worth the benefit, as evidenced by the reduced memo slots in the fixtures. ghstack-source-id: ad3e516fa474235ced8c5d56f4541d2a7c413608 Pull Request resolved: https://github.com/facebook/react/pull/29882

Joe Savona committed Jun 21, 2024 at 16:48 UTC 2caaa05c08c345f1edddc952cef3fa53c177d612
4 files changed +247 -122
compiler/packages/babel-plugin-react-compiler/src/Optimization/InstructionReordering.ts
+213 -65
@@ -13,16 +13,19 @@ import {
13 HIRFunction,
14 IdentifierId,
15 Instruction,
16 + InstructionId,
17 + Place,
18 isExpressionBlockKind,
19 + makeInstructionId,
20 markInstructionIds,
21 } from "../HIR";
22 import { printInstruction } from "../HIR/PrintHIR";
23 import {
24 + eachInstructionLValue,
25 eachInstructionValueLValue,
26 eachInstructionValueOperand,
27 eachTerminalOperand,
28 } from "../HIR/visitors";
25 -import { mayAllocate } from "../ReactiveScopes/InferReactiveScopeVariables";
29 import { getOrInsertWith } from "../Utils/utils";
30
31 /**
@@ -69,8 +72,9 @@ import { getOrInsertWith } from "../Utils/utils";
72 export function instructionReordering(fn: HIRFunction): void {
73 // Shared nodes are emitted when they are first used
74 const shared: Nodes = new Map();
75 + const references = findReferencedRangeOfTemporaries(fn);
76 for (const [, block] of fn.body.blocks) {
73 - reorderBlock(fn.env, block, shared);
77 + reorderBlock(fn.env, block, shared, references);
78 }
79 CompilerError.invariant(shared.size === 0, {
80 reason: `InstructionReordering: expected all reorderable nodes to have been emitted`,
@@ -88,13 +92,79 @@ type Nodes = Map<IdentifierId, Node>;
92 type Node = {
93 instruction: Instruction | null;
94 dependencies: Set<IdentifierId>;
95 + reorderability: Reorderability;
96 depth: number | null;
97 };
98
99 +// Inclusive start and end
100 +type References = {
101 + singleUseIdentifiers: SingleUseIdentifiers;
102 + lastAssignments: LastAssignments;
103 +};
104 +type LastAssignments = Map<string, InstructionId>;
105 +type SingleUseIdentifiers = Set<IdentifierId>;
106 +enum ReferenceKind {
107 + Read,
108 + Write,
109 +}
110 +function findReferencedRangeOfTemporaries(fn: HIRFunction): References {
111 + const singleUseIdentifiers = new Map<IdentifierId, number>();
112 + const lastAssignments: LastAssignments = new Map();
113 + function reference(
114 + instr: InstructionId,
115 + place: Place,
116 + kind: ReferenceKind
117 + ): void {
118 + if (
119 + place.identifier.name !== null &&
120 + place.identifier.name.kind === "named"
121 + ) {
122 + if (kind === ReferenceKind.Write) {
123 + const name = place.identifier.name.value;
124 + const previous = lastAssignments.get(name);
125 + if (previous === undefined) {
126 + lastAssignments.set(name, instr);
127 + } else {
128 + lastAssignments.set(
129 + name,
130 + makeInstructionId(Math.max(previous, instr))
131 + );
132 + }
133 + }
134 + return;
135 + } else if (kind === ReferenceKind.Read) {
136 + const previousCount = singleUseIdentifiers.get(place.identifier.id) ?? 0;
137 + singleUseIdentifiers.set(place.identifier.id, previousCount + 1);
138 + }
139 + }
140 + for (const [, block] of fn.body.blocks) {
141 + for (const instr of block.instructions) {
142 + for (const operand of eachInstructionValueLValue(instr.value)) {
143 + reference(instr.id, operand, ReferenceKind.Read);
144 + }
145 + for (const lvalue of eachInstructionLValue(instr)) {
146 + reference(instr.id, lvalue, ReferenceKind.Write);
147 + }
148 + }
149 + for (const operand of eachTerminalOperand(block.terminal)) {
150 + reference(block.terminal.id, operand, ReferenceKind.Read);
151 + }
152 + }
153 + return {
154 + singleUseIdentifiers: new Set(
155 + [...singleUseIdentifiers]
156 + .filter(([, count]) => count === 1)
157 + .map(([id]) => id)
158 + ),
159 + lastAssignments,
160 + };
161 +}
162 +
163 function reorderBlock(
164 env: Environment,
165 block: BasicBlock,
97 - shared: Nodes
166 + shared: Nodes,
167 + references: References
168 ): void {
169 const locals: Nodes = new Map();
170 const named: Map<string, IdentifierId> = new Map();
@@ -102,6 +172,7 @@ function reorderBlock(
172 for (const instr of block.instructions) {
173 const { lvalue, value } = instr;
174 // Get or create a node for this lvalue
175 + const reorderability = getReorderability(instr, references);
176 const node = getOrInsertWith(
177 locals,
178 lvalue.identifier.id,
@@ -109,6 +180,7 @@ function reorderBlock(
180 ({
181 instruction: instr,
182 dependencies: new Set(),
183 + reorderability,
184 depth: null,
185 }) as Node
186 );
@@ -116,7 +188,7 @@ function reorderBlock(
188 * Ensure non-reoderable instructions have their order retained by
189 * adding explicit dependencies to the previous such instruction.
190 */
119 - if (getReoderability(instr) === Reorderability.Nonreorderable) {
191 + if (reorderability === Reorderability.Nonreorderable) {
192 if (previous !== null) {
193 node.dependencies.add(previous);
194 }
@@ -172,66 +244,125 @@ function reorderBlock(
244
245 DEBUG && console.log(`bb${block.id}`);
246
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.
247 + /**
248 + * The ideal order for emitting instructions may change the final instruction,
249 + * but value blocks have special semantics for the final instruction of a block -
250 + * that's the expression's value!. So we choose between a less optimal strategy
251 + * for value blocks which preserves the final instruction order OR a more optimal
252 + * ordering for statement-y blocks.
253 */
187 - if (isExpressionBlockKind(block.kind) && block.instructions.length !== 0) {
188 - DEBUG && console.log(`(block value)`);
189 - DEBUG &&
190 - print(
254 + if (isExpressionBlockKind(block.kind)) {
255 + // First emit everything that can't be reordered
256 + if (previous !== null) {
257 + DEBUG && console.log(`(last non-reorderable instruction)`);
258 + DEBUG && print(env, locals, shared, seen, previous);
259 + emit(env, locals, shared, nextInstructions, previous);
260 + }
261 + /*
262 + * For "value" blocks the final instruction represents its value, so we have to be
263 + * careful to not change the ordering. Emit the last instruction explicitly.
264 + * Any non-reorderable instructions will get emitted first, and any unused
265 + * reorderable instructions can be deferred to the shared node list.
266 + */
267 + if (block.instructions.length !== 0) {
268 + DEBUG && console.log(`(block value)`);
269 + DEBUG &&
270 + print(
271 + env,
272 + locals,
273 + shared,
274 + seen,
275 + block.instructions.at(-1)!.lvalue.identifier.id
276 + );
277 + emit(
278 env,
279 locals,
280 shared,
194 - seen,
281 + nextInstructions,
282 block.instructions.at(-1)!.lvalue.identifier.id
283 );
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;
284 }
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`,
285 + /*
286 + * Then emit the dependencies of the terminal operand. In many cases they will have
287 + * already been emitted in the previous step and this is a no-op.
288 + * TODO: sort the dependencies based on weight, like we do for other nodes. Not a big
289 + * deal though since most terminals have a single operand
290 + */
291 + for (const operand of eachTerminalOperand(block.terminal)) {
292 + DEBUG && console.log(`(terminal operand)`);
293 + DEBUG && print(env, locals, shared, seen, operand.identifier.id);
294 + emit(env, locals, shared, nextInstructions, operand.identifier.id);
295 + }
296 + // Anything not emitted yet is globally reorderable
297 + for (const [id, node] of locals) {
298 + if (node.instruction == null) {
299 + continue;
300 }
232 - );
233 - DEBUG && console.log(`save shared: $${id}`);
234 - shared.set(id, node);
301 + CompilerError.invariant(
302 + node.reorderability === Reorderability.Reorderable,
303 + {
304 + reason: `Expected all remaining instructions to be reorderable`,
305 + loc: node.instruction?.loc ?? block.terminal.loc,
306 + description:
307 + node.instruction != null
308 + ? `Instruction [${node.instruction.id}] was not emitted yet but is not reorderable`
309 + : `Lvalue $${id} was not emitted yet but is not reorderable`,
310 + }
311 + );
312 +
313 + DEBUG && console.log(`save shared: $${id}`);
314 + shared.set(id, node);
315 + }
316 + } else {
317 + /**
318 + * If this is not a value block, then the order within the block doesn't matter
319 + * and we can optimize more. The observation is that blocks often have instructions
320 + * such as:
321 + *
322 + * ```
323 + * t$0 = nonreorderable
324 + * t$1 = nonreorderable <-- this gets in the way of merging t$0 and t$2
325 + * t$2 = reorderable deps[ t$0 ]
326 + * return t$2
327 + * ```
328 + *
329 + * Ie where there is some pair of nonreorderable+reorderable values, with some intervening
330 + * also non-reorderable instruction. If we emit all non-reorderable instructions first,
331 + * then we'll keep the original order. But reordering instructions doesn't just mean moving
332 + * them later: we can also move them _earlier_. By starting from terminal operands we
333 + * end up emitting:
334 + *
335 + * ```
336 + * t$0 = nonreorderable // dep of t$2
337 + * t$2 = reorderable deps[ t$0 ]
338 + * t$1 = nonreorderable <-- not in the way of merging anymore!
339 + * return t$2
340 + * ```
341 + *
342 + * Ie all nonreorderable transitive deps of the terminal operands will get emitted first,
343 + * but we'll be able to intersperse the depending reorderable instructions in between
344 + * them in a way that works better with scope merging.
345 + */
346 + for (const operand of eachTerminalOperand(block.terminal)) {
347 + DEBUG && console.log(`(terminal operand)`);
348 + DEBUG && print(env, locals, shared, seen, operand.identifier.id);
349 + emit(env, locals, shared, nextInstructions, operand.identifier.id);
350 + }
351 + // Anything not emitted yet is globally reorderable
352 + for (const id of Array.from(locals.keys()).reverse()) {
353 + const node = locals.get(id);
354 + if (node === undefined) {
355 + continue;
356 + }
357 + if (node.reorderability === Reorderability.Reorderable) {
358 + DEBUG && console.log(`save shared: $${id}`);
359 + shared.set(id, node);
360 + } else {
361 + DEBUG && console.log("leftover");
362 + DEBUG && print(env, locals, shared, seen, id);
363 + emit(env, locals, shared, nextInstructions, id);
364 + }
365 + }
366 }
367
368 block.instructions = nextInstructions;
@@ -247,8 +378,7 @@ function getDepth(env: Environment, nodes: Nodes, id: IdentifierId): number {
378 return node.depth;
379 }
380 node.depth = 0; // in case of cycles
250 - let depth =
251 - node.instruction != null && mayAllocate(env, node.instruction) ? 1 : 0;
381 + let depth = node.reorderability === Reorderability.Reorderable ? 1 : 10;
382 for (const dep of node.dependencies) {
383 depth += getDepth(env, nodes, dep);
384 }
@@ -265,7 +395,7 @@ function print(
395 depth: number = 0
396 ): void {
397 if (seen.has(id)) {
268 - console.log(`${"| ".repeat(depth)}$${id} <skipped>`);
398 + DEBUG && console.log(`${"| ".repeat(depth)}$${id} <skipped>`);
399 return;
400 }
401 seen.add(id);
@@ -282,11 +412,12 @@ function print(
412 for (const dep of deps) {
413 print(env, locals, shared, seen, dep, depth + 1);
414 }
285 - console.log(
286 - `${"| ".repeat(depth)}$${id} ${printNode(node)} deps=[${deps
287 - .map((x) => `$${x}`)
288 - .join(", ")}]`
289 - );
415 + DEBUG &&
416 + console.log(
417 + `${"| ".repeat(depth)}$${id} ${printNode(node)} deps=[${deps
418 + .map((x) => `$${x}`)
419 + .join(", ")}] depth=${node.depth}`
420 + );
421 }
422
423 function printNode(node: Node): string {
@@ -336,7 +467,10 @@ enum Reorderability {
467 Reorderable,
468 Nonreorderable,
469 }
339 -function getReoderability(instr: Instruction): Reorderability {
470 +function getReorderability(
471 + instr: Instruction,
472 + references: References
473 +): Reorderability {
474 switch (instr.value.kind) {
475 case "JsxExpression":
476 case "JsxFragment":
@@ -348,6 +482,20 @@ function getReoderability(instr: Instruction): Reorderability {
482 case "UnaryExpression": {
483 return Reorderability.Reorderable;
484 }
485 + case "LoadLocal": {
486 + const name = instr.value.place.identifier.name;
487 + if (name !== null && name.kind === "named") {
488 + const lastAssignment = references.lastAssignments.get(name.value);
489 + if (
490 + lastAssignment !== undefined &&
491 + lastAssignment < instr.id &&
492 + references.singleUseIdentifiers.has(instr.lvalue.identifier.id)
493 + ) {
494 + return Reorderability.Reorderable;
495 + }
496 + }
497 + return Reorderability.Nonreorderable;
498 + }
499 default: {
500 return Reorderability.Nonreorderable;
501 }
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts
+1 -4
@@ -186,10 +186,7 @@ export function isMutable({ id }: Instruction, place: Place): boolean {
186 return id >= range.start && id < range.end;
187 }
188
189 -export function mayAllocate(
190 - env: Environment,
191 - instruction: Instruction
192 -): boolean {
189 +function mayAllocate(env: Environment, instruction: Instruction): boolean {
190 const { value } = instruction;
191 switch (value.kind) {
192 case "Destructure": {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-consecutive-scopes-reordering.expect.md
+21 -33
@@ -34,59 +34,47 @@ import { useState } from "react";
34 import { Stringify } from "shared-runtime";
35
36 function Component() {
37 - const $ = _c(10);
37 + const $ = _c(7);
38 const [state, setState] = useState(0);
39 let t0;
40 + let t1;
41 if ($[0] !== state) {
41 - t0 = () => setState(state + 1);
42 + t0 = (
43 + <button data-testid="button" onClick={() => setState(state + 1)}>
44 + increment
45 + </button>
46 + );
47 + t1 = <span>{state}</span>;
48 $[0] = state;
49 $[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 + t0 = $[1];
53 t1 = $[2];
54 }
55 let t2;
55 - if ($[3] !== state) {
56 - t2 = <span>{state}</span>;
57 - $[3] = state;
58 - $[4] = t2;
56 + if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
57 + t2 = <Stringify text="Counter" />;
58 + $[3] = t2;
59 } else {
60 - t2 = $[4];
60 + t2 = $[3];
61 }
62 let t3;
63 - if ($[5] !== t0) {
63 + if ($[4] !== t1 || $[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 = (
65 <div>
78 - {t1}
66 {t2}
80 - {t3}
67 + {t1}
68 + {t0}
69 </div>
70 );
83 - $[7] = t2;
84 - $[8] = t3;
85 - $[9] = t4;
71 + $[4] = t1;
72 + $[5] = t0;
73 + $[6] = t3;
74 } else {
87 - t4 = $[9];
75 + t3 = $[6];
76 }
89 - return t4;
77 + return t3;
78 }
79
80 export const FIXTURE_ENTRYPOINT = {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.expect.md
+12 -20
@@ -27,7 +27,7 @@ import { c as _c } from "react/compiler-runtime"; // @enableInstructionReorderin
27 import { useState } from "react";
28
29 function Component() {
30 - const $ = _c(6);
30 + const $ = _c(4);
31 const [state, setState] = useState(0);
32 let t0;
33 if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -40,34 +40,26 @@ function Component() {
40 }
41 const onClick = t0;
42 let t1;
43 - if ($[1] !== state) {
44 - t1 = <span>Count: {state}</span>;
45 - $[1] = state;
46 - $[2] = t1;
43 + if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
44 + t1 = <button onClick={onClick}>Increment</button>;
45 + $[1] = t1;
46 } else {
48 - t1 = $[2];
47 + t1 = $[1];
48 }
49 let t2;
51 - if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
52 - t2 = <button onClick={onClick}>Increment</button>;
53 - $[3] = t2;
54 - } else {
55 - t2 = $[3];
56 - }
57 - let t3;
58 - if ($[4] !== t1) {
59 - t3 = (
50 + if ($[2] !== state) {
51 + t2 = (
52 <>
53 + <span>Count: {state}</span>
54 {t1}
62 - {t2}
55 </>
56 );
65 - $[4] = t1;
66 - $[5] = t3;
57 + $[2] = state;
58 + $[3] = t2;
59 } else {
68 - t3 = $[5];
60 + t2 = $[3];
61 }
70 - return t3;
62 + return t2;
63 }
64
65 ```