From 2850d7236c398d0d86571e911d46677e51e49bf2 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Mon, 25 Sep 2023 16:23:18 -0500
Subject: [PATCH] Optimize wasm calls ever-so-slightly (#7084)

* Optimize wasm calls ever-so-slightly

* Fix riscv64

* Remove stray comment

* Shuffle around where `unsafe` lies
---
 crates/runtime/src/instance.rs |  1 +
 crates/wasmtime/Cargo.toml     |  4 +++-
 crates/wasmtime/src/func.rs    | 39 +++++++++++++++++++++++++++++++++-
 3 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/crates/runtime/src/instance.rs b/crates/runtime/src/instance.rs
index 7ec1470af249..8035f5db7f4a 100644
--- 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
index 6aa56e8f54f1..d735b74b6feb 100644
--- 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
index 60e2a5047903..3ec071c72980 100644
--- 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.