Optimizations to egraph framework:

- Save elaborated results by canonical value, not latest value (union
  value). Previously we were artificially skipping and re-elaborating
  some values we already had because we were not finding them in the
  map.

- Make some changes to handling of icmp results: when icmp became
  I8-typed (when bools went away), many uses became `(uextend $I32 (icmp
  $I8 ...))`, and so patterns in lowering backends were no longer
  matching.

  This PR includes an x64-specific change to match `(brz (uextend (icmp
  ...)))` and similarly for `brnz`, but it also takes advantage of the
  ability to write rules easily in the egraph mid-end to rewrite selects
  with icmp inputs appropriately.

- Extend constprop to understand selects in the egraph mid-end.

With these changes, bz2.wasm sees a ~1% speedup, and spidermonkey.wasm
with a fib.js input sees a 16.8% speedup:

```
$ time taskset 1 target/release/wasmtime run --allow-precompiled --dir=. ./spidermonkey.base.cwasm ./fib.js
1346269
taskset 1 target/release/wasmtime run --allow-precompiled --dir=.  ./fib.js  2.14s user 0.01s system 99% cpu 2.148 total
$ time taskset 1 target/release/wasmtime run --allow-precompiled --dir=. ./spidermonkey.egraphs.cwasm ./fib.js
1346269
taskset 1 target/release/wasmtime run --allow-precompiled --dir=.  ./fib.js  1.78s user 0.01s system 99% cpu 1.788 total
```

Author: cfallin

cranelift/codegen/src/egraph/elaborate.rs

@@ -277,7 +277,7 @@ impl<'a> Elaborator<'a> {
                     let value = self.func.dfg.resolve_aliases(value);
 
                     self.stats.elaborate_visit_node += 1;
-                    let canonical_value = self.eclasses.find(value);
+                    let canonical_value = self.eclasses.find_and_update(value);
                     debug_assert_ne!(canonical_value, Value::reserved_value());
                     trace!(
                         "elaborate: value {} canonical {} before {}",
@@ -518,8 +518,9 @@ impl<'a> Elaborator<'a> {
                                 value: new_result,
                                 in_block: insert_block,
                             };
+                            let canonical_result = self.eclasses.find_and_update(result);
                             self.value_to_elaborated_value.insert_if_absent_with_depth(
-                                result,
+                                canonical_result,
                                 elab_value,
                                 scope_depth,
                             );
@@ -545,8 +546,9 @@ impl<'a> Elaborator<'a> {
                                 value: result,
                                 in_block: insert_block,
                             };
+                            let canonical_result = self.eclasses.find_and_update(result);
                             self.value_to_elaborated_value.insert_if_absent_with_depth(
-                                result,
+                                canonical_result,
                                 elab_value,
                                 scope_depth,
                             );
@@ -623,8 +625,9 @@ impl<'a> Elaborator<'a> {
             // map now.
             for &result in self.func.dfg.inst_results(inst) {
                 trace!(" -> result {}", result);
+                let canonical_result = self.eclasses.find_and_update(result);
                 self.value_to_elaborated_value.insert_if_absent(
-                    result,
+                    canonical_result,
                     ElaboratedValue {
                         in_block: block,
                         value: result,

cranelift/codegen/src/isa/x64/lower.isle

@@ -2903,7 +2903,11 @@
       (let ((cmp IcmpCondResult (invert_icmp_cond_result (emit_cmp cc a b))))
         (side_effect (jmp_cond_icmp cmp taken not_taken))))
 
-(rule 2 (lower_branch (brz (fcmp cc a b) _ _) (two_targets taken not_taken))
+(rule 2 (lower_branch (brz (uextend (icmp cc a b)) _ _) (two_targets taken not_taken))
+      (let ((cmp IcmpCondResult (invert_icmp_cond_result (emit_cmp cc a b))))
+        (side_effect (jmp_cond_icmp cmp taken not_taken))))
+
+(rule 2 (lower_branch (brz (uextend (fcmp cc a b)) _ _) (two_targets taken not_taken))
       (let ((cmp FcmpCondResult (emit_fcmp (floatcc_inverse cc) a b)))
         (side_effect (jmp_cond_fcmp cmp taken not_taken))))
 
@@ -2923,6 +2927,13 @@
       (let ((cmp FcmpCondResult (emit_fcmp cc a b)))
         (side_effect (jmp_cond_fcmp cmp taken not_taken))))
 
+(rule 2 (lower_branch (brnz (uextend (icmp cc a b)) _ _) (two_targets taken not_taken))
+      (side_effect (jmp_cond_icmp (emit_cmp cc a b) taken not_taken)))
+
+(rule 2 (lower_branch (brnz (uextend (fcmp cc a b)) _ _) (two_targets taken not_taken))
+      (let ((cmp FcmpCondResult (emit_fcmp cc a b)))
+        (side_effect (jmp_cond_fcmp cmp taken not_taken))))
+
 (rule 1 (lower_branch (brnz val @ (value_type $I128) _ _) (two_targets taken not_taken))
       (side_effect (jmp_cond_icmp (cmp_zero_i128 (CC.Z) val) taken not_taken)))
 

cranelift/codegen/src/isle_prelude.rs

@@ -607,5 +607,11 @@ macro_rules! isle_common_prelude_methods {
         fn pack_value_array_3(&mut self, a: Value, b: Value, c: Value) -> ValueArray3 {
             [a, b, c]
         }
+
+        #[inline]
+        fn intcc_invert(&mut self, cc: &IntCC) -> IntCC {
+            use crate::ir::condcodes::CondCode;
+            cc.inverse()
+        }
     };
 }

cranelift/codegen/src/opts/algebraic.isle

@@ -7,8 +7,8 @@
 ;; can freely rewrite e.g. `x+y-y` to `x`.
 
 ;; uextend/sextend of a constant.
-(rule (simplify (uextend $I64 (iconst $I32 (u64_from_imm64 imm))))
-      (iconst $I64 (imm64 (u64_uextend_u32 imm))))
+(rule (simplify (uextend ty (iconst _ (u64_from_imm64 imm))))
+      (iconst ty (imm64 imm)))
 (rule (simplify (sextend $I64 (iconst $I32 (u64_from_imm64 imm))))
       (iconst $I64 (imm64 (u64_sextend_u32 imm))))
 
@@ -185,3 +185,25 @@
       (remat x))
 (rule (simplify x @ (f64const _ _))
       (remat x))
+
+;; Optimize icmp-of-icmp.
+(rule (simplify (icmp ty
+                      (IntCC.NotEqual)
+                      (uextend _ inner @ (icmp ty _ _ _))
+                      (iconst _ (u64_from_imm64 0))))
+      (subsume inner))
+
+(rule (simplify (icmp ty
+                      (IntCC.Equal)
+                      (uextend _ (icmp ty cc x y))
+                      (iconst _ (u64_from_imm64 0))))
+      (subsume (icmp ty (intcc_invert cc) x y)))
+
+;; Optimize select-of-uextend-of-icmp to select-of-icmp, because
+;; select can take an I8 condition too.
+(rule (simplify
+       (select ty (uextend _ c @ (icmp _ _ _ _)) x y))
+      (select ty c x y))
+(rule (simplify
+       (select ty (uextend _ c @ (icmp _ _ _ _)) x y))
+      (select ty c x y))

cranelift/codegen/src/opts/cprop.isle

@@ -130,5 +130,12 @@
        (bxor ty (bxor ty x k1 @ (iconst ty _)) k2 @ (iconst ty _)))
       (bxor ty x (bxor ty k1 k2)))
 
+(rule (simplify
+       (select ty (iconst _ (u64_from_imm64 (u64_nonzero _))) x y))
+      x)
+(rule (simplify
+       (select ty (iconst _ (u64_from_imm64 0)) x y))
+      y)
+
 ;; TODO: fadd, fsub, fmul, fdiv, fneg, fabs
 

cranelift/codegen/src/prelude.isle

@@ -383,6 +383,10 @@
 (decl pure signed_cond_code (IntCC) IntCC)
 (extern constructor signed_cond_code signed_cond_code)
 
+;; Constructor that inverts a condcode.
+(decl intcc_invert (IntCC) IntCC)
+(extern constructor intcc_invert intcc_invert)
+
 ;;;; Helpers for Working with TrapCode ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 (decl trap_code_division_by_zero () TrapCode)