| //! Egraph-based mid-end optimization framework. |
| |
| use crate::dominator_tree::DominatorTree; |
| use crate::egraph::stores::PackedMemoryState; |
| use crate::flowgraph::ControlFlowGraph; |
| use crate::loop_analysis::{LoopAnalysis, LoopLevel}; |
| use crate::trace; |
| use crate::{ |
| fx::{FxHashMap, FxHashSet}, |
| inst_predicates::has_side_effect, |
| ir::{Block, Function, Inst, InstructionData, InstructionImms, Opcode, Type}, |
| }; |
| use alloc::vec::Vec; |
| use core::ops::Range; |
| use cranelift_egraph::{EGraph, Id, Language, NewOrExisting}; |
| use cranelift_entity::EntityList; |
| use cranelift_entity::SecondaryMap; |
| |
| mod domtree; |
| mod elaborate; |
| mod node; |
| mod stores; |
| |
| use elaborate::Elaborator; |
| pub use node::{Node, NodeCtx}; |
| pub use stores::{AliasAnalysis, MemoryState}; |
| |
| pub struct FuncEGraph<'a> { |
| /// Dominator tree, used for elaboration pass. |
| domtree: &'a DominatorTree, |
| /// Loop analysis results, used for built-in LICM during elaboration. |
| loop_analysis: &'a LoopAnalysis, |
| /// Last-store tracker for integrated alias analysis during egraph build. |
| alias_analysis: AliasAnalysis, |
| /// The egraph itself. |
| pub(crate) egraph: EGraph<NodeCtx, Analysis>, |
| /// "node context", containing arenas for node data. |
| pub(crate) node_ctx: NodeCtx, |
| /// Ranges in `side_effect_ids` for sequences of side-effecting |
| /// eclasses per block. |
| side_effects: SecondaryMap<Block, Range<u32>>, |
| side_effect_ids: Vec<Id>, |
| /// Map from store instructions to their nodes; used for store-to-load forwarding. |
| pub(crate) store_nodes: FxHashMap<Inst, (Type, Id)>, |
| /// Ranges in `blockparam_ids_tys` for sequences of blockparam |
| /// eclass IDs and types per block. |
| blockparams: SecondaryMap<Block, Range<u32>>, |
| blockparam_ids_tys: Vec<(Id, Type)>, |
| /// Which canonical node IDs do we want to rematerialize in each |
| /// block where they're used? |
| pub(crate) remat_ids: FxHashSet<Id>, |
| /// Which canonical node IDs have an enode whose value subsumes |
| /// all others it's unioned with? |
| pub(crate) subsume_ids: FxHashSet<Id>, |
| /// Statistics recorded during the process of building, |
| /// optimizing, and lowering out of this egraph. |
| pub(crate) stats: Stats, |
| /// Current rewrite-recursion depth. Used to enforce a finite |
| /// limit on rewrite rule application so that we don't get stuck |
| /// in an infinite chain. |
| pub(crate) rewrite_depth: usize, |
| } |
| |
| #[derive(Clone, Debug, Default)] |
| pub(crate) struct Stats { |
| pub(crate) node_created: u64, |
| pub(crate) node_param: u64, |
| pub(crate) node_result: u64, |
| pub(crate) node_pure: u64, |
| pub(crate) node_inst: u64, |
| pub(crate) node_load: u64, |
| pub(crate) node_dedup_query: u64, |
| pub(crate) node_dedup_hit: u64, |
| pub(crate) node_dedup_miss: u64, |
| pub(crate) node_ctor_created: u64, |
| pub(crate) node_ctor_deduped: u64, |
| pub(crate) node_union: u64, |
| pub(crate) node_subsume: u64, |
| pub(crate) store_map_insert: u64, |
| pub(crate) side_effect_nodes: u64, |
| pub(crate) rewrite_rule_invoked: u64, |
| pub(crate) rewrite_depth_limit: u64, |
| pub(crate) store_to_load_forward: u64, |
| pub(crate) elaborate_visit_node: u64, |
| pub(crate) elaborate_memoize_hit: u64, |
| pub(crate) elaborate_memoize_miss: u64, |
| pub(crate) elaborate_memoize_miss_remat: u64, |
| pub(crate) elaborate_licm_hoist: u64, |
| pub(crate) elaborate_func: u64, |
| pub(crate) elaborate_func_pre_insts: u64, |
| pub(crate) elaborate_func_post_insts: u64, |
| } |
| |
| impl<'a> FuncEGraph<'a> { |
| /// Create a new EGraph for the given function. Requires the |
| /// domtree to be precomputed as well; the domtree is used for |
| /// scheduling when lowering out of the egraph. |
| pub fn new( |
| func: &Function, |
| domtree: &'a DominatorTree, |
| loop_analysis: &'a LoopAnalysis, |
| cfg: &ControlFlowGraph, |
| ) -> FuncEGraph<'a> { |
| let num_values = func.dfg.num_values(); |
| let num_blocks = func.dfg.num_blocks(); |
| let node_count_estimate = num_values * 2; |
| let alias_analysis = AliasAnalysis::new(func, cfg); |
| let mut this = Self { |
| domtree, |
| loop_analysis, |
| alias_analysis, |
| egraph: EGraph::with_capacity(node_count_estimate, Some(Analysis)), |
| node_ctx: NodeCtx::with_capacity_for_dfg(&func.dfg), |
| side_effects: SecondaryMap::with_capacity(num_blocks), |
| side_effect_ids: Vec::with_capacity(node_count_estimate), |
| store_nodes: FxHashMap::default(), |
| blockparams: SecondaryMap::with_capacity(num_blocks), |
| blockparam_ids_tys: Vec::with_capacity(num_blocks * 10), |
| remat_ids: FxHashSet::default(), |
| subsume_ids: FxHashSet::default(), |
| stats: Default::default(), |
| rewrite_depth: 0, |
| }; |
| this.store_nodes.reserve(func.dfg.num_values() / 8); |
| this.remat_ids.reserve(func.dfg.num_values() / 4); |
| this.subsume_ids.reserve(func.dfg.num_values() / 4); |
| this.build(func); |
| this |
| } |
| |
| fn build(&mut self, func: &Function) { |
| // Mapping of SSA `Value` to eclass ID. |
| let mut value_to_id = FxHashMap::default(); |
| |
| // For each block in RPO, create an enode for block entry, for |
| // each block param, and for each instruction. |
| for &block in self.domtree.cfg_postorder().iter().rev() { |
| let loop_level = self.loop_analysis.loop_level(block); |
| let blockparam_start = |
| u32::try_from(self.blockparam_ids_tys.len()).expect("Overflow in blockparam count"); |
| for (i, &value) in func.dfg.block_params(block).iter().enumerate() { |
| let ty = func.dfg.value_type(value); |
| let param = self |
| .egraph |
| .add( |
| Node::Param { |
| block, |
| index: i |
| .try_into() |
| .expect("blockparam index should fit in Node::Param"), |
| ty, |
| loop_level, |
| }, |
| &mut self.node_ctx, |
| ) |
| .get(); |
| value_to_id.insert(value, param); |
| self.blockparam_ids_tys.push((param, ty)); |
| self.stats.node_created += 1; |
| self.stats.node_param += 1; |
| } |
| let blockparam_end = |
| u32::try_from(self.blockparam_ids_tys.len()).expect("Overflow in blockparam count"); |
| self.blockparams[block] = blockparam_start..blockparam_end; |
| |
| let side_effect_start = |
| u32::try_from(self.side_effect_ids.len()).expect("Overflow in side-effect count"); |
| for inst in func.layout.block_insts(block) { |
| // Build args from SSA values. |
| let args = EntityList::from_iter( |
| func.dfg.inst_args(inst).iter().map(|&arg| { |
| let arg = func.dfg.resolve_aliases(arg); |
| *value_to_id |
| .get(&arg) |
| .expect("Must have seen def before this use") |
| }), |
| &mut self.node_ctx.args, |
| ); |
| |
| let results = func.dfg.inst_results(inst); |
| let ty = if results.len() == 1 { |
| func.dfg.value_type(results[0]) |
| } else { |
| crate::ir::types::INVALID |
| }; |
| |
| let load_mem_state = self.alias_analysis.get_state_for_load(inst); |
| let is_readonly_load = match func.dfg[inst] { |
| InstructionData::Load { |
| opcode: Opcode::Load, |
| flags, |
| .. |
| } => flags.readonly() && flags.notrap(), |
| _ => false, |
| }; |
| |
| // Create the egraph node. |
| let op = InstructionImms::from(&func.dfg[inst]); |
| let opcode = op.opcode(); |
| let srcloc = func.srclocs[inst]; |
| let arity = u16::try_from(results.len()) |
| .expect("More than 2^16 results from an instruction"); |
| |
| let node = if is_readonly_load { |
| self.stats.node_created += 1; |
| self.stats.node_pure += 1; |
| Node::Pure { |
| op, |
| args, |
| ty, |
| arity, |
| } |
| } else if let Some(load_mem_state) = load_mem_state { |
| let addr = args.as_slice(&self.node_ctx.args)[0]; |
| trace!("load at inst {} has mem state {:?}", inst, load_mem_state); |
| self.stats.node_created += 1; |
| self.stats.node_load += 1; |
| Node::Load { |
| op, |
| ty, |
| addr, |
| mem_state: load_mem_state, |
| srcloc, |
| } |
| } else if has_side_effect(func, inst) || opcode.can_load() { |
| self.stats.node_created += 1; |
| self.stats.node_inst += 1; |
| Node::Inst { |
| op, |
| args, |
| ty, |
| arity, |
| srcloc, |
| loop_level, |
| } |
| } else { |
| self.stats.node_created += 1; |
| self.stats.node_pure += 1; |
| Node::Pure { |
| op, |
| args, |
| ty, |
| arity, |
| } |
| }; |
| let dedup_needed = self.node_ctx.needs_dedup(&node); |
| let is_pure = matches!(node, Node::Pure { .. }); |
| |
| let mut id = self.egraph.add(node, &mut self.node_ctx); |
| |
| if dedup_needed { |
| self.stats.node_dedup_query += 1; |
| match id { |
| NewOrExisting::New(_) => { |
| self.stats.node_dedup_miss += 1; |
| } |
| NewOrExisting::Existing(_) => { |
| self.stats.node_dedup_hit += 1; |
| } |
| } |
| } |
| |
| if opcode == Opcode::Store { |
| let store_data_ty = func.dfg.value_type(func.dfg.inst_args(inst)[0]); |
| self.store_nodes.insert(inst, (store_data_ty, id.get())); |
| self.stats.store_map_insert += 1; |
| } |
| |
| // Loads that did not already merge into an existing |
| // load: try to forward from a store (store-to-load |
| // forwarding). |
| if let NewOrExisting::New(new_id) = id { |
| if load_mem_state.is_some() { |
| let opt_id = crate::opts::store_to_load(new_id, self); |
| trace!("store_to_load: {} -> {}", new_id, opt_id); |
| if opt_id != new_id { |
| id = NewOrExisting::Existing(opt_id); |
| } |
| } |
| } |
| |
| // Now either optimize (for new pure nodes), or add to |
| // the side-effecting list (for all other new nodes). |
| let id = match id { |
| NewOrExisting::Existing(id) => id, |
| NewOrExisting::New(id) if is_pure => { |
| // Apply all optimization rules immediately; the |
| // aegraph (acyclic egraph) works best when we do |
| // this so all uses pick up the eclass with all |
| // possible enodes. |
| crate::opts::optimize_eclass(id, self) |
| } |
| NewOrExisting::New(id) => { |
| self.side_effect_ids.push(id); |
| self.stats.side_effect_nodes += 1; |
| id |
| } |
| }; |
| |
| // Create results and save in Value->Id map. |
| match results { |
| &[] => {} |
| &[one_result] => { |
| trace!("build: value {} -> id {}", one_result, id); |
| value_to_id.insert(one_result, id); |
| } |
| many_results => { |
| debug_assert!(many_results.len() > 1); |
| for (i, &result) in many_results.iter().enumerate() { |
| let ty = func.dfg.value_type(result); |
| let projection = self |
| .egraph |
| .add( |
| Node::Result { |
| value: id, |
| result: i, |
| ty, |
| }, |
| &mut self.node_ctx, |
| ) |
| .get(); |
| self.stats.node_created += 1; |
| self.stats.node_result += 1; |
| trace!("build: value {} -> id {}", result, projection); |
| value_to_id.insert(result, projection); |
| } |
| } |
| } |
| } |
| |
| let side_effect_end = |
| u32::try_from(self.side_effect_ids.len()).expect("Overflow in side-effect count"); |
| let side_effect_range = side_effect_start..side_effect_end; |
| self.side_effects[block] = side_effect_range; |
| } |
| } |
| |
| /// Scoped elaboration: compute a final ordering of op computation |
| /// for each block and replace the given Func body. |
| /// |
| /// This works in concert with the domtree. We do a preorder |
| /// traversal of the domtree, tracking a scoped map from Id to |
| /// (new) Value. The map's scopes correspond to levels in the |
| /// domtree. |
| /// |
| /// At each block, we iterate forward over the side-effecting |
| /// eclasses, and recursively generate their arg eclasses, then |
| /// emit the ops themselves. |
| /// |
| /// To use an eclass in a given block, we first look it up in the |
| /// scoped map, and get the Value if already present. If not, we |
| /// need to generate it. We emit the extracted enode for this |
| /// eclass after recursively generating its args. Eclasses are |
| /// thus computed "as late as possible", but then memoized into |
| /// the Id-to-Value map and available to all dominated blocks and |
| /// for the rest of this block. (This subsumes GVN.) |
| pub fn elaborate(&mut self, func: &mut Function) { |
| let mut elab = Elaborator::new( |
| func, |
| self.domtree, |
| self.loop_analysis, |
| &self.egraph, |
| &self.node_ctx, |
| &self.remat_ids, |
| &mut self.stats, |
| ); |
| elab.elaborate( |
| |block| { |
| let blockparam_range = self.blockparams[block].clone(); |
| &self.blockparam_ids_tys |
| [blockparam_range.start as usize..blockparam_range.end as usize] |
| }, |
| |block| { |
| let side_effect_range = self.side_effects[block].clone(); |
| &self.side_effect_ids |
| [side_effect_range.start as usize..side_effect_range.end as usize] |
| }, |
| ); |
| } |
| } |
| |
| /// State for egraph analysis that computes all needed properties. |
| pub(crate) struct Analysis; |
| |
| /// Analysis results for each eclass id. |
| #[derive(Clone, Debug)] |
| pub(crate) struct AnalysisValue { |
| pub(crate) loop_level: LoopLevel, |
| } |
| |
| impl Default for AnalysisValue { |
| fn default() -> Self { |
| Self { |
| loop_level: LoopLevel::root(), |
| } |
| } |
| } |
| |
| impl cranelift_egraph::Analysis for Analysis { |
| type L = NodeCtx; |
| type Value = AnalysisValue; |
| |
| fn for_node( |
| &self, |
| ctx: &NodeCtx, |
| n: &Node, |
| values: &SecondaryMap<Id, AnalysisValue>, |
| ) -> AnalysisValue { |
| let loop_level = match n { |
| &Node::Pure { ref args, .. } => args |
| .as_slice(&ctx.args) |
| .iter() |
| .map(|&arg| values[arg].loop_level) |
| .max() |
| .unwrap_or(LoopLevel::root()), |
| &Node::Load { addr, .. } => values[addr].loop_level, |
| &Node::Result { value, .. } => values[value].loop_level, |
| &Node::Inst { loop_level, .. } | &Node::Param { loop_level, .. } => loop_level, |
| }; |
| |
| AnalysisValue { loop_level } |
| } |
| |
| fn meet(&self, _ctx: &NodeCtx, v1: &AnalysisValue, v2: &AnalysisValue) -> AnalysisValue { |
| AnalysisValue { |
| loop_level: std::cmp::max(v1.loop_level, v2.loop_level), |
| } |
| } |
| } |
