Repository: bytecodealliance/wasmtime
Pull request #7450 (2023-11-01 21:52:15)

Cranelift: Reassociate constants out of nested shifts
This allows for more constant propagation.

<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->



Patch:
diff --git a/cranelift/codegen/src/opts/cprop.isle b/cranelift/codegen/src/opts/cprop.isle
index 19f5f2aac..26580a7ea 100644
--- a/cranelift/codegen/src/opts/cprop.isle
+++ b/cranelift/codegen/src/opts/cprop.isle
@@ -198,3 +198,13 @@
 (rule (splat32 n) (splat64 (u64_or n (u64_shl n 32))))
 (decl splat64 (u64) Constant)
 (extern constructor splat64 splat64)
+
+;; Reassociate nested shifts of constants to put constants together for cprop.
+;;
+;; ((A shift b) shift C) ==> ((A shift C) shift b)
+(rule (simplify (ishl ty (ishl ty a@(iconst _ _) b) c@(iconst _ _)))
+      (ishl ty (ishl ty a c) b))
+(rule (simplify (ushr ty (ushr ty a@(iconst _ _) b) c@(iconst _ _)))
+      (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))
diff --git a/crates/wasi-preview1-component-adapter/src/lib.rs b/crates/wasi-preview1-component-adapter/src/lib.rs
index 66d56da23..e9f5c6f0f 100644
--- a/crates/wasi-preview1-component-adapter/src/lib.rs
+++ b/crates/wasi-preview1-component-adapter/src/lib.rs
@@ -94,6 +94,14 @@ impl<T, E> TrappingUnwrap<T> for Result<T, E> {
     }
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn reset_adapter_state() {
+    let state = get_state_ptr();
+    if !state.is_null() {
+        State::init(state)
+    }
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn cabi_import_realloc(
     old_ptr: *mut u8,
@@ -2385,8 +2393,8 @@ enum AllocationState {
 
 #[allow(improper_ctypes)]
 extern "C" {
-    fn get_state_ptr() -> *const State;
-    fn set_state_ptr(state: *const State);
+    fn get_state_ptr() -> *mut State;
+    fn set_state_ptr(state: *mut State);
     fn get_allocation_state() -> AllocationState;
     fn set_allocation_state(state: AllocationState);
 }
@@ -2415,7 +2423,7 @@ impl State {
     }
 
     #[cold]
-    fn new() -> &'static State {
+    fn new() -> *mut State {
         #[link(wasm_import_module = "__main_module__")]
         extern "C" {
             fn cabi_realloc(
@@ -2445,31 +2453,37 @@ impl State {
         unsafe { set_allocation_state(AllocationState::StateAllocated) };
 
         unsafe {
-            ret.write(State {
-                magic1: MAGIC,
-                magic2: MAGIC,
-                import_alloc: ImportAlloc::new(),
-                descriptors: RefCell::new(None),
-                path_buf: UnsafeCell::new(MaybeUninit::uninit()),
-                long_lived_arena: BumpArena::new(),
-                args: Cell::new(None),
-                env_vars: Cell::new(None),
-                dirent_cache: DirentCache {
-                    stream: Cell::new(None),
-                    for_fd: Cell::new(0),
-                    cookie: Cell::new(wasi::DIRCOOKIE_START),
-                    cached_dirent: Cell::new(wasi::Dirent {
-                        d_next: 0,
-                        d_ino: 0,
-                        d_type: FILETYPE_UNKNOWN,
-                        d_namlen: 0,
-                    }),
-                    path_data: UnsafeCell::new(MaybeUninit::uninit()),
-                },
-                dotdot: [UnsafeCell::new(b'.'), UnsafeCell::new(b'.')],
-            });
-            &*ret
+            Self::init(ret);
         }
+
+        ret
+    }
+
+    #[cold]
+    unsafe fn init(state: *mut State) {
+        state.write(State {
+            magic1: MAGIC,
+            magic2: MAGIC,
+            import_alloc: ImportAlloc::new(),
+            descriptors: RefCell::new(None),
+            path_buf: UnsafeCell::new(MaybeUninit::uninit()),
+            long_lived_arena: BumpArena::new(),
+            args: Cell::new(None),
+            env_vars: Cell::new(None),
+            dirent_cache: DirentCache {
+                stream: Cell::new(None),
+                for_fd: Cell::new(0),
+                cookie: Cell::new(wasi::DIRCOOKIE_START),
+                cached_dirent: Cell::new(wasi::Dirent {
+                    d_next: 0,
+                    d_ino: 0,
+                    d_type: FILETYPE_UNKNOWN,
+                    d_namlen: 0,
+                }),
+                path_data: UnsafeCell::new(MaybeUninit::uninit()),
+            },
+            dotdot: [UnsafeCell::new(b'.'), UnsafeCell::new(b'.')],
+        });
     }
 
     /// Accessor for the descriptors member that ensures it is properly initialized


Test patch:
diff --git a/cranelift/filetests/filetests/egraph/reassociate-constants-in-shifts.clif b/cranelift/filetests/filetests/egraph/reassociate-constants-in-shifts.clif
new file mode 100644
index 000000000..9d50399d5
--- /dev/null
+++ b/cranelift/filetests/filetests/egraph/reassociate-constants-in-shifts.clif
@@ -0,0 +1,41 @@
+test optimize
+set opt_level=speed
+target x86_64
+
+;; Test egraph rewrite rules that reassociate constants out of nested shifts.
+
+function %a(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 1
+    v2 = ishl v1, v0
+    v3 = iconst.i32 2
+    v4 = ishl v2, v3
+    return v4
+; check:  v6 = iconst.i32 4
+; nextln: v7 = ishl v6, v0
+; check:  return v7
+}
+
+function %b(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 8
+    v2 = ushr v1, v0
+    v3 = iconst.i32 2
+    v4 = ushr v2, v3
+    return v4
+; check:  v3 = iconst.i32 2
+; nextln: v6 = ushr v3, v0
+; check:  return v6
+}
+
+function %c(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 8
+    v2 = sshr v1, v0
+    v3 = iconst.i32 2
+    v4 = sshr v2, v3
+    return v4
+; check:  v3 = iconst.i32 2
+; nextln: v6 = sshr v3, v0
+; check:  return v6
+}