Cranelift: Reassociate long and narrow chains of operations into shallow and wide trees

cranelift/codegen/src/egraph.rs

@@ -424,7 +424,7 @@ impl<'a> EgraphPass<'a> {
                 }
             }
         }
-        trace!("stats: {:?}", self.stats);
+        trace!("stats: {:#?}", self.stats);
         self.elaborate();
     }
 

cranelift/codegen/src/opts/arithmetic.isle

@@ -126,3 +126,41 @@
 ;; both of them since they cancel out.
 (rule (simplify (fmul ty (fneg ty x) (fneg ty y)))
       (fmul ty x y))
+
+;; (a op (b op (c op d))) ==> ((a op b) op (c op d))
+;;
+;; and
+;;
+;; (((a op b) op c) op d) ==> ((a op b) op (c op d))
+;;
+;; where `op` is an associative operation: `iadd`, `imul`, `band`, or `bxor`.
+;;
+;; This increases instruction-level parallelism and shrinks live ranges. It also
+;; canonicalizes into the shallow-and-wide form for reassociating constants
+;; together for cprop.
+;;
+;; NB: We subsume to avoid exponential e-node blow up due to reassociating very
+;; large chains of operations.
+;;
+;; TODO: We should add `bor` rules for this as well. Unfortunately, they
+;; conflict with our `bswap` recognizing rules when we `subsume`.
+
+(rule (simplify (iadd ty a (iadd ty b (iadd ty c d))))
+      (subsume (iadd ty (iadd ty a b) (iadd ty c d))))
+(rule (simplify (iadd ty (iadd ty (iadd ty a b) c) d))
+      (subsume (iadd ty (iadd ty a b) (iadd ty c d))))
+
+(rule (simplify (imul ty a (imul ty b (imul ty c d))))
+      (subsume (imul ty (imul ty a b) (imul ty c d))))
+(rule (simplify (imul ty (imul ty (imul ty a b) c) d))
+      (subsume (imul ty (imul ty a b) (imul ty c d))))
+
+(rule (simplify (band ty a (band ty b (band ty c d))))
+      (subsume (band ty (band ty a b) (band ty c d))))
+(rule (simplify (band ty (band ty (band ty a b) c) d))
+      (subsume (band ty (band ty a b) (band ty c d))))
+
+(rule (simplify (bxor ty a (bxor ty b (bxor ty c d))))
+      (subsume (bxor ty (bxor ty a b) (bxor ty c d))))
+(rule (simplify (bxor ty (bxor ty (bxor ty a b) c) d))
+      (subsume (bxor ty (bxor ty a b) (bxor ty c d))))

cranelift/codegen/src/opts/cprop.isle

@@ -208,3 +208,30 @@
       (ushr ty (ushr ty a c) b))
 (rule (simplify (sshr ty (sshr ty a@(iconst _ _) b) c@(iconst _ _)))
       (sshr ty (sshr ty a c) b))
+
+;; When we operations that are both commutative and associative, reassociate
+;; constants together for cprop:
+;;
+;; ((a op B) op (c op D)) ==> ((a op c) op (B op D))
+;;
+;; Where `op` is one of: `iadd`, `imul`, `band`, `bor`, or `bxor`.
+(rule (simplify (iadd ty
+                      (iadd ty a b@(iconst _ _))
+                      (iadd ty c d@(iconst _ _))))
+      (iadd ty (iadd ty a c) (iadd ty b d)))
+(rule (simplify (imul ty
+                      (imul ty a b@(iconst _ _))
+                      (imul ty c d@(iconst _ _))))
+      (imul ty (imul ty a c) (imul ty b d)))
+(rule (simplify (band ty
+                      (band ty a b@(iconst _ _))
+                      (band ty c d@(iconst _ _))))
+      (band ty (band ty a c) (band ty b d)))
+(rule (simplify (bor ty
+                      (bor ty a b@(iconst _ _))
+                      (bor ty c d@(iconst _ _))))
+      (bor ty (bor ty a c) (bor ty b d)))
+(rule (simplify (bxor ty
+                      (bxor ty a b@(iconst _ _))
+                      (bxor ty c d@(iconst _ _))))
+      (bxor ty (bxor ty a c) (bxor ty b d)))

cranelift/codegen/src/unionfind.rs

@@ -13,6 +13,7 @@ pub struct UnionFind<Idx: EntityRef> {
 
 #[derive(Clone, Debug, PartialEq)]
 struct Val<Idx>(Idx);
+
 impl<Idx: EntityRef + ReservedValue> Default for Val<Idx> {
     fn default() -> Self {
         Self(Idx::reserved_value())

cranelift/filetests/filetests/egraph/associative-and-commutative.clif

@@ -0,0 +1,81 @@
+test optimize
+set opt_level=speed
+target x86_64
+
+function %iadd_shallow_and_wide(i32, i32, i32, i32) -> i32 {
+block0(v0: i32, v1: i32, v2: i32, v3: i32):
+    v4 = iadd v0, v1
+    v5 = iadd v4, v2
+    v6 = iadd v5, v3
+    return v6
+; check:  v4 = iadd v0, v1
+; nextln: v7 = iadd v2, v3
+; nextln: v8 = iadd v4, v7
+; check: return v8
+}
+
+function %iadd_shallow_and_wide_twice(i32, i32, i32, i32, i32, i32) -> i32 {
+block0(v0: i32, v1: i32, v2: i32, v3: i32, v4: i32, v5: i32):
+    v6 = iadd v0, v1
+    v7 = iadd v6, v2
+    v8 = iadd v7, v3
+    v9 = iadd v8, v4
+    v10 = iadd v9, v5
+    return v10
+; check:  v6 = iadd v0, v1
+; nextln: v11 = iadd v2, v3
+; nextln: v12 = iadd v6, v11
+; nextln: v16 = iadd v4, v5
+; nextln: v23 = iadd v12, v16
+; nextln: return v23
+}
+
+function %iadd_reassociate_then_cprop(i32, i32) -> i32 {
+block0(v0: i32, v1: i32):
+    v2 = iconst.i32 42
+    v3 = iconst.i32 36
+    v4 = iadd v0, v2
+    v5 = iadd v1, v3
+    v6 = iadd v4, v5
+    return v6
+; check:  v7 = iadd v0, v1
+; nextln: iconst.i32 78
+; nextln: v16 = iadd v7, v14  ; v14 = 78
+; check:  return v16
+}
+
+function %imul_shallow_and_wide(i32, i32, i32, i32) -> i32 {
+block0(v0: i32, v1: i32, v2: i32, v3: i32):
+    v4 = imul v0, v1
+    v5 = imul v4, v2
+    v6 = imul v5, v3
+    return v6
+; check:  v4 = imul v0, v1
+; nextln: v7 = imul v2, v3
+; nextln: v8 = imul v4, v7
+; check: return v8
+}
+
+function %band_shallow_and_wide(i32, i32, i32, i32) -> i32 {
+block0(v0: i32, v1: i32, v2: i32, v3: i32):
+    v4 = band v0, v1
+    v5 = band v4, v2
+    v6 = band v5, v3
+    return v6
+; check:  v4 = band v0, v1
+; nextln: v7 = band v2, v3
+; nextln: v8 = band v4, v7
+; check: return v8
+}
+
+function %bxor_shallow_and_wide(i32, i32, i32, i32) -> i32 {
+block0(v0: i32, v1: i32, v2: i32, v3: i32):
+    v4 = bxor v0, v1
+    v5 = bxor v4, v2
+    v6 = bxor v5, v3
+    return v6
+; check:  v4 = bxor v0, v1
+; nextln: v7 = bxor v2, v3
+; nextln: v8 = bxor v4, v7
+; check: return v8
+}

cranelift/filetests/filetests/egraph/icmp.clif

@@ -53,4 +53,3 @@ block0:
 ;     v5 = icmp ne v2, v3  ; v3 = 0
 ;     return v5
 ; }
-