Optimize wasm calls ever-so-slightly (#7084)

* Optimize wasm calls ever-so-slightly * Fix riscv64 * Remove stray comment * Shuffle around where `unsafe` lies
bytecodealliance · Sep 25, 2023 · 2850d72 · 2850d72
1 parent 6a7ef27
commit 2850d72
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 2 deletions.
diff --git a/crates/runtime/src/instance.rs b/crates/runtime/src/instance.rs
@@ -426,6 +426,7 @@ impl Instance {
     }
 
     /// Return a pointer to the interrupts structure
+    #[inline]
     pub fn runtime_limits(&mut self) -> *mut *const VMRuntimeLimits {
         unsafe { self.vmctx_plus_offset_mut(self.offsets().vmctx_runtime_limits()) }
     }

diff --git a/crates/wasmtime/Cargo.toml b/crates/wasmtime/Cargo.toml
@@ -41,7 +41,6 @@ serde_json = { workspace = true }
 bincode = "1.2.1"
 indexmap = { workspace = true }
 paste = "1.0.3"
-psm = "0.1.11"
 once_cell = { workspace = true }
 rayon = { version = "1.0", optional = true }
 object = { workspace = true }
@@ -56,6 +55,9 @@ features = [
   "Win32_System_Diagnostics_Debug",
 ]
 
+[target.'cfg(target_arch = "s390x")'.dependencies]
+psm = "0.1.11"
+
 [dev-dependencies]
 tempfile = "3.0"
 wasmtime-wasi = { path = "../wasi", default-features = true }

diff --git a/crates/wasmtime/src/func.rs b/crates/wasmtime/src/func.rs
@@ -1408,7 +1408,7 @@ fn enter_wasm<T>(store: &mut StoreContextMut<'_, T>) -> Option<usize> {
         return None;
     }
 
-    let stack_pointer = psm::stack_pointer() as usize;
+    let stack_pointer = stack_pointer();
 
     // Determine the stack pointer where, after which, any wasm code will
     // immediately trap. This is checked on the entry to all wasm functions.
@@ -1435,6 +1435,43 @@ fn enter_wasm<T>(store: &mut StoreContextMut<'_, T>) -> Option<usize> {
     Some(prev_stack)
 }
 
+#[inline]
+fn stack_pointer() -> usize {
+    let stack_pointer: usize;
+    cfg_if::cfg_if! {
+        if #[cfg(target_arch = "x86_64")] {
+            unsafe {
+                std::arch::asm!(
+                    "mov {}, rsp",
+                    out(reg) stack_pointer,
+                    options(nostack,nomem),
+                );
+            }
+        } else if #[cfg(target_arch = "aarch64")] {
+            unsafe {
+                std::arch::asm!(
+                    "mov {}, sp",
+                    out(reg) stack_pointer,
+                    options(nostack,nomem),
+                );
+            }
+        } else if #[cfg(target_arch = "riscv64")] {
+            unsafe {
+                std::arch::asm!(
+                    "mv {}, sp",
+                    out(reg) stack_pointer,
+                    options(nostack,nomem),
+                );
+            }
+        } else if #[cfg(target_arch = "s390x")] {
+            stack_pointer = psm::stack_pointer() as usize;
+        } else {
+            compile_error!("unsupported platform");
+        }
+    }
+    stack_pointer
+}
+
 fn exit_wasm<T>(store: &mut StoreContextMut<'_, T>, prev_stack: Option<usize>) {
     // If we don't have a previous stack pointer to restore, then there's no
     // cleanup we need to perform here.