unionfind: robustly avoid changing `Idx`s in the GVN map (#7746)

* unionfind: robustly avoid changing `Idx`s in the GVN map

Stop hoping that keeping the smallest Idx as canonical will yield good
results, and instead explicitly inform the UnionFind of what we expect
not to move.

Fixes #6126

* unionfind: track unions of pinned indices as stats

Emitting a warning in this situation is too much.

cranelift/codegen/src/context.rs

@@ -367,6 +367,7 @@ impl Context {
         );
         pass.run();
         log::debug!("egraph stats: {:?}", pass.stats);
+        trace!("pinned_union_count: {}", pass.eclasses.pinned_union_count);
         trace!("After egraph optimization:\n{}", self.func.display());
 
         self.verify_if(fisa)

cranelift/codegen/src/egraph.rs

@@ -67,7 +67,7 @@ pub struct EgraphPass<'a> {
     pub(crate) stats: Stats,
     /// Union-find that maps all members of a Union tree (eclass) back
     /// to the *oldest* (lowest-numbered) `Value`.
-    eclasses: UnionFind<Value>,
+    pub(crate) eclasses: UnionFind<Value>,
 }
 
 // The maximum number of rewrites we will take from a single call into ISLE.
@@ -109,7 +109,7 @@ impl NewOrExistingInst {
             NewOrExistingInst::New(data, ty) => (*ty, *data),
             NewOrExistingInst::Existing(inst) => {
                 let ty = dfg.ctrl_typevar(*inst);
-                (ty, dfg.insts[*inst].clone())
+                (ty, dfg.insts[*inst])
             }
         }
     }
@@ -195,15 +195,17 @@ where
             };
 
             let opt_value = self.optimize_pure_enode(inst);
+
+            for &argument in self.func.dfg.inst_args(inst) {
+                self.eclasses.pin_index(argument);
+            }
+
             let gvn_context = GVNContext {
                 union_find: self.eclasses,
                 value_lists: &self.func.dfg.value_lists,
             };
-            self.gvn_map.insert(
-                (ty, self.func.dfg.insts[inst].clone()),
-                opt_value,
-                &gvn_context,
-            );
+            self.gvn_map
+                .insert((ty, self.func.dfg.insts[inst]), opt_value, &gvn_context);
             self.value_to_opt_value[result] = opt_value;
             opt_value
         }
@@ -434,6 +436,7 @@ impl<'a> EgraphPass<'a> {
             }
         }
         trace!("stats: {:#?}", self.stats);
+        trace!("pinned_union_count: {}", self.eclasses.pinned_union_count);
         self.elaborate();
     }
 

cranelift/codegen/src/unionfind.rs

@@ -3,12 +3,43 @@
 use crate::trace;
 use cranelift_entity::{packed_option::ReservedValue, EntityRef, SecondaryMap};
 use std::hash::Hash;
+use std::mem::swap;
 
 /// A union-find data structure. The data structure can allocate
-/// `Id`s, indicating eclasses, and can merge eclasses together.
+/// `Idx`s, indicating eclasses, and can merge eclasses together.
+///
+/// Running `union(a, b)` will change the canonical `Idx` of `a` or `b`.
+/// Usually, this is chosen based on what will minimize path lengths,
+/// but it is also possible to _pin_ an eclass, such that its canonical `Idx`
+/// won't change unless it gets unioned with another pinned eclass.
+///
+/// In the context of the egraph pass, merging two pinned eclasses
+/// is very unlikely to happen – we do not know a single concrete test case
+/// where it does. The only situation where it might happen looks as follows:
+///
+/// 1. We encounter terms `A` and `B`, and the optimizer does not find any
+///    reason to union them together.
+/// 2. We encounter a term `C`, and we rewrite `C -> A`, and separately, `C -> B`.
+///
+/// Unless `C` somehow includes some crucial hint without which it is hard to
+/// notice that `A = B`, there's probably a rewrite rule that we should add.
+///
+/// Worst case, if we do merge two pinned eclasses, some nodes will essentially
+/// disappear from the GVN map, which only affects the quality of the generated
+/// code.
 #[derive(Clone, Debug, PartialEq)]
 pub struct UnionFind<Idx: EntityRef> {
     parent: SecondaryMap<Idx, Val<Idx>>,
+    /// The `rank` table is used to perform the union operations optimally,
+    /// without creating unnecessarily long paths. Pins are represented by
+    /// eclasses with a rank of `u8::MAX`.
+    ///
+    /// `rank[x]` is the upper bound on the height of the subtree rooted at `x`.
+    /// The subtree is guaranteed to have at least `2**rank[x]` elements,
+    /// unless `rank` has been artificially inflated by pinning.
+    rank: SecondaryMap<Idx, u8>,
+
+    pub(crate) pinned_union_count: u64,
 }
 
 #[derive(Clone, Debug, PartialEq)]
@@ -25,6 +56,8 @@ impl<Idx: EntityRef + Hash + std::fmt::Display + Ord + ReservedValue> UnionFind<
     pub fn with_capacity(cap: usize) -> Self {
         UnionFind {
             parent: SecondaryMap::with_capacity(cap),
+            rank: SecondaryMap::with_capacity(cap),
+            pinned_union_count: 0,
         }
     }
 
@@ -61,15 +94,50 @@ impl<Idx: EntityRef + Hash + std::fmt::Display + Ord + ReservedValue> UnionFind<
         node
     }
 
+    /// Request a stable identifier for `node`.
+    ///
+    /// After an `union` operation, the canonical representative of one
+    /// of the eclasses being merged together necessarily changes. If a pinned
+    /// eclass is merged with a non-pinned eclass, it'll be the other eclass
+    /// whose representative will change.
+    ///
+    /// If two pinned eclasses are unioned, one of the pins gets broken,
+    /// which is reported in the statistics for the pass. No concrete test case
+    /// which triggers this is known.
+    pub fn pin_index(&mut self, mut node: Idx) -> Idx {
+        node = self.find_and_update(node);
+        self.rank[node] = u8::MAX;
+        node
+    }
+
     /// Merge the equivalence classes of the two `Idx`s.
     pub fn union(&mut self, a: Idx, b: Idx) {
-        let a = self.find_and_update(a);
-        let b = self.find_and_update(b);
-        let (a, b) = (std::cmp::min(a, b), std::cmp::max(a, b));
-        if a != b {
-            // Always canonicalize toward lower IDs.
-            self.parent[b] = Val(a);
-            trace!("union: {}, {}", a, b);
+        let mut a = self.find_and_update(a);
+        let mut b = self.find_and_update(b);
+
+        if a == b {
+            return;
+        }
+
+        if self.rank[a] < self.rank[b] {
+            swap(&mut a, &mut b);
+        } else if self.rank[a] == self.rank[b] {
+            self.rank[a] = self.rank[a].checked_add(1).unwrap_or_else(
+                #[cold]
+                || {
+                    // Both `a` and `b` are pinned.
+                    //
+                    // This should only occur if we rewrite X -> Y and X -> Z,
+                    // yet neither Y -> Z nor Z -> Y can be established without
+                    // the "hint" provided by X. This probably means we're
+                    // missing an optimization rule.
+                    self.pinned_union_count += 1;
+                    u8::MAX
+                },
+            );
         }
+
+        self.parent[b] = Val(a);
+        trace!("union: {}, {}", a, b);
     }
 }

cranelift/filetests/filetests/egraph/shifts.clif

@@ -701,8 +701,8 @@ block0(v0: i8):
 }
 
 ; check: v2 = iconst.i8 1
-; check: v16 = rotr v0, v2  ; v2 = 1
-; check: return v16
+; check: v15 = rotr v0, v2  ; v2 = 1
+; check: return v15
 
 function %rotl_rotr_add(i8, i8, i8) -> i8 {
 block0(v0: i8, v1: i8, v2: i8):
@@ -725,8 +725,8 @@ block0(v0: i8):
 }
 
 ; check: v2 = iconst.i8 1
-; check: v16 = rotl v0, v2  ; v2 = 1
-; check: return v16
+; check: v15 = rotl v0, v2  ; v2 = 1
+; check: return v15
 
 function %rotl_rotr_add(i8, i8, i8) -> i8 {
 block0(v0: i8, v1: i8, v2: i8):

cranelift/filetests/filetests/egraph/unordered.clif

@@ -0,0 +1,24 @@
+test optimize precise-output
+set opt_level=speed
+target x86_64
+
+; https://github.com/bytecodealliance/wasmtime/issues/6126
+;
+; In this example, iconst defines v2, and later an identical iconst defines v1. 
+; In between, v2 is used in an iadd instruction that's added to the GVN map.
+; Finally, v1 is used in another iadd which should unify with the first one.
+function %unordered(i32) -> i32, i32 {
+block0(v0: i32):
+    v2 = iconst.i32 42
+    v3 = iadd v0, v2
+    v1 = iconst.i32 42
+    v4 = iadd v0, v1
+    return v3, v4
+}
+
+; function %unordered(i32) -> i32, i32 fast {
+; block0(v0: i32):
+;     v2 = iconst.i32 42
+;     v3 = iadd v0, v2  ; v2 = 42
+;     return v3, v3
+; }