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