Skip to main content

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}