wasm2c: Segue support for CPUs without wrgsbase instructions

3 files changed, 112 insertions, 24 deletions

diff --git a/wasm2c/examples/fac/fac.c b/wasm2c/examples/fac/fac.c
index d9f821df..e253cc2a 100644
--- a/wasm2c/examples/fac/fac.c
+++ b/wasm2c/examples/fac/fac.c
@@ -47,14 +47,15 @@
 // (1) Segue is allowed using WASM_RT_ALLOW_SEGUE
 // (2) on x86_64 without WABT_BIG_ENDIAN enabled
 // (3) the Wasm module uses a single unshared imported or exported memory
-// (4) the compiler supports: intrinsics for (rd|wr)(fs|gs)base, "address
-//     namespaces" for accessing pointers, and supports memcpy on pointers with
-//     custom "address namespaces". GCC does not support the memcpy requirement,
-//     so this leaves only clang for now.
-#if WASM_RT_ALLOW_SEGUE && !WABT_BIG_ENDIAN &&               \
-    (defined(__x86_64__) || defined(_M_X64)) &&              \
-    WASM_RT_MODULE_IS_SINGLE_UNSHARED_MEMORY && __clang__ && \
-    __has_builtin(__builtin_ia32_wrgsbase64)
+// (4) the compiler supports: intrinsics for (rd|wr)gsbase, "address namespaces"
+//     for accessing pointers, and supports memcpy on pointers with custom
+//     "address namespaces". GCC does not support the memcpy requirement, so
+//     this leaves only clang for now.
+// (5) The OS doesn't replace the segment register on context switch which
+//     eliminates windows for now
+#if WASM_RT_ALLOW_SEGUE && !WABT_BIG_ENDIAN &&                               \
+    (defined(__x86_64__) || defined(_M_X64)) && IS_SINGLE_UNSHARED_MEMORY && \
+    __clang__ && __has_builtin(__builtin_ia32_wrgsbase64) && !defined(_WIN32)
 #define WASM_RT_USE_SEGUE 1
 #else
 #define WASM_RT_USE_SEGUE 0
@@ -62,21 +63,35 @@
 #endif
 
 #if WASM_RT_USE_SEGUE
-// Different segments are free on different platforms
-// Windows uses GS for TLS, FS is free
-// Linux uses FS for TLS, GS is free
-#if defined(__WIN32)
-#define WASM_RT_SEGUE_READ_BASE() __builtin_ia32_rdfsbase64()
-#define WASM_RT_SEGUE_WRITE_BASE(base) \
-  __builtin_ia32_wrfsbase64((uintptr_t)base)
-#define MEM_ADDR_MEMOP(mem, addr, n) ((uint8_t __seg_fs*)(uintptr_t)addr)
-#else
-// POSIX style OS
-#define WASM_RT_SEGUE_READ_BASE() __builtin_ia32_rdgsbase64()
-#define WASM_RT_SEGUE_WRITE_BASE(base) \
-  __builtin_ia32_wrgsbase64((uintptr_t)base)
+// POSIX uses FS for TLS, GS is free
+#include <asm/prctl.h>
+#include <stdio.h>
+#include <sys/syscall.h>
+#include <unistd.h>
+
+static inline void* wasm_rt_segue_read_base() {
+  if (wasm_rt_fsgsbase_inst_supported) {
+    return (void*)__builtin_ia32_rdgsbase64();
+  } else {
+    void* base;
+    if (syscall(SYS_arch_prctl, ARCH_GET_GS, &base) != 0) {
+      perror("Syscall SYS_arch_prctl error");
+      abort();
+    }
+    return base;
+  }
+}
+static inline void wasm_rt_segue_write_base(void* base) {
+  if (wasm_rt_fsgsbase_inst_supported) {
+    __builtin_ia32_wrgsbase64((uintptr_t)base);
+  } else {
+    if (syscall(SYS_arch_prctl, ARCH_SET_GS, (uintptr_t)base) != 0) {
+      perror("Syscall SYS_arch_prctl error");
+      abort();
+    }
+  }
+}
 #define MEM_ADDR_MEMOP(mem, addr, n) ((uint8_t __seg_gs*)(uintptr_t)addr)
-#endif
 #else
 #define MEM_ADDR_MEMOP(mem, addr, n) MEM_ADDR(mem, addr, n)
 #endif
@@ -128,10 +143,23 @@ static inline bool func_types_eq(const wasm_rt_func_type_t a,
     TRAP(OOB);
 #endif
 
+#if WASM_RT_USE_SEGUE && WASM_RT_SANITY_CHECKS
+#include <stdio.h>
+#define WASM_RT_CHECK_BASE(mem)                                               \
+  if (((uintptr_t)((mem)->data)) != ((uintptr_t)wasm_rt_segue_read_base())) { \
+    puts("Segment register mismatch\n");                                      \
+    abort();                                                                  \
+  }
+#else
+#define WASM_RT_CHECK_BASE(mem)
+#endif
+
 #if WASM_RT_MEMCHECK_GUARD_PAGES
-#define MEMCHECK(mem, a, t)
+#define MEMCHECK(mem, a, t) WASM_RT_CHECK_BASE(mem);
 #else
-#define MEMCHECK(mem, a, t) RANGE_CHECK(mem, a, sizeof(t))
+#define MEMCHECK(mem, a, t) \
+  WASM_RT_CHECK_BASE(mem);  \
+  RANGE_CHECK(mem, a, sizeof(t))
 #endif
 
 #ifdef __GNUC__
diff --git a/wasm2c/wasm-rt-impl.c b/wasm2c/wasm-rt-impl.c
index 52a42e85..851a4c5a 100644
--- a/wasm2c/wasm-rt-impl.c
+++ b/wasm2c/wasm-rt-impl.c
@@ -49,6 +49,19 @@ static void* g_sig_handler_handle = 0;
 #endif
 #endif
 
+#if WASM_RT_USE_SEGUE || WASM_RT_ALLOW_SEGUE
+// Currently Segue is used only for linux
+#include <sys/auxv.h>
+#ifdef __GLIBC__
+#include <gnu/libc-version.h>
+#endif
+bool wasm_rt_fsgsbase_inst_supported = false;
+
+#include <asm/prctl.h>    // For ARCH_SET_GS
+#include <sys/syscall.h>  // For SYS_arch_prctl
+#include <unistd.h>       // For syscall
+#endif
+
 #if WASM_RT_STACK_DEPTH_COUNT
 WASM_RT_THREAD_LOCAL uint32_t wasm_rt_call_stack_depth;
 WASM_RT_THREAD_LOCAL uint32_t wasm_rt_saved_call_stack_depth;
@@ -220,6 +233,15 @@ void wasm_rt_init(void) {
     os_install_signal_handler();
   }
 #endif
+
+#if WASM_RT_USE_SEGUE || WASM_RT_ALLOW_SEGUE
+#if defined(__GLIBC__) && __GLIBC__ >= 2 && __GLIBC_MINOR__ >= 18
+  // Check for support for userspace wrgsbase instructions
+  unsigned long val = getauxval(AT_HWCAP2);
+  wasm_rt_fsgsbase_inst_supported = val & (1 << 1);
+#endif
+#endif
+
   assert(wasm_rt_is_initialized());
 }
 
@@ -257,6 +279,23 @@ void wasm_rt_free_thread(void) {
 #endif
 }
 
+#if WASM_RT_USE_SEGUE || WASM_RT_ALLOW_SEGUE
+void wasm_rt_syscall_set_segue_base(void* base) {
+  if (syscall(SYS_arch_prctl, ARCH_SET_GS, base) != 0) {
+    perror("wasm_rt_syscall_set_segue_base error");
+    abort();
+  }
+}
+void* wasm_rt_syscall_get_segue_base() {
+  void* base;
+  if (syscall(SYS_arch_prctl, ARCH_GET_GS, &base) != 0) {
+    perror("wasm_rt_syscall_get_segue_base error");
+    abort();
+  }
+  return base;
+}
+#endif
+
 // Include table operations for funcref
 #define WASM_RT_TABLE_OPS_FUNCREF
 #include "wasm-rt-impl-tableops.inc"
diff --git a/wasm2c/wasm-rt.h b/wasm2c/wasm-rt.h
index 5c07ff44..c489c850 100644
--- a/wasm2c/wasm-rt.h
+++ b/wasm2c/wasm-rt.h
@@ -294,6 +294,27 @@ extern WASM_RT_THREAD_LOCAL uint32_t wasm_rt_call_stack_depth;
 
 #endif
 
+#if WASM_RT_USE_SEGUE || WASM_RT_ALLOW_SEGUE
+/**
+ * The segue optimization uses x86 segments to point to a linear memory. If
+ * used, the runtime must query whether it can use the fast userspace wrgsbase
+ * instructions or whether it must invoke syscalls to set the segment base,
+ * depending on the supported CPU features. The result of this query is saved in
+ * this variable.
+ */
+extern bool wasm_rt_fsgsbase_inst_supported;
+/**
+ * If fast userspace wrgsbase instructions don't exist, the runtime most provide
+ * a function that invokes the OS' underlying syscall to set the segment base.
+ */
+void wasm_rt_syscall_set_segue_base(void* base);
+/**
+ * If fast userspace rdgsbase instructions don't exist, the runtime most provide
+ * a function that invokes the OS' underlying syscall to get the segment base.
+ */
+void* wasm_rt_syscall_get_segue_base();
+#endif
+
 #if defined(_MSC_VER)
 #define WASM_RT_NO_RETURN __declspec(noreturn)
 #else