wasm2c: Add segue option to make exclusive use of the segment register

14 files changed, 154 insertions, 31 deletions

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 3a4442b3..92b5a088 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -203,7 +203,7 @@ jobs:
     env:
       USE_NINJA: "1"
       WASM2C_CC: "clang"
-      WASM2C_CFLAGS: "-DWASM_RT_USE_MMAP=1 -DWASM_RT_SKIP_SIGNAL_RECOVERY=1 -DWASM_RT_NONCONFORMING_UNCHECKED_STACK_EXHAUSTION=1 -DWASM2C_TEST_EMBEDDER_SIGNAL_HANDLING -DWASM_RT_ALLOW_SEGUE=1 -mfsgsbase -DWASM_RT_SANITY_CHECKS=1 -Wno-pass-failed"
+      WASM2C_CFLAGS: "-DWASM_RT_USE_MMAP=1 -DWASM_RT_SKIP_SIGNAL_RECOVERY=1 -DWASM_RT_NONCONFORMING_UNCHECKED_STACK_EXHAUSTION=1 -DWASM2C_TEST_EMBEDDER_SIGNAL_HANDLING -DWASM_RT_ALLOW_SEGUE=1 -DWASM_RT_SEGUE_FREE_SEGMENT=1 -mfsgsbase -DWASM_RT_SANITY_CHECKS=1 -Wno-pass-failed"
     steps:
     - uses: actions/setup-python@v1
       with:
diff --git a/src/c-writer.cc b/src/c-writer.cc
index 986ace55..cbffa1d4 100644
--- a/src/c-writer.cc
+++ b/src/c-writer.cc
@@ -2449,7 +2449,9 @@ bool CWriter::IsSingleUnsharedMemory() {
 void CWriter::InstallSegueBase(Memory* memory, bool save_old_value) {
   NonIndented([&] { Write("#if WASM_RT_USE_SEGUE", Newline()); });
   if (save_old_value) {
+    NonIndented([&] { Write("#if !WASM_RT_SEGUE_FREE_SEGMENT", Newline()); });
     Write("void* segue_saved_base = wasm_rt_segue_read_base();", Newline());
+    NonIndented([&] { Write("#endif", Newline()); });
   }
   auto primary_memory =
       ExternalInstanceRef(ModuleFieldType::Memory, memory->name);
@@ -2458,7 +2460,9 @@ void CWriter::InstallSegueBase(Memory* memory, bool save_old_value) {
 }
 
 void CWriter::RestoreSegueBase() {
-  NonIndented([&] { Write("#if WASM_RT_USE_SEGUE", Newline()); });
+  NonIndented([&] {
+    Write("#if WASM_RT_USE_SEGUE && !WASM_RT_SEGUE_FREE_SEGMENT", Newline());
+  });
   Write("wasm_rt_segue_write_base(segue_saved_base);", Newline());
   NonIndented([&] { Write("#endif", Newline()); });
 }
diff --git a/src/prebuilt/wasm2c_source_declarations.cc b/src/prebuilt/wasm2c_source_declarations.cc
index 000b25c8..86c07ded 100644
--- a/src/prebuilt/wasm2c_source_declarations.cc
+++ b/src/prebuilt/wasm2c_source_declarations.cc
@@ -115,7 +115,21 @@ R"w2c_template(}
 )w2c_template"
 R"w2c_template(static inline void wasm_rt_segue_write_base(void* base) {
 )w2c_template"
-R"w2c_template(  if (wasm_rt_fsgsbase_inst_supported) {
+R"w2c_template(#if WASM_RT_SEGUE_FREE_SEGMENT
+)w2c_template"
+R"w2c_template(  if (wasm_rt_last_segment_val == base) {
+)w2c_template"
+R"w2c_template(    return;
+)w2c_template"
+R"w2c_template(  }
+)w2c_template"
+R"w2c_template(
+  wasm_rt_last_segment_val = base;
+)w2c_template"
+R"w2c_template(#endif
+)w2c_template"
+R"w2c_template(
+  if (wasm_rt_fsgsbase_inst_supported) {
 )w2c_template"
 R"w2c_template(    __builtin_ia32_wrgsbase64((uintptr_t)base);
 )w2c_template"
diff --git a/src/template/wasm2c.declarations.c b/src/template/wasm2c.declarations.c
index efe9dbbe..43cd1d1b 100644
--- a/src/template/wasm2c.declarations.c
+++ b/src/template/wasm2c.declarations.c
@@ -58,6 +58,14 @@ static inline void* wasm_rt_segue_read_base() {
   }
 }
 static inline void wasm_rt_segue_write_base(void* base) {
+#if WASM_RT_SEGUE_FREE_SEGMENT
+  if (wasm_rt_last_segment_val == base) {
+    return;
+  }
+
+  wasm_rt_last_segment_val = base;
+#endif
+
   if (wasm_rt_fsgsbase_inst_supported) {
     __builtin_ia32_wrgsbase64((uintptr_t)base);
   } else {
diff --git a/test/wasm2c/add.txt b/test/wasm2c/add.txt
index 7a576d9c..db581f49 100644
--- a/test/wasm2c/add.txt
+++ b/test/wasm2c/add.txt
@@ -125,6 +125,14 @@ static inline void* wasm_rt_segue_read_base() {
   }
 }
 static inline void wasm_rt_segue_write_base(void* base) {
+#if WASM_RT_SEGUE_FREE_SEGMENT
+  if (wasm_rt_last_segment_val == base) {
+    return;
+  }
+
+  wasm_rt_last_segment_val = base;
+#endif
+
   if (wasm_rt_fsgsbase_inst_supported) {
     __builtin_ia32_wrgsbase64((uintptr_t)base);
   } else {
diff --git a/test/wasm2c/check-imports.txt b/test/wasm2c/check-imports.txt
index a5d3ae61..f2fe508e 100644
--- a/test/wasm2c/check-imports.txt
+++ b/test/wasm2c/check-imports.txt
@@ -149,6 +149,14 @@ static inline void* wasm_rt_segue_read_base() {
   }
 }
 static inline void wasm_rt_segue_write_base(void* base) {
+#if WASM_RT_SEGUE_FREE_SEGMENT
+  if (wasm_rt_last_segment_val == base) {
+    return;
+  }
+
+  wasm_rt_last_segment_val = base;
+#endif
+
   if (wasm_rt_fsgsbase_inst_supported) {
     __builtin_ia32_wrgsbase64((uintptr_t)base);
   } else {
@@ -825,11 +833,13 @@ void wasm2c_test_instantiate(w2c_test* instance, struct w2c_env* w2c_env_instanc
   init_tables(instance);
   init_memories(instance);
 #if WASM_RT_USE_SEGUE
+#if !WASM_RT_SEGUE_FREE_SEGMENT
   void* segue_saved_base = wasm_rt_segue_read_base();
+#endif
   wasm_rt_segue_write_base((*instance->w2c_env_0x5F_linear_memory).data);
 #endif
   init_elem_instances(instance);
-#if WASM_RT_USE_SEGUE
+#if WASM_RT_USE_SEGUE && !WASM_RT_SEGUE_FREE_SEGMENT
   wasm_rt_segue_write_base(segue_saved_base);
 #endif
 }
diff --git a/test/wasm2c/export-names.txt b/test/wasm2c/export-names.txt
index f5762a19..a58bee0a 100644
--- a/test/wasm2c/export-names.txt
+++ b/test/wasm2c/export-names.txt
@@ -149,6 +149,14 @@ static inline void* wasm_rt_segue_read_base() {
   }
 }
 static inline void wasm_rt_segue_write_base(void* base) {
+#if WASM_RT_SEGUE_FREE_SEGMENT
+  if (wasm_rt_last_segment_val == base) {
+    return;
+  }
+
+  wasm_rt_last_segment_val = base;
+#endif
+
   if (wasm_rt_fsgsbase_inst_supported) {
     __builtin_ia32_wrgsbase64((uintptr_t)base);
   } else {
@@ -796,11 +804,13 @@ static void init_memories(w2c_test* instance) {
 /* export: '' */
 void w2c_test_(w2c_test* instance) {
 #if WASM_RT_USE_SEGUE
+#if !WASM_RT_SEGUE_FREE_SEGMENT
   void* segue_saved_base = wasm_rt_segue_read_base();
+#endif
   wasm_rt_segue_write_base((*instance->w2c_0x5Cmodule_import0x200x2A0x2F).data);
 #endif
   w2c_test__0(instance);
-#if WASM_RT_USE_SEGUE
+#if WASM_RT_USE_SEGUE && !WASM_RT_SEGUE_FREE_SEGMENT
   wasm_rt_segue_write_base(segue_saved_base);
 #endif
 }
@@ -808,11 +818,13 @@ void w2c_test_(w2c_test* instance) {
 /* export: '*\2F' */
 void w2c_test_0x2A0x2F(w2c_test* instance) {
 #if WASM_RT_USE_SEGUE
+#if !WASM_RT_SEGUE_FREE_SEGMENT
   void* segue_saved_base = wasm_rt_segue_read_base();
+#endif
   wasm_rt_segue_write_base((*instance->w2c_0x5Cmodule_import0x200x2A0x2F).data);
 #endif
   w2c_test__0(instance);
-#if WASM_RT_USE_SEGUE
+#if WASM_RT_USE_SEGUE && !WASM_RT_SEGUE_FREE_SEGMENT
   wasm_rt_segue_write_base(segue_saved_base);
 #endif
 }
@@ -820,11 +832,13 @@ void w2c_test_0x2A0x2F(w2c_test* instance) {
 /* export: '\3F\3F\2F' */
 void w2c_test_0x3F0x3F0x2F(w2c_test* instance) {
 #if WASM_RT_USE_SEGUE
+#if !WASM_RT_SEGUE_FREE_SEGMENT
   void* segue_saved_base = wasm_rt_segue_read_base();
+#endif
   wasm_rt_segue_write_base((*instance->w2c_0x5Cmodule_import0x200x2A0x2F).data);
 #endif
   w2c_test__0(instance);
-#if WASM_RT_USE_SEGUE
+#if WASM_RT_USE_SEGUE && !WASM_RT_SEGUE_FREE_SEGMENT
   wasm_rt_segue_write_base(segue_saved_base);
 #endif
 }
@@ -832,11 +846,13 @@ void w2c_test_0x3F0x3F0x2F(w2c_test* instance) {
 /* export: '\0A' */
 void w2c_test_0x0A(w2c_test* instance) {
 #if WASM_RT_USE_SEGUE
+#if !WASM_RT_SEGUE_FREE_SEGMENT
   void* segue_saved_base = wasm_rt_segue_read_base();
+#endif
   wasm_rt_segue_write_base((*instance->w2c_0x5Cmodule_import0x200x2A0x2F).data);
 #endif
   w2c_test__0(instance);
-#if WASM_RT_USE_SEGUE
+#if WASM_RT_USE_SEGUE && !WASM_RT_SEGUE_FREE_SEGMENT
   wasm_rt_segue_write_base(segue_saved_base);
 #endif
 }
@@ -844,11 +860,13 @@ void w2c_test_0x0A(w2c_test* instance) {
 /* export: '\E2\9D\A4\EF\B8\8F' */
 void w2c_test_0xE20x9D0xA40xEF0xB80x8F(w2c_test* instance) {
 #if WASM_RT_USE_SEGUE
+#if !WASM_RT_SEGUE_FREE_SEGMENT
   void* segue_saved_base = wasm_rt_segue_read_base();
+#endif
   wasm_rt_segue_write_base((*instance->w2c_0x5Cmodule_import0x200x2A0x2F).data);
 #endif
   w2c_test__0(instance);
-#if WASM_RT_USE_SEGUE
+#if WASM_RT_USE_SEGUE && !WASM_RT_SEGUE_FREE_SEGMENT
   wasm_rt_segue_write_base(segue_saved_base);
 #endif
 }
@@ -866,10 +884,12 @@ void wasm2c_test_instantiate(w2c_test* instance, struct w2c_0x5Cmodule* w2c_0x5C
   init_instance_import(instance, w2c_0x5Cmodule_instance);
   init_memories(instance);
 #if WASM_RT_USE_SEGUE
+#if !WASM_RT_SEGUE_FREE_SEGMENT
   void* segue_saved_base = wasm_rt_segue_read_base();
+#endif
   wasm_rt_segue_write_base((*instance->w2c_0x5Cmodule_import0x200x2A0x2F).data);
 #endif
-#if WASM_RT_USE_SEGUE
+#if WASM_RT_USE_SEGUE && !WASM_RT_SEGUE_FREE_SEGMENT
   wasm_rt_segue_write_base(segue_saved_base);
 #endif
 }
diff --git a/test/wasm2c/hello.txt b/test/wasm2c/hello.txt
index a9beeede..2a0d79f0 100644
--- a/test/wasm2c/hello.txt
+++ b/test/wasm2c/hello.txt
@@ -157,6 +157,14 @@ static inline void* wasm_rt_segue_read_base() {
   }
 }
 static inline void wasm_rt_segue_write_base(void* base) {
+#if WASM_RT_SEGUE_FREE_SEGMENT
+  if (wasm_rt_last_segment_val == base) {
+    return;
+  }
+
+  wasm_rt_last_segment_val = base;
+#endif
+
   if (wasm_rt_fsgsbase_inst_supported) {
     __builtin_ia32_wrgsbase64((uintptr_t)base);
   } else {
@@ -833,11 +841,13 @@ wasm_rt_memory_t* w2c_test_memory(w2c_test* instance) {
 /* export: '_start' */
 void w2c_test_0x5Fstart(w2c_test* instance) {
 #if WASM_RT_USE_SEGUE
+#if !WASM_RT_SEGUE_FREE_SEGMENT
   void* segue_saved_base = wasm_rt_segue_read_base();
+#endif
   wasm_rt_segue_write_base(instance->w2c_memory.data);
 #endif
   w2c_test_0x5Fstart_0(instance);
-#if WASM_RT_USE_SEGUE
+#if WASM_RT_USE_SEGUE && !WASM_RT_SEGUE_FREE_SEGMENT
   wasm_rt_segue_write_base(segue_saved_base);
 #endif
 }
@@ -852,12 +862,14 @@ void wasm2c_test_instantiate(w2c_test* instance, struct w2c_wasi__snapshot__prev
   init_tables(instance);
   init_memories(instance);
 #if WASM_RT_USE_SEGUE
+#if !WASM_RT_SEGUE_FREE_SEGMENT
   void* segue_saved_base = wasm_rt_segue_read_base();
+#endif
   wasm_rt_segue_write_base(instance->w2c_memory.data);
 #endif
   init_elem_instances(instance);
   init_data_instances(instance);
-#if WASM_RT_USE_SEGUE
+#if WASM_RT_USE_SEGUE && !WASM_RT_SEGUE_FREE_SEGMENT
   wasm_rt_segue_write_base(segue_saved_base);
 #endif
 }
diff --git a/test/wasm2c/minimal.txt b/test/wasm2c/minimal.txt
index 3fb8ad8b..7ffa3bca 100644
--- a/test/wasm2c/minimal.txt
+++ b/test/wasm2c/minimal.txt
@@ -119,6 +119,14 @@ static inline void* wasm_rt_segue_read_base() {
   }
 }
 static inline void wasm_rt_segue_write_base(void* base) {
+#if WASM_RT_SEGUE_FREE_SEGMENT
+  if (wasm_rt_last_segment_val == base) {
+    return;
+  }
+
+  wasm_rt_last_segment_val = base;
+#endif
+
   if (wasm_rt_fsgsbase_inst_supported) {
     __builtin_ia32_wrgsbase64((uintptr_t)base);
   } else {
diff --git a/test/wasm2c/tail-calls.txt b/test/wasm2c/tail-calls.txt
index c0e940f4..b27d7417 100644
--- a/test/wasm2c/tail-calls.txt
+++ b/test/wasm2c/tail-calls.txt
@@ -149,6 +149,14 @@ static inline void* wasm_rt_segue_read_base() {
   }
 }
 static inline void wasm_rt_segue_write_base(void* base) {
+#if WASM_RT_SEGUE_FREE_SEGMENT
+  if (wasm_rt_last_segment_val == base) {
+    return;
+  }
+
+  wasm_rt_last_segment_val = base;
+#endif
+
   if (wasm_rt_fsgsbase_inst_supported) {
     __builtin_ia32_wrgsbase64((uintptr_t)base);
   } else {
diff --git a/wasm2c/README.md b/wasm2c/README.md
index 6d8b9999..f032b4e2 100644
--- a/wasm2c/README.md
+++ b/wasm2c/README.md
@@ -151,8 +151,9 @@ high performance overhead, and is thus only recommended for debug builds.
 
 Wasm2c can use the "Segue" optimization if allowed. The segue optimization uses
 an x86 segment register to store the location of Wasm's linear memory, when
-compiling a Wasm module with clang, running on x86_64 Linux, and the macro
-`WASM_RT_ALLOW_SEGUE` is defined. Segue is not used if
+compiling a Wasm module with clang, running on x86_64 Linux, the macro
+`WASM_RT_ALLOW_SEGUE` is defined, and the flag `-mfsgsbase` is passed to clang.
+Segue is not used if
 
 1. The Wasm module uses a more than a single unshared imported or exported
    memory
@@ -173,6 +174,12 @@ However, any host functions written in assembly that clobber the free segment
 register must restore the value of this register prior to executing or returning
 control to wasm2c generated code.
 
+As an additional optimization, if the host program does not use the `%gs`
+segment register for any other purpose (which is typically the case in most
+programs), you can additionally allow wasm2c to unconditionally overwrite the
+value of the `%gs` register without restoring the old value. This can be done
+defining the macro `WASM_RT_SEGUE_FREE_SEGMENT`.
+
 You can test the performance of the Segue optimization by running Dhrystone with
 and without Segue:
 
diff --git a/wasm2c/examples/fac/fac.c b/wasm2c/examples/fac/fac.c
index e253cc2a..38c5e823 100644
--- a/wasm2c/examples/fac/fac.c
+++ b/wasm2c/examples/fac/fac.c
@@ -51,11 +51,16 @@
 //     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
+// (5) The OS provides a way to query if (rd|wr)gsbase is allowed by the kernel
+// or the implementation has to use a syscall for this.
+// (6) The OS doesn't replace the segment register on context switch which
 //     eliminates windows for now
+//
+// While more OS can be supported in the future, we only support linux 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)
+    __clang__ && __has_builtin(__builtin_ia32_wrgsbase64) &&                 \
+    !defined(_WIN32) && defined(__linux__)
 #define WASM_RT_USE_SEGUE 1
 #else
 #define WASM_RT_USE_SEGUE 0
@@ -64,31 +69,26 @@
 
 #if WASM_RT_USE_SEGUE
 // 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;
+    return wasm_rt_syscall_get_segue_base();
   }
 }
 static inline void wasm_rt_segue_write_base(void* base) {
+#if WASM_RT_SEGUE_FREE_SEGMENT
+  if (wasm_rt_last_segment_val == base) {
+    return;
+  }
+
+  wasm_rt_last_segment_val = base;
+#endif
+
   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();
-    }
+    wasm_rt_syscall_set_segue_base(base);
   }
 }
 #define MEM_ADDR_MEMOP(mem, addr, n) ((uint8_t __seg_gs*)(uintptr_t)addr)
diff --git a/wasm2c/wasm-rt-impl.c b/wasm2c/wasm-rt-impl.c
index 851a4c5a..36bb9690 100644
--- a/wasm2c/wasm-rt-impl.c
+++ b/wasm2c/wasm-rt-impl.c
@@ -62,6 +62,10 @@ bool wasm_rt_fsgsbase_inst_supported = false;
 #include <unistd.h>       // For syscall
 #endif
 
+#if WASM_RT_SEGUE_FREE_SEGMENT
+WASM_RT_THREAD_LOCAL void* wasm_rt_last_segment_val = NULL;
+#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;
diff --git a/wasm2c/wasm-rt.h b/wasm2c/wasm-rt.h
index c489c850..b72d12df 100644
--- a/wasm2c/wasm-rt.h
+++ b/wasm2c/wasm-rt.h
@@ -217,6 +217,18 @@ extern "C" {
 #endif
 
 /**
+ * The segue optimization restores x86 segments to their old values when exiting
+ * wasm2c code. If WASM_RT_SEGUE_FREE_SEGMENT is defined, wasm2c assumes it has
+ * exclusive use of the segment and optimizes performance in two ways. First, it
+ * does not restore the "old" value of the segment during exits. Second, wasm2c
+ * only sets the segment register if the value has changed since the last time
+ * it was set.
+ */
+#ifndef WASM_RT_SEGUE_FREE_SEGMENT
+#define WASM_RT_SEGUE_FREE_SEGMENT 0
+#endif
+
+/**
  * This macro, if defined, allows the embedder to disable all stack exhaustion
  * checks. This a non conformant configuration, i.e., this does not respect
  * Wasm's specification, and may compromise security. Use with caution.
@@ -313,6 +325,14 @@ void wasm_rt_syscall_set_segue_base(void* base);
  * a function that invokes the OS' underlying syscall to get the segment base.
  */
 void* wasm_rt_syscall_get_segue_base();
+/**
+ * If WASM_RT_SEGUE_FREE_SEGMENT is defined, we must only set the segment
+ * register if it was changed since the last time it was set. The last value set
+ * on the segment register is stored in this variable.
+ */
+#if WASM_RT_SEGUE_FREE_SEGMENT
+extern WASM_RT_THREAD_LOCAL void* wasm_rt_last_segment_val;
+#endif
 #endif
 
 #if defined(_MSC_VER)