cranelift_codegen/egraph/elaborate.rs
1//! Elaboration phase: lowers EGraph back to sequences of operations
2//! in CFG nodes.
3
4use super::Stats;
5use super::cost::Cost;
6use crate::ctxhash::NullCtx;
7use crate::dominator_tree::DominatorTreePreorder;
8use crate::hash_map::Entry as HashEntry;
9use crate::inst_predicates::is_pure_for_egraph;
10use crate::ir::{Block, Function, Inst, Value, ValueDef};
11use crate::loop_analysis::{Loop, LoopAnalysis};
12use crate::scoped_hash_map::ScopedHashMap;
13use crate::trace;
14use alloc::vec::Vec;
15use cranelift_control::ControlPlane;
16use cranelift_entity::{SecondaryMap, packed_option::ReservedValue};
17use rustc_hash::{FxHashMap, FxHashSet};
18use smallvec::{SmallVec, smallvec};
19
20pub(crate) struct Elaborator<'a> {
21 func: &'a mut Function,
22 domtree: &'a DominatorTreePreorder,
23 loop_analysis: &'a LoopAnalysis,
24 /// Map from Value that is produced by a pure Inst (and was thus
25 /// not in the side-effecting skeleton) to the value produced by
26 /// an elaborated inst (placed in the layout) to whose results we
27 /// refer in the final code.
28 ///
29 /// The first time we use some result of an instruction during
30 /// elaboration, we can place it and insert an identity map (inst
31 /// results to that same inst's results) in this scoped
32 /// map. Within that block and its dom-tree children, that mapping
33 /// is visible and we can continue to use it. This allows us to
34 /// avoid cloning the instruction. However, if we pop that scope
35 /// and use it somewhere else as well, we will need to
36 /// duplicate. We detect this case by checking, when a value that
37 /// we want is not present in this map, whether the producing inst
38 /// is already placed in the Layout. If so, we duplicate, and
39 /// insert non-identity mappings from the original inst's results
40 /// to the cloned inst's results.
41 ///
42 /// Note that as values may refer to unions that represent a subset
43 /// of a larger eclass, it's not valid to walk towards the root of a
44 /// union tree: doing so would potentially equate values that fall
45 /// on different branches of the dominator tree.
46 value_to_elaborated_value: ScopedHashMap<Value, ElaboratedValue>,
47 /// Map from Value to the best (lowest-cost) Value in its eclass
48 /// (tree of union value-nodes).
49 value_to_best_value: SecondaryMap<Value, BestEntry>,
50 /// Stack of blocks and loops in current elaboration path.
51 loop_stack: SmallVec<[LoopStackEntry; 8]>,
52 /// The current block into which we are elaborating.
53 cur_block: Block,
54 /// Values that opt rules have indicated should be rematerialized
55 /// in every block they are used (e.g., immediates or other
56 /// "cheap-to-compute" ops).
57 remat_values: &'a FxHashSet<Value>,
58 /// Explicitly-unrolled value elaboration stack.
59 elab_stack: Vec<ElabStackEntry>,
60 /// Results from the elab stack.
61 elab_result_stack: Vec<ElaboratedValue>,
62 /// Explicitly-unrolled block elaboration stack.
63 block_stack: Vec<BlockStackEntry>,
64 /// Copies of values that have been rematerialized.
65 remat_copies: FxHashMap<(Block, Value), Value>,
66 /// Stats for various events during egraph processing, to help
67 /// with optimization of this infrastructure.
68 stats: &'a mut Stats,
69 /// Chaos-mode control-plane so we can test that we still get
70 /// correct results when our heuristics make bad decisions.
71 ctrl_plane: &'a mut ControlPlane,
72}
73
74#[derive(Clone, Copy, Debug, PartialEq, Eq)]
75struct BestEntry(Cost, Value);
76
77impl PartialOrd for BestEntry {
78 fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
79 Some(self.cmp(other))
80 }
81}
82
83impl Ord for BestEntry {
84 #[inline]
85 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
86 self.0.cmp(&other.0).then_with(|| {
87 // Note that this comparison is reversed. When costs are equal,
88 // prefer the value with the bigger index. This is a heuristic that
89 // prefers results of rewrites to the original value, since we
90 // expect that our rewrites are generally improvements.
91 self.1.cmp(&other.1).reverse()
92 })
93 }
94}
95
96#[derive(Clone, Copy, Debug)]
97struct ElaboratedValue {
98 in_block: Block,
99 value: Value,
100}
101
102#[derive(Clone, Debug)]
103struct LoopStackEntry {
104 /// The loop identifier.
105 lp: Loop,
106 /// The hoist point: a block that immediately dominates this
107 /// loop. May not be an immediate predecessor, but will be a valid
108 /// point to place all loop-invariant ops: they must depend only
109 /// on inputs that dominate the loop, so are available at (the end
110 /// of) this block.
111 hoist_block: Block,
112 /// The depth in the scope map.
113 scope_depth: u32,
114}
115
116#[derive(Clone, Debug)]
117enum ElabStackEntry {
118 /// Next action is to resolve this value into an elaborated inst
119 /// (placed into the layout) that produces the value, and
120 /// recursively elaborate the insts that produce its args.
121 ///
122 /// Any inserted ops should be inserted before `before`, which is
123 /// the instruction demanding this value.
124 Start { value: Value, before: Inst },
125 /// Args have been pushed; waiting for results.
126 PendingInst {
127 inst: Inst,
128 result_idx: usize,
129 num_args: usize,
130 before: Inst,
131 },
132}
133
134#[derive(Clone, Debug)]
135enum BlockStackEntry {
136 Elaborate { block: Block, idom: Option<Block> },
137 Pop,
138}
139
140impl<'a> Elaborator<'a> {
141 pub(crate) fn new(
142 func: &'a mut Function,
143 domtree: &'a DominatorTreePreorder,
144 loop_analysis: &'a LoopAnalysis,
145 remat_values: &'a FxHashSet<Value>,
146 stats: &'a mut Stats,
147 ctrl_plane: &'a mut ControlPlane,
148 ) -> Self {
149 let num_values = func.dfg.num_values();
150 let mut value_to_best_value =
151 SecondaryMap::with_default(BestEntry(Cost::infinity(), Value::reserved_value()));
152 value_to_best_value.resize(num_values);
153 Self {
154 func,
155 domtree,
156 loop_analysis,
157 value_to_elaborated_value: ScopedHashMap::with_capacity(num_values),
158 value_to_best_value,
159 loop_stack: smallvec![],
160 cur_block: Block::reserved_value(),
161 remat_values,
162 elab_stack: vec![],
163 elab_result_stack: vec![],
164 block_stack: vec![],
165 remat_copies: FxHashMap::default(),
166 stats,
167 ctrl_plane,
168 }
169 }
170
171 fn start_block(&mut self, idom: Option<Block>, block: Block) {
172 trace!(
173 "start_block: block {:?} with idom {:?} at loop depth {:?} scope depth {}",
174 block,
175 idom,
176 self.loop_stack.len(),
177 self.value_to_elaborated_value.depth()
178 );
179
180 // Pop any loop levels we're no longer in.
181 while let Some(inner_loop) = self.loop_stack.last() {
182 if self.loop_analysis.is_in_loop(block, inner_loop.lp) {
183 break;
184 }
185 self.loop_stack.pop();
186 }
187
188 // Note that if the *entry* block is a loop header, we will
189 // not make note of the loop here because it will not have an
190 // immediate dominator. We must disallow this case because we
191 // will skip adding the `LoopStackEntry` here but our
192 // `LoopAnalysis` will otherwise still make note of this loop
193 // and loop depths will not match.
194 if let Some(idom) = idom {
195 if let Some(lp) = self.loop_analysis.is_loop_header(block) {
196 self.loop_stack.push(LoopStackEntry {
197 lp,
198 // Any code hoisted out of this loop will have code
199 // placed in `idom`, and will have def mappings
200 // inserted in to the scoped hashmap at that block's
201 // level.
202 hoist_block: idom,
203 scope_depth: (self.value_to_elaborated_value.depth() - 1) as u32,
204 });
205 trace!(
206 " -> loop header, pushing; depth now {}",
207 self.loop_stack.len()
208 );
209 }
210 } else {
211 debug_assert!(
212 self.loop_analysis.is_loop_header(block).is_none(),
213 "Entry block (domtree root) cannot be a loop header!"
214 );
215 }
216
217 trace!("block {}: loop stack is {:?}", block, self.loop_stack);
218
219 self.cur_block = block;
220 }
221
222 fn compute_best_values(&mut self) {
223 let best = &mut self.value_to_best_value;
224
225 // We can't make random decisions inside the fixpoint loop below because
226 // that could cause values to change on every iteration of the loop,
227 // which would make the loop never terminate. So in chaos testing
228 // mode we need a form of making suboptimal decisions that is fully
229 // deterministic. We choose to simply make the worst decision we know
230 // how to do instead of the best.
231 let use_worst = self.ctrl_plane.get_decision();
232
233 // Do a fixpoint loop to compute the best value for each eclass.
234 //
235 // The maximum number of iterations is the length of the longest chain
236 // of `vNN -> vMM` edges in the dataflow graph where `NN < MM`, so this
237 // is *technically* quadratic, but `cranelift-frontend` won't construct
238 // any such edges. NaN canonicalization will introduce some of these
239 // edges, but they are chains of only two or three edges. So in
240 // practice, we *never* do more than a handful of iterations here unless
241 // (a) we parsed the CLIF from text and the text was funkily numbered,
242 // which we don't really care about, or (b) the CLIF producer did
243 // something weird, in which case it is their responsibility to stop
244 // doing that.
245 trace!(
246 "Entering fixpoint loop to compute the {} values for each eclass",
247 if use_worst {
248 "worst (chaos mode)"
249 } else {
250 "best"
251 }
252 );
253 let mut keep_going = true;
254 while keep_going {
255 keep_going = false;
256 trace!(
257 "fixpoint iteration {}",
258 self.stats.elaborate_best_cost_fixpoint_iters
259 );
260 self.stats.elaborate_best_cost_fixpoint_iters += 1;
261
262 for (value, def) in self.func.dfg.values_and_defs() {
263 trace!("computing best for value {:?} def {:?}", value, def);
264 let orig_best_value = best[value];
265
266 match def {
267 ValueDef::Union(x, y) => {
268 // Pick the best of the two options based on
269 // min-cost. This works because each element of `best`
270 // is a `(cost, value)` tuple; `cost` comes first so
271 // the natural comparison works based on cost, and
272 // breaks ties based on value number.
273 best[value] = if use_worst {
274 if best[x].1.is_reserved_value() {
275 best[y]
276 } else if best[y].1.is_reserved_value() {
277 best[x]
278 } else {
279 std::cmp::max(best[x], best[y])
280 }
281 } else {
282 std::cmp::min(best[x], best[y])
283 };
284 trace!(
285 " -> best of union({:?}, {:?}) = {:?}",
286 best[x], best[y], best[value]
287 );
288 }
289 ValueDef::Param(_, _) => {
290 best[value] = BestEntry(Cost::zero(), value);
291 }
292 // If the Inst is inserted into the layout (which is,
293 // at this point, only the side-effecting skeleton),
294 // then it must be computed and thus we give it zero
295 // cost.
296 ValueDef::Result(inst, _) => {
297 if let Some(_) = self.func.layout.inst_block(inst) {
298 best[value] = BestEntry(Cost::zero(), value);
299 } else {
300 let inst_data = &self.func.dfg.insts[inst];
301 // N.B.: at this point we know that the opcode is
302 // pure, so `pure_op_cost`'s precondition is
303 // satisfied.
304 let cost = Cost::of_pure_op(
305 inst_data.opcode(),
306 self.func.dfg.inst_values(inst).map(|value| best[value].0),
307 );
308 best[value] = BestEntry(cost, value);
309 trace!(" -> cost of value {} = {:?}", value, cost);
310 }
311 }
312 };
313
314 // Keep on iterating the fixpoint loop while we are finding new
315 // best values.
316 keep_going |= orig_best_value != best[value];
317 }
318 }
319
320 if cfg!(any(feature = "trace-log", debug_assertions)) {
321 trace!("finished fixpoint loop to compute best value for each eclass");
322 for value in self.func.dfg.values() {
323 trace!("-> best for eclass {:?}: {:?}", value, best[value]);
324 debug_assert_ne!(best[value].1, Value::reserved_value());
325 // You might additionally be expecting an assert that the best
326 // cost is not infinity, however infinite cost *can* happen in
327 // practice. First, note that our cost function doesn't know
328 // about any shared structure in the dataflow graph, it only
329 // sums operand costs. (And trying to avoid that by deduping a
330 // single operation's operands is a losing game because you can
331 // always just add one indirection and go from `add(x, x)` to
332 // `add(foo(x), bar(x))` to hide the shared structure.) Given
333 // that blindness to sharing, we can make cost grow
334 // exponentially with a linear sequence of operations:
335 //
336 // v0 = iconst.i32 1 ;; cost = 1
337 // v1 = iadd v0, v0 ;; cost = 3 + 1 + 1
338 // v2 = iadd v1, v1 ;; cost = 3 + 5 + 5
339 // v3 = iadd v2, v2 ;; cost = 3 + 13 + 13
340 // v4 = iadd v3, v3 ;; cost = 3 + 29 + 29
341 // v5 = iadd v4, v4 ;; cost = 3 + 61 + 61
342 // v6 = iadd v5, v5 ;; cost = 3 + 125 + 125
343 // ;; etc...
344 //
345 // Such a chain can cause cost to saturate to infinity. How do
346 // we choose which e-node is best when there are multiple that
347 // have saturated to infinity? It doesn't matter. As long as
348 // invariant (2) for optimization rules is upheld by our rule
349 // set (see `cranelift/codegen/src/opts/README.md`) it is safe
350 // to choose *any* e-node in the e-class. At worst we will
351 // produce suboptimal code, but never an incorrectness.
352 }
353 }
354 }
355
356 /// Elaborate use of an eclass, inserting any needed new
357 /// instructions before the given inst `before`. Should only be
358 /// given values corresponding to results of instructions or
359 /// blockparams.
360 fn elaborate_eclass_use(&mut self, value: Value, before: Inst) -> ElaboratedValue {
361 debug_assert_ne!(value, Value::reserved_value());
362
363 // Kick off the process by requesting this result
364 // value.
365 self.elab_stack
366 .push(ElabStackEntry::Start { value, before });
367
368 // Now run the explicit-stack recursion until we reach
369 // the root.
370 self.process_elab_stack();
371 debug_assert_eq!(self.elab_result_stack.len(), 1);
372 self.elab_result_stack.pop().unwrap()
373 }
374
375 /// Possibly rematerialize the instruction producing the value in
376 /// `arg` and rewrite `arg` to refer to it, if needed. Returns
377 /// `true` if a rewrite occurred.
378 fn maybe_remat_arg(
379 remat_values: &FxHashSet<Value>,
380 func: &mut Function,
381 remat_copies: &mut FxHashMap<(Block, Value), Value>,
382 insert_block: Block,
383 before: Inst,
384 arg: &mut ElaboratedValue,
385 stats: &mut Stats,
386 ) -> bool {
387 // TODO (#7313): we may want to consider recursive
388 // rematerialization as well. We could process the arguments of
389 // the rematerialized instruction up to a certain depth. This
390 // would affect, e.g., adds-with-one-constant-arg, which are
391 // currently rematerialized. Right now we don't do this, to
392 // avoid the need for another fixpoint loop here.
393 if arg.in_block != insert_block && remat_values.contains(&arg.value) {
394 let new_value = match remat_copies.entry((insert_block, arg.value)) {
395 HashEntry::Occupied(o) => *o.get(),
396 HashEntry::Vacant(v) => {
397 let inst = func.dfg.value_def(arg.value).inst().unwrap();
398 debug_assert_eq!(func.dfg.inst_results(inst).len(), 1);
399 let new_inst = func.dfg.clone_inst(inst);
400 func.layout.insert_inst(new_inst, before);
401 let new_result = func.dfg.inst_results(new_inst)[0];
402 *v.insert(new_result)
403 }
404 };
405 trace!("rematerialized {} as {}", arg.value, new_value);
406 arg.value = new_value;
407 stats.elaborate_remat += 1;
408 true
409 } else {
410 false
411 }
412 }
413
414 fn process_elab_stack(&mut self) {
415 while let Some(entry) = self.elab_stack.pop() {
416 match entry {
417 ElabStackEntry::Start { value, before } => {
418 debug_assert!(self.func.dfg.value_is_real(value));
419
420 self.stats.elaborate_visit_node += 1;
421
422 // Get the best option; we use `value` (latest
423 // value) here so we have a full view of the
424 // eclass.
425 trace!("looking up best value for {}", value);
426 let BestEntry(_, best_value) = self.value_to_best_value[value];
427 trace!("elaborate: value {} -> best {}", value, best_value);
428 debug_assert_ne!(best_value, Value::reserved_value());
429
430 if let Some(elab_val) =
431 self.value_to_elaborated_value.get(&NullCtx, &best_value)
432 {
433 // Value is available; use it.
434 trace!("elaborate: value {} -> {:?}", value, elab_val);
435 self.stats.elaborate_memoize_hit += 1;
436 self.elab_result_stack.push(*elab_val);
437 continue;
438 }
439
440 self.stats.elaborate_memoize_miss += 1;
441
442 // Now resolve the value to its definition to see
443 // how we can compute it.
444 let (inst, result_idx) = match self.func.dfg.value_def(best_value) {
445 ValueDef::Result(inst, result_idx) => {
446 trace!(
447 " -> value {} is result {} of {}",
448 best_value, result_idx, inst
449 );
450 (inst, result_idx)
451 }
452 ValueDef::Param(in_block, _) => {
453 // We don't need to do anything to compute
454 // this value; just push its result on the
455 // result stack (blockparams are already
456 // available).
457 trace!(" -> value {} is a blockparam", best_value);
458 self.elab_result_stack.push(ElaboratedValue {
459 in_block,
460 value: best_value,
461 });
462 continue;
463 }
464 ValueDef::Union(_, _) => {
465 panic!("Should never have a Union value as the best value");
466 }
467 };
468
469 trace!(
470 " -> result {} of inst {:?}",
471 result_idx, self.func.dfg.insts[inst]
472 );
473
474 // We're going to need to use this instruction
475 // result, placing the instruction into the
476 // layout. First, enqueue all args to be
477 // elaborated. Push state to receive the results
478 // and later elab this inst.
479 let num_args = self.func.dfg.inst_values(inst).count();
480 self.elab_stack.push(ElabStackEntry::PendingInst {
481 inst,
482 result_idx,
483 num_args,
484 before,
485 });
486
487 // Push args in reverse order so we process the
488 // first arg first.
489 for arg in self.func.dfg.inst_values(inst).rev() {
490 debug_assert_ne!(arg, Value::reserved_value());
491 self.elab_stack
492 .push(ElabStackEntry::Start { value: arg, before });
493 }
494 }
495
496 ElabStackEntry::PendingInst {
497 inst,
498 result_idx,
499 num_args,
500 before,
501 } => {
502 trace!(
503 "PendingInst: {} result {} args {} before {}",
504 inst, result_idx, num_args, before
505 );
506
507 // We should have all args resolved at this
508 // point. Grab them and drain them out, removing
509 // them.
510 let arg_idx = self.elab_result_stack.len() - num_args;
511 let arg_values = &mut self.elab_result_stack[arg_idx..];
512
513 // Compute max loop depth.
514 //
515 // Note that if there are no arguments then this instruction
516 // is allowed to get hoisted up one loop. This is not
517 // usually used since no-argument values are things like
518 // constants which are typically rematerialized, but for the
519 // `vconst` instruction 128-bit constants aren't as easily
520 // rematerialized. They're hoisted out of inner loops but
521 // not to the function entry which may run the risk of
522 // placing too much register pressure on the entire
523 // function. This is modeled with the `.saturating_sub(1)`
524 // as the default if there's otherwise no maximum.
525 let loop_hoist_level = arg_values
526 .iter()
527 .map(|&value| {
528 // Find the outermost loop level at which
529 // the value's defining block *is not* a
530 // member. This is the loop-nest level
531 // whose hoist-block we hoist to.
532 let hoist_level = self
533 .loop_stack
534 .iter()
535 .position(|loop_entry| {
536 !self.loop_analysis.is_in_loop(value.in_block, loop_entry.lp)
537 })
538 .unwrap_or(self.loop_stack.len());
539 trace!(
540 " -> arg: elab_value {:?} hoist level {:?}",
541 value, hoist_level
542 );
543 hoist_level
544 })
545 .max()
546 .unwrap_or(self.loop_stack.len().saturating_sub(1));
547 trace!(
548 " -> loop hoist level: {:?}; cur loop depth: {:?}, loop_stack: {:?}",
549 loop_hoist_level,
550 self.loop_stack.len(),
551 self.loop_stack,
552 );
553
554 // We know that this is a pure inst, because
555 // non-pure roots have already been placed in the
556 // value-to-elab'd-value map, so they will not
557 // reach this stage of processing.
558 //
559 // We now must determine the location at which we
560 // place the instruction. This is the current
561 // block *unless* we hoist above a loop when all
562 // args are loop-invariant (and this op is pure).
563 let (scope_depth, before, insert_block) =
564 if loop_hoist_level == self.loop_stack.len() {
565 // Depends on some value at the current
566 // loop depth, or remat forces it here:
567 // place it at the current location.
568 (
569 self.value_to_elaborated_value.depth(),
570 before,
571 self.func.layout.inst_block(before).unwrap(),
572 )
573 } else {
574 // Does not depend on any args at current
575 // loop depth: hoist out of loop.
576 self.stats.elaborate_licm_hoist += 1;
577 let data = &self.loop_stack[loop_hoist_level];
578 // `data.hoist_block` should dominate `before`'s block.
579 let before_block = self.func.layout.inst_block(before).unwrap();
580 debug_assert!(self.domtree.dominates(data.hoist_block, before_block));
581 // Determine the instruction at which we
582 // insert in `data.hoist_block`.
583 let before = self.func.layout.last_inst(data.hoist_block).unwrap();
584 (data.scope_depth as usize, before, data.hoist_block)
585 };
586
587 trace!(
588 " -> decided to place: before {} insert_block {}",
589 before, insert_block
590 );
591
592 // Now that we have the location for the
593 // instruction, check if any of its args are remat
594 // values. If so, and if we don't have a copy of
595 // the rematerializing instruction for this block
596 // yet, create one.
597 let mut remat_arg = false;
598 for arg_value in arg_values.iter_mut() {
599 if Self::maybe_remat_arg(
600 &self.remat_values,
601 &mut self.func,
602 &mut self.remat_copies,
603 insert_block,
604 before,
605 arg_value,
606 &mut self.stats,
607 ) {
608 remat_arg = true;
609 }
610 }
611
612 // Now we need to place `inst` at the computed
613 // location (just before `before`). Note that
614 // `inst` may already have been placed somewhere
615 // else, because a pure node may be elaborated at
616 // more than one place. In this case, we need to
617 // duplicate the instruction (and return the
618 // `Value`s for that duplicated instance instead).
619 //
620 // Also clone if we rematerialized, because we
621 // don't want to rewrite the args in the original
622 // copy.
623 trace!("need inst {} before {}", inst, before);
624 let inst = if self.func.layout.inst_block(inst).is_some() || remat_arg {
625 // Clone the inst!
626 let new_inst = self.func.dfg.clone_inst(inst);
627 trace!(
628 " -> inst {} already has a location; cloned to {}",
629 inst, new_inst
630 );
631 // Create mappings in the
632 // value-to-elab'd-value map from original
633 // results to cloned results.
634 for (&result, &new_result) in self
635 .func
636 .dfg
637 .inst_results(inst)
638 .iter()
639 .zip(self.func.dfg.inst_results(new_inst).iter())
640 {
641 let elab_value = ElaboratedValue {
642 value: new_result,
643 in_block: insert_block,
644 };
645 let best_result = self.value_to_best_value[result];
646 self.value_to_elaborated_value.insert_if_absent_with_depth(
647 &NullCtx,
648 best_result.1,
649 elab_value,
650 scope_depth,
651 );
652
653 self.value_to_best_value[new_result] = best_result;
654
655 trace!(
656 " -> cloned inst has new result {} for orig {}",
657 new_result, result
658 );
659 }
660 new_inst
661 } else {
662 trace!(" -> no location; using original inst");
663 // Create identity mappings from result values
664 // to themselves in this scope, since we're
665 // using the original inst.
666 for &result in self.func.dfg.inst_results(inst) {
667 let elab_value = ElaboratedValue {
668 value: result,
669 in_block: insert_block,
670 };
671 let best_result = self.value_to_best_value[result];
672 self.value_to_elaborated_value.insert_if_absent_with_depth(
673 &NullCtx,
674 best_result.1,
675 elab_value,
676 scope_depth,
677 );
678 trace!(" -> inserting identity mapping for {}", result);
679 }
680 inst
681 };
682
683 // Place the inst just before `before`.
684 assert!(
685 is_pure_for_egraph(self.func, inst),
686 "something has gone very wrong if we are elaborating effectful \
687 instructions, they should have remained in the skeleton"
688 );
689 self.func.layout.insert_inst(inst, before);
690
691 // Update the inst's arguments.
692 self.func
693 .dfg
694 .overwrite_inst_values(inst, arg_values.into_iter().map(|ev| ev.value));
695
696 // Now that we've consumed the arg values, pop
697 // them off the stack.
698 self.elab_result_stack.truncate(arg_idx);
699
700 // Push the requested result index of the
701 // instruction onto the elab-results stack.
702 self.elab_result_stack.push(ElaboratedValue {
703 in_block: insert_block,
704 value: self.func.dfg.inst_results(inst)[result_idx],
705 });
706 }
707 }
708 }
709 }
710
711 fn elaborate_block(&mut self, elab_values: &mut Vec<Value>, idom: Option<Block>, block: Block) {
712 trace!("elaborate_block: block {}", block);
713 self.start_block(idom, block);
714
715 // Iterate over the side-effecting skeleton using the linked
716 // list in Layout. We will insert instructions that are
717 // elaborated *before* `inst`, so we can always use its
718 // next-link to continue the iteration.
719 let mut next_inst = self.func.layout.first_inst(block);
720 let mut first_branch = None;
721 while let Some(inst) = next_inst {
722 trace!(
723 "elaborating inst {} with results {:?}",
724 inst,
725 self.func.dfg.inst_results(inst)
726 );
727 // Record the first branch we see in the block; all
728 // elaboration for args of *any* branch must be inserted
729 // before the *first* branch, because the branch group
730 // must remain contiguous at the end of the block.
731 if self.func.dfg.insts[inst].opcode().is_branch() && first_branch == None {
732 first_branch = Some(inst);
733 }
734
735 // Determine where elaboration inserts insts.
736 let before = first_branch.unwrap_or(inst);
737 trace!(" -> inserting before {}", before);
738
739 elab_values.extend(self.func.dfg.inst_values(inst));
740 for arg in elab_values.iter_mut() {
741 trace!(" -> arg {}", *arg);
742 // Elaborate the arg, placing any newly-inserted insts
743 // before `before`. Get the updated value, which may
744 // be different than the original.
745 let mut new_arg = self.elaborate_eclass_use(*arg, before);
746 Self::maybe_remat_arg(
747 &self.remat_values,
748 &mut self.func,
749 &mut self.remat_copies,
750 block,
751 inst,
752 &mut new_arg,
753 &mut self.stats,
754 );
755 trace!(" -> rewrote arg to {:?}", new_arg);
756 *arg = new_arg.value;
757 }
758 self.func
759 .dfg
760 .overwrite_inst_values(inst, elab_values.drain(..));
761
762 // We need to put the results of this instruction in the
763 // map now.
764 for &result in self.func.dfg.inst_results(inst) {
765 trace!(" -> result {}", result);
766 let best_result = self.value_to_best_value[result];
767 self.value_to_elaborated_value.insert_if_absent(
768 &NullCtx,
769 best_result.1,
770 ElaboratedValue {
771 in_block: block,
772 value: result,
773 },
774 );
775 }
776
777 next_inst = self.func.layout.next_inst(inst);
778 }
779 }
780
781 fn elaborate_domtree(&mut self, domtree: &DominatorTreePreorder) {
782 self.block_stack.push(BlockStackEntry::Elaborate {
783 block: self.func.layout.entry_block().unwrap(),
784 idom: None,
785 });
786
787 // A temporary workspace for elaborate_block, allocated here to maximize the use of the
788 // allocation.
789 let mut elab_values = Vec::new();
790
791 while let Some(top) = self.block_stack.pop() {
792 match top {
793 BlockStackEntry::Elaborate { block, idom } => {
794 self.block_stack.push(BlockStackEntry::Pop);
795 self.value_to_elaborated_value.increment_depth();
796
797 self.elaborate_block(&mut elab_values, idom, block);
798
799 // Push children. We are doing a preorder
800 // traversal so we do this after processing this
801 // block above.
802 let block_stack_end = self.block_stack.len();
803 for child in self.ctrl_plane.shuffled(domtree.children(block)) {
804 self.block_stack.push(BlockStackEntry::Elaborate {
805 block: child,
806 idom: Some(block),
807 });
808 }
809 // Reverse what we just pushed so we elaborate in
810 // original block order. (The domtree iter is a
811 // single-ended iter over a singly-linked list so
812 // we can't `.rev()` above.)
813 self.block_stack[block_stack_end..].reverse();
814 }
815 BlockStackEntry::Pop => {
816 self.value_to_elaborated_value.decrement_depth();
817 }
818 }
819 }
820 }
821
822 pub(crate) fn elaborate(&mut self) {
823 self.stats.elaborate_func += 1;
824 self.stats.elaborate_func_pre_insts += self.func.dfg.num_insts() as u64;
825 self.compute_best_values();
826 self.elaborate_domtree(&self.domtree);
827 self.stats.elaborate_func_post_insts += self.func.dfg.num_insts() as u64;
828 }
829}