Refactor stack management on wasip3

.github/workflows/main.yml

@@ -128,7 +128,7 @@ jobs:
             clang_version: 22
             test: true
             upload: wasm32-wasip3
-            args: -DTARGET_TRIPLE=wasm32-wasip3
+            args: -DTARGET_TRIPLE=wasm32-wasip3 -DCHECK_SYMBOLS=OFF
 
           - name: Test wasm32-wasip3 (debug)
             os: ubuntu-24.04

CMakeLists.txt

@@ -204,17 +204,6 @@ add_link_options(
   -resource-dir ${tmp_resource_dir}
 )
 
-if (HAVE_WASM_LIBCALL_THREAD_CONTEXT)
-  # TODO: shouldn't have to pass these as it means it's required by all end
-  # users as well. Ideally this would be part of the driver in rustc/clang, but
-  # that's left for a future change.
-  add_link_options(
-    -Wl,--export-if-defined=__wasm_init_tls
-    -Wl,--export-if-defined=__tls_size
-    -Wl,--export-if-defined=__tls_align
-  )
-endif()
-
 # Expose the public headers to the implementation. We use `-isystem` for
 # purpose for two reasons:
 #

cmake/bindings.cmake

@@ -152,7 +152,6 @@ function(wit_bindgen_edit p)
     COMMAND sed ${SED_INPLACE_ARGS} "'s_#include .wasi${p}\.h._#include \"wasi/wasi${p}.h\"_'" ${bottom_half}/sources/wasi${p}.c
     COMMAND sed ${SED_INPLACE_ARGS} "s/extern void exit_exit/_Noreturn extern void exit_exit/" ${bottom_half}/headers/public/wasi/__generated_wasi${p}.h
     COMMAND sed ${SED_INPLACE_ARGS} "s/extern void __wasm_import_exit_exit/_Noreturn extern void __wasm_import_exit_exit/" ${bottom_half}/sources/wasi${p}.c
-    COMMAND sed ${SED_INPLACE_ARGS} "s/__attribute__.*\"cabi_realloc\".*/#include \\\"cabi_realloc_augment.h\\\"/" ${bottom_half}/sources/wasi${p}.c
     DEPENDS bindings-${p}
   )
   add_dependencies(bindings bindings-${p}-edit)

cmake/wasm-component-ld.cmake

@@ -3,7 +3,7 @@ if (NOT USE_WASM_COMPONENT_LD)
   return()
 endif()
 
-set(WASM_COMPONENT_LD_REQUIRED_VERSION "0.5.29")
+set(WASM_COMPONENT_LD_REQUIRED_VERSION "0.5.30")
 
 find_program(WASM_COMPONENT_LD_EXECUTABLE NAMES wasm-component-ld)
 

cmake/wasm-tools.cmake

@@ -5,7 +5,7 @@ if (NOT WASM_TOOLS_EXECUTABLE)
   ba_download(
     wasm-tools
     "https://github.com/bytecodealliance/wasm-tools"
-    "1.256.0"
+    "1.257.0"
   )
   ExternalProject_Get_Property(wasm-tools SOURCE_DIR)
   set(wasm_tools "${SOURCE_DIR}/wasm-tools")

expected/wasm32-wasip3-coop/defined-symbols.txt

@@ -387,14 +387,14 @@ __wasilibc_rename_newat
 __wasilibc_rename_oldat
 __wasilibc_reset_preopens
 __wasilibc_rmdirat
-__wasilibc_set_task_tls
 __wasilibc_setsockopt_timeout
 __wasilibc_sockaddr_to_wasi
 __wasilibc_sockaddr_validate
 __wasilibc_stat
 __wasilibc_stream_block_on_timeout
 __wasilibc_subtask_await_nonblocking
 __wasilibc_subtask_block_on_and_drop
+__wasilibc_task_hook
 __wasilibc_tell
 __wasilibc_unlinkat
 __wasilibc_unspecified_addr
@@ -406,6 +406,8 @@ __wasilibc_write
 __wasilibc_write_poll
 __wasilibc_write_state_close
 __wasm_call_dtors
+__wasm_library_tls_info
+__wasm_task_hook
 __wcscoll_l
 __wcsftime_l
 __wcsxfrm_l

expected/wasm32-wasip3/defined-symbols.txt

@@ -368,6 +368,7 @@ __wasilibc_stat
 __wasilibc_stream_block_on_timeout
 __wasilibc_subtask_await_nonblocking
 __wasilibc_subtask_block_on_and_drop
+__wasilibc_task_hook
 __wasilibc_tell
 __wasilibc_unlinkat
 __wasilibc_unspecified_addr
@@ -379,6 +380,8 @@ __wasilibc_write
 __wasilibc_write_poll
 __wasilibc_write_state_close
 __wasm_call_dtors
+__wasm_library_tls_info
+__wasm_task_hook
 __wcscoll_l
 __wcsftime_l
 __wcsxfrm_l

libc-bottom-half/CMakeLists.txt

@@ -184,10 +184,7 @@ if (WASI STREQUAL "p3")
 endif()
 
 if (HAVE_WASM_LIBCALL_THREAD_CONTEXT)
-  list(APPEND bottom_half_sources
-    sources/wasip3_tls.c
-    sources/__wasm_init_task.S
-  )
+  list(APPEND bottom_half_sources sources/__wasm_task_hook.S)
 endif()
 
 # Don't export symbols of generated code in shared libraries, so specifically

libc-bottom-half/crt/wasip3_symbol_references.h

@@ -5,16 +5,38 @@
 #ifdef __wasm_libcall_thread_context__
 
 #include <wasi/version.h>
+#include <wasi/wasip3_tls.h>
 
 // Force some symbols to be linked in for wasip3
-extern void __wasm_init_task(void);
-extern void __wasm_init_async_task(void);
+extern void __wasm_task_hook(uint32_t);
 extern void cabi_realloc(void);
-__attribute__((used)) static void *__wasm_init_task_ref = __wasm_init_task;
-__attribute__((used)) static void *__wasm_init_async_task_ref =
-    __wasm_init_async_task;
+__attribute__((used)) static void *__wasm_task_hook_ref = __wasm_task_hook;
 __attribute__((used)) static void *cabi_realloc_ref = cabi_realloc;
 
+// Export a symbol from this object which describes the TLS information
+// required for this object. This encapsulates the size/align information as
+// well as an initialization function.
+//
+// Note that being part of `crt1-*.o` this is exported from all libraries,
+// which is what we want because each library has a unique return value from
+// `__builtin_wasm_tls_*` intrinsics and a separate `__wasm_init_tls` function.
+//
+// For more information on this see
+// https://github.com/WebAssembly/wasi-libc/issues/857 which has since been
+// adjusted with this scheme.
+static size_t tls_size_and_align(size_t *align) {
+  *align = __builtin_wasm_tls_align();
+  return __builtin_wasm_tls_size();
+}
+
+void __wasm_init_tls(void *base);
+
+__attribute__((visibility("default")))
+struct __wasilibc_library_tls_info __wasm_library_tls_info = {
+  .tls_size_and_align = tls_size_and_align,
+  .init_tls = __wasm_init_tls,
+};
+
 // Force `__wasm_{g,s}et_{stack_pointer,tls_base}` to exist as defined symbols.
 // These end up as imported functions which `wit-component` recognizes, and what
 // exactly they're hooked up to will depend on `wit-component` when this is

libc-bottom-half/sources/__cabi_realloc_wrapper.S

@@ -1,129 +0,0 @@
-#include <wasi/version.h>
-
-#ifdef __wasm_libcall_thread_context__
-// The size, in bytes, of the stack statically allocated for `cabi_realloc`.
-// This doesn't need to be too big as it's just running `cabi_realloc` as defined
-// in wasi-libc. The current number is determined by:
-//
-// * Start with 16 bytes.
-// * Double until tests don't crash in a debug build.
-// * Double again.
-//
-// The goal is to be pretty modest since this is part of every module.
-#define REALLOC_STACK_SIZE 512
-
-// Declare the `__wasm_{g,s}et_stack_pointer` functions.
-.functype   __wasm_set_stack_pointer (i32) -> ()
-.functype   __wasm_get_stack_pointer () -> (i32)
-.type       __wasm_set_stack_pointer,@function
-.type       __wasm_get_stack_pointer,@function
-.functype   __cabi_realloc_impl (i32, i32, i32, i32) -> (i32)
-.type       __cabi_realloc_impl,@function
-
-// Go ahead and generate an `-fPIC`-compatible addressing mode for
-// `cabi_realloc_stack` so this is compatible with both PIC and not builds.
-.type        __memory_base,@global
-.globaltype  __memory_base, i32, immutable
-
-// For coop-threads builds we need to configure TLS to the initial TLS
-// value for all tasks/etc, so import those symbols here.
-//
-// Note that this saves and restores context slot 1 itself rather than this
-// module's TLS base. When several modules are linked together the slot holds an
-// array of per-module base pointers, so clobbering just libc's entry would both
-// leave the other modules pointing at the wrong TLS and scribble through
-// whatever stale pointer the slot happened to contain.
-#ifdef __wasi_cooperative_threads__
-.functype   wasip3_context_set_1 (i32) -> ()
-.functype   wasip3_context_get_1 () -> (i32)
-.functype   __wasilibc_set_task_tls (i32) -> ()
-.type       wasip3_context_set_1,@function
-.type       wasip3_context_get_1,@function
-.type       __wasilibc_set_task_tls,@function
-.import_module wasip3_context_get_1, "$root"
-.import_name wasip3_context_get_1, "[context-get-1]"
-.import_module wasip3_context_set_1, "$root"
-.import_name wasip3_context_set_1, "[context-set-1]"
-.type       __init_tls_base,@global
-#ifdef __PIC__
-.globaltype __init_tls_base, i32
-#else
-.globaltype __init_tls_base, i32, immutable
-#endif
-#endif
-
-// This wrapper takes over the `cabi_realloc` symbol itself and the
-// `cabi_realloc` core module export. The C implementation it forwards to is
-// renamed to `__cabi_realloc_impl` by `cabi_realloc_augment.h`.
-//
-// Note that it's important that the symbol name matches the export name here.
-// In a `BUILD_SHARED=ON` build `libc.so` is dynamically linked against by other
-// shared libraries and by `crt1-{command,reactor}.o`, and the name a symbol is
-// resolved by in that context is its wasm export name. If `.export_name` were
-// used to rename the export then this function would be unresolvable by name
-// from anything linking against `libc.so`.
-.globl cabi_realloc
-.export_name cabi_realloc, "cabi_realloc"
-.type cabi_realloc,@function
-cabi_realloc:
-  .functype cabi_realloc (i32, i32, i32, i32) -> (i32)
-
-  // Prior to bytecodealliance/wasmtime#13949 wasmtime's handling of
-  // `cabi_realloc` and context slots was a bit buggy. That means that prior
-  // to that PR the slots aren't guaranteed to be 0 and the value present will
-  // be used by some future task. This export is going to clobber these
-  // slots, so while that Wasmtime change is percolating this preserves the
-  // slots around the invocation of `cabi_realloc`. Once #13949 percolates and
-  // ships then this workaround code, and the restore down below, can be
-  // deleted.
-  .local i32, i32
-  call __wasm_get_stack_pointer
-  local.set 4
-#ifdef __wasi_cooperative_threads__
-  call wasip3_context_get_1
-  local.set 5
-#endif
-
-  // First up, configure the stack pointer. Get the base of the stack from the
-  // data symbol defined below, add the stack size, and that's what we're
-  // starting with.
-  i32.const cabi_realloc_stack@MBREL
-  global.get __memory_base
-  i32.add
-  i32.const  REALLOC_STACK_SIZE
-  i32.add
-  call __wasm_set_stack_pointer
-
-  // Next up configure TLS. This is only required in coop-threads builds.
-#ifdef __wasi_cooperative_threads__
-  global.get   __init_tls_base
-  call __wasilibc_set_task_tls
-#endif
-
-  // And finally forward to the actual implementation of `cabi_realloc`
-  // defined in `wasip3.c`
-  local.get 0
-  local.get 1
-  local.get 2
-  local.get 3
-  call __cabi_realloc_impl
-
-  local.get 4
-  call __wasm_set_stack_pointer
-#ifdef __wasi_cooperative_threads__
-  local.get 5
-  call wasip3_context_set_1
-#endif
-
-  end_function
-
-// Define a data symbol which is the stack that `cabi_realloc` runs on.
-// Note that this is aligned to 16-bytes, the expected stack alignment.
-.type cabi_realloc_stack,@object
-.section .bss.cabi_realloc_stack,"",@
-.p2align 4, 0x0
-cabi_realloc_stack:
-  .skip  REALLOC_STACK_SIZE
-  .size cabi_realloc_stack,  REALLOC_STACK_SIZE
-
-#endif // __wasm_libcall_thread_context__