cranelift egraphs

[pca006132]

Just opening this to discuss how we want to add it. I thought I can do it in an afternoon, but apparently there are some issues here...

1. Just add the patch to cranelift in this repo, or do we want to upstream the code that serialize the egraph?
2. How should we deal with root e-class? The main issue is that in cranelift, they don't extract a single thing, but they compute the best extraction cost for each e-class (yeah, they don't care about sharing...).

If someone wants to have a look, here is their code: https://github.com/bytecodealliance/wasmtime/blob/8ea25906f561d6d2e16e6152963f3a87af5a4e46/cranelift/codegen/src/egraph/elaborate.rs#L221

There are some minor annoyance, e.g. they use union nodes instead of e-classes, and e-nodes may belong to one union node or no union node. Haven't yet checked if they can belong to multiple union nodes or something else can reference an e-node with a union node without going through that union node. But this can be canonicalized easily.

[pca006132]

@mwillsey ping

[mwillsey]

Sorry for the delay!

1. No, let's keep the pathced cranelift somewhere else. Just link it in the PR. Just adding the serialized e-graphs here is good enough!
2. The serialization format supports multiple roots! Surely we don't care about every e-class, since many are generated dynamically. Could the patched cranelift record the root e-classes (on insertion time maybe)?

[pca006132]

Was busy previously, I just implemented this. The issue is that there are 26k files, i.e. 1 for each function for each benchmark in sightglass, so I am not sure whether uploading those files directly is a good idea. Maybe I can upload a zip file? Once this is decided I will make a PR.

Regarding root e-classes, my current approach is to extract all the nodes with no parent in the egraph, because the graph is acyclic in cranelift we can do that. The issue is that due to the nature of those graphs (a lot of values in the function can have no parent), there can be a lot of root eclasses...

Here is the (really ugly) patch to wasmtime:

diff --git a/cranelift/codegen/Cargo.toml b/cranelift/codegen/Cargo.toml
index 05a54e8a6..bf113e6b1 100644
--- a/cranelift/codegen/Cargo.toml
+++ b/cranelift/codegen/Cargo.toml
@@ -42,6 +42,11 @@ regalloc2 = { workspace = true, features = ["checker"] }
 souper-ir = { version = "2.1.0", optional = true }
 sha2 = { version = "0.10.2", optional = true }
 rustc-hash  = { workspace = true }
+
+[dependencies.egraph-serialize]
+git = "https://github.com/egraphs-good/egraph-serialize"
+rev = "951b829a434f4008c7b45ba4ac0da1037d2da90"
+
 # 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
diff --git a/cranelift/codegen/src/egraph/cost.rs b/cranelift/codegen/src/egraph/cost.rs
index 28aab40e9..a1d12ac3e 100644
--- a/cranelift/codegen/src/egraph/cost.rs
+++ b/cranelift/codegen/src/egraph/cost.rs
@@ -31,7 +31,7 @@ use crate::ir::Opcode;
 /// that cannot be computed, or otherwise serve as a sentinel when
 /// performing search for the lowest-cost representation of a value.
 #[derive(Clone, Copy, PartialEq, Eq)]
-pub(crate) struct Cost(u32);
+pub(crate) struct Cost(pub u32);
 
 impl core::fmt::Debug for Cost {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
@@ -143,7 +143,7 @@ impl std::ops::Add<Cost> for Cost {
 ///
 /// Caller is responsible for checking that the opcode came from an instruction
 /// that satisfies `inst_predicates::is_pure_for_egraph()`.
-fn pure_op_cost(op: Opcode) -> Cost {
+pub fn pure_op_cost(op: Opcode) -> Cost {
     match op {
         // Constants.
         Opcode::Iconst | Opcode::F32const | Opcode::F64const => Cost::new(1, 0),
diff --git a/cranelift/codegen/src/egraph/elaborate.rs b/cranelift/codegen/src/egraph/elaborate.rs
index a35c2ac25..d628767ad 100644
--- a/cranelift/codegen/src/egraph/elaborate.rs
+++ b/cranelift/codegen/src/egraph/elaborate.rs
@@ -1,6 +1,8 @@
 //! Elaboration phase: lowers EGraph back to sequences of operations
 //! in CFG nodes.
 
+use std::path::Path;
+
 use super::cost::Cost;
 use super::Stats;
 use crate::dominator_tree::DominatorTreePreorder;
@@ -13,6 +15,7 @@ use crate::trace;
 use alloc::vec::Vec;
 use cranelift_control::ControlPlane;
 use cranelift_entity::{packed_option::ReservedValue, SecondaryMap};
+use egraph_serialize::{ClassId, EGraph, Node};
 use rustc_hash::{FxHashMap, FxHashSet};
 use smallvec::{smallvec, SmallVec};
 
@@ -352,6 +355,103 @@ impl<'a> Elaborator<'a> {
                 // produce suboptimal code, but never an incorrectness.
             }
         }
+
+        let mut value2class = FxHashMap::<Value, ClassId>::default();
+        let mut class2value = FxHashMap::<ClassId, Vec<Value>>::default();
+        let mut value2parent = FxHashMap::<Value, Value>::default();
+        // compute the classes
+        for (value, def) in self.func.dfg.values_and_defs() {
+            match def {
+                ValueDef::Union(x, y) => {
+                    let x_id = value2class
+                        .entry(x)
+                        .or_insert_with(|| {
+                            let id: ClassId = format!("{}", value.0).into();
+                            class2value.insert(id.clone(), vec![x]);
+                            id
+                        })
+                        .clone();
+                    value2class.insert(value, x_id.clone());
+                    class2value.get_mut(&x_id).unwrap().push(value);
+                    value2parent.insert(x, value);
+                    value2parent.insert(y, value);
+                    if let Some(y_id) = value2class.get(&y) {
+                        if x_id == *y_id {
+                            continue;
+                        }
+                        // union find can be faster but this is just serialization...
+                        let values = class2value.remove(y_id).unwrap();
+                        let x_id_values = class2value.get_mut(&x_id).unwrap();
+                        for v in values {
+                            value2class.insert(v, x_id.clone());
+                            x_id_values.push(v);
+                        }
+                    } else {
+                        class2value.get_mut(&x_id).unwrap().push(y);
+                        value2class.insert(y, x_id.clone());
+                    }
+                }
+                _ if value2class.get(&value).is_none() => {
+                    let id: ClassId = format!("{}", value.0).into();
+                    class2value.insert(id.clone(), vec![value]);
+                    value2class.insert(value, id.clone());
+                }
+                _ => {}
+            }
+        }
+        // actually construct the egraph
+        let mut egraph = EGraph::default();
+        let get_value_node_id = |value: Value| {
+            let mut current = value;
+            while let ValueDef::Union(x, _) = self.func.dfg.value_def(current) {
+                current = x;
+            }
+            format!("v{}", value.0).into()
+        };
+        for (value, def) in self.func.dfg.values_and_defs() {
+            let (children, cost) = match def {
+                ValueDef::Union(_, _) => continue,
+                ValueDef::Param(_, _) => (vec![], 0.0),
+                ValueDef::Result(inst, _) => {
+                    if let Some(_) = self.func.layout.inst_block(inst) {
+                        (vec![], 0.0)
+                    } else {
+                        let inst_data = &self.func.dfg.insts[inst];
+                        for child in self.func.dfg.inst_values(inst) {
+                            value2parent.insert(child, value);
+                        }
+                        (
+                            self.func
+                                .dfg
+                                .inst_values(inst)
+                                .map(get_value_node_id)
+                                .collect(),
+                            crate::egraph::cost::pure_op_cost(inst_data.opcode()).0 as f64,
+                        )
+                    }
+                }
+            };
+            egraph.add_node(
+                format!("v{}", value.0),
+                Node {
+                    op: format!("v{}", value.0),
+                    children,
+                    eclass: value2class.get(&value).unwrap().clone(),
+                    cost: cost.try_into().unwrap(),
+                },
+            );
+        }
+        let mut roots = FxHashSet::default();
+        for value in self.func.dfg.values() {
+            let mut current = value;
+            while let Some(&parent) = value2parent.get(&current) {
+                current = parent;
+            }
+            roots.insert(value2class.get(&current).unwrap().clone());
+        }
+        egraph.root_eclasses = roots.into_iter().collect();
+        let name = format!("{}.json", self.func.name);
+        egraph.to_json_file(Path::new(&name)).unwrap();
     }
 
     /// Elaborate use of an eclass, inserting any needed new
diff --git a/cranelift/codegen/src/ir/entities.rs b/cranelift/codegen/src/ir/entities.rs
index 005007471..0d5031792 100644
--- a/cranelift/codegen/src/ir/entities.rs
+++ b/cranelift/codegen/src/ir/entities.rs
@@ -34,7 +34,7 @@ use serde_derive::{Deserialize, Serialize};
 /// While the order is stable, it is arbitrary and does not necessarily resemble the layout order.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
-pub struct Block(u32);
+pub struct Block(pub u32);
 entity_impl!(Block, "block");
 
 impl Block {
@@ -68,7 +68,7 @@ impl Block {
 /// While the order is stable, it is arbitrary.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
-pub struct Value(u32);
+pub struct Value(pub u32);
 entity_impl!(Value, "v");
 
 impl Value {

My script for running everything in sightglass and get the json files:

#!/usr/bin/env bash

for f in $( find .. -name "*.wasm" ); do
  ff=${f/.wasm/}
  ff=${ff/..\/benchmarks\//}
  ff=${ff/\//--}
  $WASMTIME compile $f
  mkdir $ff
  for j in *.json; do
    mv $j $ff/${ff}-${j}
  done
done

for f in *; do
  if [ -d $f ]; then
    rm -r $f
  fi
done

[mwillsey]

I think having them all somewhere is a probably good, but you're right, I think having all 26k is maybe a bit much. Perhaps the largest 100?

I think we need to have some notion of roots. In cranelift, so all the un-parented nodes get extracted?

[pca006132]

I think having them all somewhere is a probably good, but you're right, I think having all 26k is maybe a bit much. Perhaps the largest 100?

Sounds good. I will put a zip file and the largest 100 files.

I think we need to have some notion of roots. In cranelift, so all the un-parented nodes get extracted?

I need to check this, I am not sure. They extract more things in additional to un-parented nodes, e.g. nodes with zero cost usually means they must be extracted. Also, in their code I just saw this comment:

Note that as values may refer to unions that represent a subset
of a larger eclass, it's not valid to walk towards the root of a
union tree: doing so would potentially equate values that fall
on different branches of the dominator tree.

So things are probably a bit more complicated...