Repository: bytecodealliance/wasmtime
Pull request #7466 (2023-11-02 21:57:04)

Cranelift: Reassociate long and narrow chains of operations into shallow and wide trees
This increases instruction-level parallelism and shrinks live ranges. It also
canonicalizes into the shallow-and-wide form for reassociating constants
together for cprop.

Depends on https://github.com/bytecodealliance/wasmtime/pull/7456

github-actions[bot]: #### Subscribe to Label Action

cc @cfallin, @fitzgen

<details>
This issue or pull request has been labeled: "cranelift", "isle"

Thus the following users have been cc'd because of the following labels:

* cfallin: isle
* fitzgen: isle

To subscribe or unsubscribe from this label, edit the <code>.github/subscribe-to-label.json</code> configuration file.

[Learn more.](https://github.com/bytecodealliance/subscribe-to-label-action)
</details>

Patch:
diff --git a/cranelift/codegen/src/egraph.rs b/cranelift/codegen/src/egraph.rs
index dab240c3a..2d5802572 100644
--- a/cranelift/codegen/src/egraph.rs
+++ b/cranelift/codegen/src/egraph.rs
@@ -424,7 +424,7 @@ impl<'a> EgraphPass<'a> {
                 }
             }
         }
-        trace!("stats: {:?}", self.stats);
+        trace!("stats: {:#?}", self.stats);
         self.elaborate();
     }
 
diff --git a/cranelift/codegen/src/opts/arithmetic.isle b/cranelift/codegen/src/opts/arithmetic.isle
index 00c56c240..6553b8674 100644
--- a/cranelift/codegen/src/opts/arithmetic.isle
+++ b/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))))
diff --git a/cranelift/codegen/src/opts/cprop.isle b/cranelift/codegen/src/opts/cprop.isle
index 26580a7ea..e44404e05 100644
--- a/cranelift/codegen/src/opts/cprop.isle
+++ b/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)))
diff --git a/cranelift/codegen/src/unionfind.rs b/cranelift/codegen/src/unionfind.rs
index a51f305c2..8e2be3ba3 100644
--- a/cranelift/codegen/src/unionfind.rs
+++ b/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())
diff --git a/crates/cache/src/lib.rs b/crates/cache/src/lib.rs
index a997846c1..279831259 100644
--- a/crates/cache/src/lib.rs
+++ b/crates/cache/src/lib.rs
@@ -2,11 +2,11 @@ use base64::Engine;
 use log::{debug, trace, warn};
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
-use std::fs;
 use std::hash::Hash;
 use std::hash::Hasher;
 use std::io::Write;
 use std::path::{Path, PathBuf};
+use std::{fs, io};
 
 #[macro_use] // for tests
 mod config;
@@ -172,7 +172,7 @@ impl<'config> ModuleCacheEntryInner<'config> {
 
         // Optimize syscalls: first, try writing to disk. It should succeed in most cases.
         // Otherwise, try creating the cache directory and retry writing to the file.
-        if fs_write_atomic(&mod_cache_path, "mod", &compressed_data) {
+        if fs_write_atomic(&mod_cache_path, "mod", &compressed_data).is_ok() {
             return Some(());
         }
 
@@ -193,10 +193,16 @@ impl<'config> ModuleCacheEntryInner<'config> {
             })
             .ok()?;
 
-        if fs_write_atomic(&mod_cache_path, "mod", &compressed_data) {
-            Some(())
-        } else {
-            None
+        match fs_write_atomic(&mod_cache_path, "mod", &compressed_data) {
+            Ok(_) => Some(()),
+            Err(err) => {
+                warn!(
+                    "Failed to write file with rename, target path: {}, err: {}",
+                    mod_cache_path.display(),
+                    err
+                );
+                None
+            }
         }
     }
 }
@@ -214,7 +220,7 @@ impl Hasher for Sha256Hasher {
 // Assumption: path inside cache directory.
 // Then, we don't have to use sound OS-specific exclusive file access.
 // Note: there's no need to remove temporary file here - cleanup task will do it later.
-fn fs_write_atomic(path: &Path, reason: &str, contents: &[u8]) -> bool {
+fn fs_write_atomic(path: &Path, reason: &str, contents: &[u8]) -> io::Result<()> {
     let lock_path = path.with_extension(format!("wip-atomic-write-{}", reason));
     fs::OpenOptions::new()
         .create_new(true) // atomic file creation (assumption: no one will open it without this flag)
@@ -223,15 +229,6 @@ fn fs_write_atomic(path: &Path, reason: &str, contents: &[u8]) -> bool {
         .and_then(|mut file| file.write_all(contents))
         // file should go out of scope and be closed at this point
         .and_then(|()| fs::rename(&lock_path, &path)) // atomic file rename
-        .map_err(|err| {
-            warn!(
-                "Failed to write file with rename, lock path: {}, target path: {}, err: {}",
-                lock_path.display(),
-                path.display(),
-                err
-            )
-        })
-        .is_ok()
 }
 
 #[cfg(test)]
diff --git a/crates/cache/src/worker.rs b/crates/cache/src/worker.rs
index 72a53f94c..64eff12f3 100644
--- a/crates/cache/src/worker.rs
+++ b/crates/cache/src/worker.rs
@@ -758,11 +758,7 @@ fn write_stats_file(path: &Path, stats: &ModuleCacheStatistics) -> bool {
             )
         })
         .and_then(|serialized| {
-            if fs_write_atomic(path, "stats", serialized.as_bytes()) {
-                Ok(())
-            } else {
-                Err(())
-            }
+            fs_write_atomic(path, "stats", serialized.as_bytes()).map_err(|_| ())
         })
         .is_ok()
 }


Test patch:
diff --git a/cranelift/filetests/filetests/egraph/associative-and-commutative.clif b/cranelift/filetests/filetests/egraph/associative-and-commutative.clif
new file mode 100644
index 000000000..37bfa8da5
--- /dev/null
+++ b/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
+}
diff --git a/cranelift/filetests/filetests/egraph/icmp.clif b/cranelift/filetests/filetests/egraph/icmp.clif
index 1a8874845..ef60cbf95 100644
--- a/cranelift/filetests/filetests/egraph/icmp.clif
+++ b/cranelift/filetests/filetests/egraph/icmp.clif
@@ -53,4 +53,3 @@ block0:
 ;     v5 = icmp ne v2, v3  ; v3 = 0
 ;     return v5
 ; }
-