wasm2c: Segue optimization for modules with a single unshared memory (#2395)

1 files changed, 77 insertions, 14 deletions

diff --git a/test/wasm2c/hello.txt b/test/wasm2c/hello.txt
index d56c7216..77608b6c 100644
--- a/test/wasm2c/hello.txt
+++ b/test/wasm2c/hello.txt
@@ -96,6 +96,7 @@ void w2c_test_0x5Fstart(w2c_test*);
 #endif
 
 #include "wasm.h"
+#define IS_SINGLE_UNSHARED_MEMORY 1
 
 // Computes a pointer to an object of the given size in a little-endian memory.
 //
@@ -118,6 +119,39 @@ void w2c_test_0x5Fstart(w2c_test*);
 #define MEM_ADDR(mem, addr, n) &(mem)->data[addr]
 #endif
 
+#ifndef WASM_RT_USE_SEGUE
+// Memory functions can use the segue optimization if allowed. The segue
+// optimization uses x86 segments to point to a linear memory. We use this
+// optimization when:
+//
+// (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)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
+#endif
+#endif
+
+#if WASM_RT_USE_SEGUE
+// POSIX uses FS for TLS, GS is free
+#define WASM_RT_SEGUE_READ_BASE() __builtin_ia32_rdgsbase64()
+#define WASM_RT_SEGUE_WRITE_BASE(base) \
+  __builtin_ia32_wrgsbase64((uintptr_t)base)
+#define MEM_ADDR_MEMOP(mem, addr, n) ((uint8_t __seg_gs*)(uintptr_t)addr)
+#else
+#define MEM_ADDR_MEMOP(mem, addr, n) MEM_ADDR(mem, addr, n)
+#endif
+
 #define TRAP(x) (wasm_rt_trap(WASM_RT_TRAP_##x), 0)
 
 #if WASM_RT_STACK_DEPTH_COUNT
@@ -165,10 +199,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__
@@ -207,20 +254,22 @@ static inline void load_data(void* dest, const void* src, size_t n) {
     load_data(MEM_ADDR(&m, o, s), i, s); \
   } while (0)
 
-#define DEFINE_LOAD(name, t1, t2, t3, force_read)                         \
-  static inline t3 name(wasm_rt_memory_t* mem, u64 addr) {                \
-    MEMCHECK(mem, addr, t1);                                              \
-    t1 result;                                                            \
-    wasm_rt_memcpy(&result, MEM_ADDR(mem, addr, sizeof(t1)), sizeof(t1)); \
-    force_read(result);                                                   \
-    return (t3)(t2)result;                                                \
+#define DEFINE_LOAD(name, t1, t2, t3, force_read)                  \
+  static inline t3 name(wasm_rt_memory_t* mem, u64 addr) {         \
+    MEMCHECK(mem, addr, t1);                                       \
+    t1 result;                                                     \
+    wasm_rt_memcpy(&result, MEM_ADDR_MEMOP(mem, addr, sizeof(t1)), \
+                   sizeof(t1));                                    \
+    force_read(result);                                            \
+    return (t3)(t2)result;                                         \
   }
 
-#define DEFINE_STORE(name, t1, t2)                                         \
-  static inline void name(wasm_rt_memory_t* mem, u64 addr, t2 value) {     \
-    MEMCHECK(mem, addr, t1);                                               \
-    t1 wrapped = (t1)value;                                                \
-    wasm_rt_memcpy(MEM_ADDR(mem, addr, sizeof(t1)), &wrapped, sizeof(t1)); \
+#define DEFINE_STORE(name, t1, t2)                                     \
+  static inline void name(wasm_rt_memory_t* mem, u64 addr, t2 value) { \
+    MEMCHECK(mem, addr, t1);                                           \
+    t1 wrapped = (t1)value;                                            \
+    wasm_rt_memcpy(MEM_ADDR_MEMOP(mem, addr, sizeof(t1)), &wrapped,    \
+                   sizeof(t1));                                        \
   }
 
 DEFINE_LOAD(i32_load, u32, u32, u32, FORCE_READ_INT)
@@ -767,7 +816,14 @@ wasm_rt_memory_t* w2c_test_memory(w2c_test* instance) {
 
 /* export: '_start' */
 void w2c_test_0x5Fstart(w2c_test* instance) {
+#if WASM_RT_USE_SEGUE
+  uintptr_t segue_saved_base = WASM_RT_SEGUE_READ_BASE();
+  WASM_RT_SEGUE_WRITE_BASE(instance->w2c_memory.data);
+#endif
   w2c_test_0x5Fstart_0(instance);
+#if WASM_RT_USE_SEGUE
+  WASM_RT_SEGUE_WRITE_BASE(segue_saved_base);
+#endif
 }
 
 static void init_instance_import(w2c_test* instance, struct w2c_wasi__snapshot__preview1* w2c_wasi__snapshot__preview1_instance) {
@@ -779,8 +835,15 @@ void wasm2c_test_instantiate(w2c_test* instance, struct w2c_wasi__snapshot__prev
   init_instance_import(instance, w2c_wasi__snapshot__preview1_instance);
   init_tables(instance);
   init_memories(instance);
+#if WASM_RT_USE_SEGUE
+  uintptr_t segue_saved_base = WASM_RT_SEGUE_READ_BASE();
+  WASM_RT_SEGUE_WRITE_BASE(instance->w2c_memory.data);
+#endif
   init_elem_instances(instance);
   init_data_instances(instance);
+#if WASM_RT_USE_SEGUE
+  WASM_RT_SEGUE_WRITE_BASE(segue_saved_base);
+#endif
 }
 
 void wasm2c_test_free(w2c_test* instance) {