Rework cranelift-egraph to use more arenas.

This moves to a strategy based on `hashbrown::raw::RawTable` and
"eq-with-context" / "hash-with-context" traits, used to allow nodes to
be stored once in a `BumpVec` (which encapsulates a range in a shared
`Vec`) and then used in an eclass, with keys in the deduplication
hashtable referring to an eclass and enode index in that eclass.

This moves back to the enodes-without-IDs strategy used in `egg`, which
makes the graph rebuild simpler, but retains the single-storage /
no-cloning property.

Along the way, carrying through the eq-with-context / hash-with-context
to the `Node` itself, we can remove the `&mut [Id]` and `&mut [Type]`
slices and replace with `BumpVec`s. Actually they could be made
`NonGrowingBumpVec`s for 8 bytes instead of 12; that is future work.

This should allow equivalent algorithmic approaches to `egg` but without
the separate allocations; all allocations for nodes are now within the
entity-component-system-style large `Vec`s.

Author: cfallin

Cargo.lock

@@ -26,7 +26,6 @@ gimli = { version = "0.26.0", default-features = false, features = ["write"], op
 smallvec = { version = "1.6.1" }
 regalloc2 = { version = "0.2.0", features = ["checker"] }
 souper-ir = { version = "2.1.0", optional = true }
-bumpalo = { version = "3.10.0" }
 # It is a goal of the cranelift-codegen crate to have minimal external dependencies.
 # Please don't add any unless they are essential to the task of creating binary
 # machine code. Integration tests that need external dependencies can be

cranelift/codegen/Cargo.toml

@@ -33,7 +33,6 @@ use crate::verifier::{verify_context, VerifierErrors, VerifierResult};
 #[cfg(feature = "souper-harvest")]
 use alloc::string::String;
 use alloc::vec::Vec;
-use bumpalo::Bump;
 
 #[cfg(feature = "souper-harvest")]
 use crate::souper_harvest::do_souper_harvest;
@@ -195,8 +194,7 @@ impl Context {
                 "About to optimize with egraph phase:\n{}",
                 self.func.display()
             );
-            let bump = Bump::new();
-            let mut eg = FuncEGraph::new(&self.func, &self.domtree, &bump);
+            let mut eg = FuncEGraph::new(&self.func, &self.domtree);
             eg.extract(&self.func);
             eg.elaborate(&mut self.func);
             log::debug!("After egraph optimization:\n{}", self.func.display());

cranelift/codegen/src/context.rs

@@ -7,9 +7,9 @@ use crate::{
     ir::{Block, Function, InstructionImms, Type},
 };
 use alloc::vec::Vec;
-use bumpalo::Bump;
 use core::ops::Range;
 use cranelift_egraph::{EGraph, Id};
+use cranelift_entity::EntityList;
 use cranelift_entity::SecondaryMap;
 
 mod domtree;
@@ -19,14 +19,13 @@ mod node;
 
 use elaborate::Elaborator;
 use extract::Extractor;
-use node::Node;
+use node::{Node, NodeCtx};
 
 pub struct FuncEGraph<'a> {
     domtree: &'a DominatorTree,
-    egraph: EGraph<Node<'a>>,
-    /// Backing arena for eclass argument arrays and type arrays for
-    /// nodes.
-    bump: &'a Bump,
+    egraph: EGraph<NodeCtx>,
+    /// "node context", containing arenas for node data.
+    node_ctx: NodeCtx,
     /// Ranges in `side_effect_ids` for sequences of side-effecting
     /// eclasses per block.
     side_effects: SecondaryMap<Block, Range<u32>>,
@@ -43,11 +42,11 @@ 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, bump: &'a Bump) -> FuncEGraph<'a> {
+    pub fn new(func: &Function, domtree: &'a DominatorTree) -> FuncEGraph<'a> {
         let mut this = Self {
             domtree,
             egraph: EGraph::new(),
-            bump,
+            node_ctx: NodeCtx::default(),
             side_effects: SecondaryMap::with_default(0..0),
             side_effect_ids: vec![],
             blockparams: SecondaryMap::with_default(0..0),
@@ -68,11 +67,14 @@ impl<'a> FuncEGraph<'a> {
             let blockparam_start = self.blockparam_ids_tys.len() as u32;
             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 as u32,
-                    ty,
-                });
+                let param = self.egraph.add(
+                    Node::Param {
+                        block,
+                        index: i as u32,
+                        ty,
+                    },
+                    &mut self.node_ctx,
+                );
                 value_to_id.insert(value, param);
                 self.blockparam_ids_tys.push((param, ty));
             }
@@ -86,20 +88,22 @@ impl<'a> FuncEGraph<'a> {
                     || func.dfg[inst].opcode().can_store();
 
                 // Build args from SSA values.
-                let args = self
-                    .bump
-                    .alloc_slice_fill_iter(func.dfg.inst_args(inst).iter().map(|&arg| {
+                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 types = self
-                    .bump
-                    .alloc_slice_fill_iter(results.iter().map(|&val| func.dfg.value_type(val)));
+                    .node_ctx
+                    .types
+                    .from_iter(results.iter().map(|&val| func.dfg.value_type(val)));
 
                 // Create the egraph node.
                 let op = InstructionImms::from(&func.dfg[inst]);
@@ -115,7 +119,7 @@ impl<'a> FuncEGraph<'a> {
                 } else {
                     Node::Pure { op, args, types }
                 };
-                let id = self.egraph.add(node);
+                let id = self.egraph.add(node, &mut self.node_ctx);
 
                 if side_effect {
                     self.side_effect_ids.push(id);
@@ -131,11 +135,14 @@ impl<'a> FuncEGraph<'a> {
                         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,
-                            });
+                            let projection = self.egraph.add(
+                                Node::Result {
+                                    value: id,
+                                    result: i,
+                                    ty,
+                                },
+                                &mut self.node_ctx,
+                            );
                             value_to_id.insert(result, projection);
                         }
                     }
@@ -214,7 +221,7 @@ impl<'a> FuncEGraph<'a> {
                 let side_effect = self.side_effect_ids[side_effect as usize];
                 let present = self
                     .extractor
-                    .visit_eclass(&self.egraph, side_effect)
+                    .visit_eclass(&self.egraph, side_effect, &self.node_ctx)
                     .is_some();
                 debug_assert!(present);
             }
@@ -241,7 +248,13 @@ impl<'a> FuncEGraph<'a> {
     /// 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.egraph, &self.extractor);
+        let mut elab = Elaborator::new(
+            func,
+            self.domtree,
+            &self.egraph,
+            &self.node_ctx,
+            &self.extractor,
+        );
         elab.elaborate(
             |block| {
                 let blockparam_range = self.blockparams[block].clone();

cranelift/codegen/src/egg.rs

@@ -3,7 +3,7 @@
 
 use super::domtree::DomTreeWithChildren;
 use super::extract::Extractor;
-use super::node::Node;
+use super::node::{Node, NodeCtx};
 use crate::dominator_tree::DominatorTree;
 use crate::ir::{Block, Function, SourceLoc, Type, Value, ValueList};
 use crate::scoped_hash_map::ScopedHashMap;
@@ -13,7 +13,8 @@ use smallvec::SmallVec;
 pub(crate) struct Elaborator<'a> {
     func: &'a mut Function,
     domtree: &'a DominatorTree,
-    egraph: &'a EGraph<Node<'a>>,
+    node_ctx: &'a NodeCtx,
+    egraph: &'a EGraph<NodeCtx>,
     extractor: &'a Extractor,
     id_to_value: ScopedHashMap<Id, IdValue>,
     cur_block: Option<Block>,
@@ -31,13 +32,15 @@ impl<'a> Elaborator<'a> {
     pub(crate) fn new(
         func: &'a mut Function,
         domtree: &'a DominatorTree,
-        egraph: &'a EGraph<Node>,
+        egraph: &'a EGraph<NodeCtx>,
+        node_ctx: &'a NodeCtx,
         extractor: &'a Extractor,
     ) -> Self {
         Self {
             func,
             domtree,
             egraph,
+            node_ctx,
             extractor,
             id_to_value: ScopedHashMap::new(),
             cur_block: None,
@@ -69,7 +72,7 @@ impl<'a> Elaborator<'a> {
         };
         let inst = self.func.dfg.make_inst(instdata);
         self.func.srclocs[inst] = srcloc;
-        for &ty in result_tys.iter() {
+        for &ty in result_tys.as_slice(&self.node_ctx.types) {
             self.func.dfg.append_result(inst, ty);
         }
         self.func.layout.append_inst(inst, self.cur_block.unwrap());
@@ -111,8 +114,9 @@ impl<'a> Elaborator<'a> {
 
         // We're going to need to emit this operator. First, elaborate
         // all args, recursively.
-        let args: SmallVec<[Value; 8]> = node
-            .children()
+        let args: SmallVec<[Value; 8]> = self
+            .node_ctx
+            .children(node)
             .iter()
             .map(|&id| self.elaborate_eclass_use(id))
             .map(|idvalue| match idvalue {

cranelift/codegen/src/egg/elaborate.rs

@@ -1,13 +1,13 @@
 //! Extraction phase: pick one enode per eclass, avoiding loops.
 
-use super::node::Node;
+use super::node::{Node, NodeCtx};
 use crate::fx::FxHashMap;
-use cranelift_egraph::{EGraph, Id, Language, NodeId};
+use cranelift_egraph::{EGraph, Id, Language};
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 enum EclassState {
     Visiting,
-    Visited { cost: u32, node: NodeId },
+    Visited { cost: u32, node_idx: u32 },
     Deleted,
 }
 
@@ -25,7 +25,12 @@ impl Extractor {
 
     /// Visit an eclass. Return `None` if deleted, or Some(cost) if
     /// present.
-    pub(crate) fn visit_eclass(&mut self, egraph: &EGraph<Node>, id: Id) -> Option<u32> {
+    pub(crate) fn visit_eclass(
+        &mut self,
+        egraph: &EGraph<NodeCtx>,
+        id: Id,
+        ctx: &NodeCtx,
+    ) -> Option<u32> {
         if let Some(state) = self.eclass_state.get(&id) {
             match state {
                 EclassState::Visiting => {
@@ -43,25 +48,20 @@ impl Extractor {
         self.eclass_state.insert(id, EclassState::Visiting);
 
         let mut best_cost_and_node = None;
-        for (node_id, node) in egraph.enodes(id) {
-            let this_cost = self.visit_enode(egraph, node);
+        for (node_idx, node) in egraph.enodes(id).iter().enumerate() {
+            let this_cost = self.visit_enode(egraph, node, ctx);
             best_cost_and_node = match (best_cost_and_node, this_cost) {
                 (None, None) => None,
-                (None, Some(c)) => Some((c, node_id)),
-                (Some((c1, _)), Some(c2)) if c2 < c1 => Some((c2, node_id)),
-                (Some((c1, node_id1)), _) => Some((c1, node_id1)),
+                (None, Some(c)) => Some((c, node_idx as u32)),
+                (Some((c1, _)), Some(c2)) if c2 < c1 => Some((c2, node_idx as u32)),
+                (Some((c1, node_idx1)), _) => Some((c1, node_idx1)),
             };
         }
 
         match best_cost_and_node {
-            Some((cost, node_id)) => {
-                self.eclass_state.insert(
-                    id,
-                    EclassState::Visited {
-                        cost,
-                        node: node_id,
-                    },
-                );
+            Some((cost, node_idx)) => {
+                self.eclass_state
+                    .insert(id, EclassState::Visited { cost, node_idx });
                 Some(cost)
             }
             None => {
@@ -71,19 +71,19 @@ impl Extractor {
         }
     }
 
-    fn visit_enode(&mut self, egraph: &EGraph<Node<'_>>, node: &Node) -> Option<u32> {
+    fn visit_enode(&mut self, egraph: &EGraph<NodeCtx>, node: &Node, ctx: &NodeCtx) -> Option<u32> {
         let mut cost = node.cost() as u32;
-        for &arg in node.children() {
-            let arg_cost = self.visit_eclass(egraph, arg)?;
+        for &arg in ctx.children(node) {
+            let arg_cost = self.visit_eclass(egraph, arg, ctx)?;
             cost += arg_cost;
         }
         Some(cost)
     }
 
-    pub(crate) fn get_node<'a>(&'a self, egraph: &'a EGraph<Node<'a>>, id: Id) -> Option<&'a Node> {
+    pub(crate) fn get_node<'a>(&'a self, egraph: &'a EGraph<NodeCtx>, id: Id) -> Option<&'a Node> {
         match self.eclass_state.get(&id)? {
             &EclassState::Visiting => unreachable!(),
-            &EclassState::Visited { node, .. } => Some(egraph.enode(node)),
+            &EclassState::Visited { node_idx, .. } => Some(&egraph.enodes(id)[node_idx as usize]),
             &EclassState::Deleted => None,
         }
     }