Disable aux stack allocations for threads spawned by wasi_thread_start

core/iwasm/common/wasm_exec_env.h

@@ -179,6 +179,13 @@ wasm_exec_env_create(struct WASMModuleInstanceCommon *module_inst,
 void
 wasm_exec_env_destroy(WASMExecEnv *exec_env);
 
+static inline bool
+wasm_exec_env_is_aux_stack_managed_by_runtime(WASMExecEnv *exec_env)
+{
+    return exec_env->aux_stack_boundary.boundary != 0
+           || exec_env->aux_stack_bottom.bottom != 0;
+}
+
 /**
  * Allocate a WASM frame from the WASM stack.
  *

core/iwasm/interpreter/wasm_interp_classic.c

@@ -1783,14 +1783,18 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                 global = globals + global_idx;
                 global_addr = get_global_addr(global_data, global);
                 aux_stack_top = *(uint32 *)(frame_sp - 1);
-                if (aux_stack_top <= exec_env->aux_stack_boundary.boundary) {
-                    wasm_set_exception(module, "wasm auxiliary stack overflow");
-                    goto got_exception;
-                }
-                if (aux_stack_top > exec_env->aux_stack_bottom.bottom) {
-                    wasm_set_exception(module,
-                                       "wasm auxiliary stack underflow");
-                    goto got_exception;
+                if (wasm_exec_env_is_aux_stack_managed_by_runtime(exec_env)) {
+                    if (aux_stack_top
+                        <= exec_env->aux_stack_boundary.boundary) {
+                        wasm_set_exception(module,
+                                           "wasm auxiliary stack overflow");
+                        goto got_exception;
+                    }
+                    if (aux_stack_top > exec_env->aux_stack_bottom.bottom) {
+                        wasm_set_exception(module,
+                                           "wasm auxiliary stack underflow");
+                        goto got_exception;
+                    }
                 }
                 *(int32 *)global_addr = aux_stack_top;
                 frame_sp--;

core/iwasm/interpreter/wasm_interp_fast.c

@@ -1576,14 +1576,18 @@ wasm_interp_call_func_bytecode(WASMModuleInstance *module,
                 global = globals + global_idx;
                 global_addr = get_global_addr(global_data, global);
                 aux_stack_top = frame_lp[GET_OFFSET()];
-                if (aux_stack_top <= exec_env->aux_stack_boundary.boundary) {
-                    wasm_set_exception(module, "wasm auxiliary stack overflow");
-                    goto got_exception;
-                }
-                if (aux_stack_top > exec_env->aux_stack_bottom.bottom) {
-                    wasm_set_exception(module,
-                                       "wasm auxiliary stack underflow");
-                    goto got_exception;
+                if (wasm_exec_env_is_aux_stack_managed_by_runtime(exec_env)) {
+                    if (aux_stack_top
+                        <= exec_env->aux_stack_boundary.boundary) {
+                        wasm_set_exception(module,
+                                           "wasm auxiliary stack overflow");
+                        goto got_exception;
+                    }
+                    if (aux_stack_top > exec_env->aux_stack_bottom.bottom) {
+                        wasm_set_exception(module,
+                                           "wasm auxiliary stack underflow");
+                        goto got_exception;
+                    }
                 }
                 *(int32 *)global_addr = aux_stack_top;
 #if WASM_ENABLE_MEMORY_PROFILING != 0

core/iwasm/libraries/lib-pthread/lib_pthread_wrapper.c

@@ -619,8 +619,9 @@ pthread_create_wrapper(wasm_exec_env_t exec_env,
     routine_args->module_inst = new_module_inst;
 
     os_mutex_lock(&exec_env->wait_lock);
-    ret = wasm_cluster_create_thread(
-        exec_env, new_module_inst, pthread_start_routine, (void *)routine_args);
+    ret =
+        wasm_cluster_create_thread(exec_env, new_module_inst, true,
+                                   pthread_start_routine, (void *)routine_args);
     if (ret != 0) {
         os_mutex_unlock(&exec_env->wait_lock);
         goto fail;

core/iwasm/libraries/lib-wasi-threads/lib_wasi_threads_wrapper.c

@@ -126,8 +126,8 @@ thread_spawn_wrapper(wasm_exec_env_t exec_env, uint32 start_arg)
     thread_start_arg->start_func = start_func;
 
     os_mutex_lock(&exec_env->wait_lock);
-    ret = wasm_cluster_create_thread(exec_env, new_module_inst, thread_start,
-                                     thread_start_arg);
+    ret = wasm_cluster_create_thread(exec_env, new_module_inst, false,
+                                     thread_start, thread_start_arg);
     if (ret != 0) {
         LOG_ERROR("Failed to spawn a new thread");
         goto thread_spawn_fail;

core/iwasm/libraries/thread-mgr/thread_manager.c

@@ -125,6 +125,10 @@ free_aux_stack(WASMExecEnv *exec_env, uint32 start)
     WASMModuleInstanceCommon *module_inst =
         wasm_exec_env_get_module_inst(exec_env);
 
+    if (!wasm_exec_env_is_aux_stack_managed_by_runtime(exec_env)) {
+        return true;
+    }
+
     bh_assert(start >= cluster->stack_size);
 
     wasm_runtime_module_free(module_inst, start - cluster->stack_size);
@@ -534,7 +538,7 @@ thread_manager_start_routine(void *arg)
 
 int32
 wasm_cluster_create_thread(WASMExecEnv *exec_env,
-                           wasm_module_inst_t module_inst,
+                           wasm_module_inst_t module_inst, bool alloc_aux_stack,
                            void *(*thread_routine)(void *), void *arg)
 {
     WASMCluster *cluster;
@@ -550,16 +554,18 @@ wasm_cluster_create_thread(WASMExecEnv *exec_env,
     if (!new_exec_env)
         return -1;
 
-    if (!allocate_aux_stack(exec_env, &aux_stack_start, &aux_stack_size)) {
-        LOG_ERROR("thread manager error: "
-                  "failed to allocate aux stack space for new thread");
-        goto fail1;
-    }
+    if (alloc_aux_stack) {
+        if (!allocate_aux_stack(exec_env, &aux_stack_start, &aux_stack_size)) {
+            LOG_ERROR("thread manager error: "
+                      "failed to allocate aux stack space for new thread");
+            goto fail1;
+        }
 
-    /* Set aux stack for current thread */
-    if (!wasm_exec_env_set_aux_stack(new_exec_env, aux_stack_start,
-                                     aux_stack_size)) {
-        goto fail2;
+        /* Set aux stack for current thread */
+        if (!wasm_exec_env_set_aux_stack(new_exec_env, aux_stack_start,
+                                         aux_stack_size)) {
+            goto fail2;
+        }
     }
 
     if (!wasm_cluster_add_exec_env(cluster, new_exec_env))
@@ -581,7 +587,8 @@ wasm_cluster_create_thread(WASMExecEnv *exec_env,
     wasm_cluster_del_exec_env(cluster, new_exec_env);
 fail2:
     /* free the allocated aux stack space */
-    free_aux_stack(exec_env, aux_stack_start);
+    if (alloc_aux_stack)
+        free_aux_stack(exec_env, aux_stack_start);
 fail1:
     wasm_exec_env_destroy(new_exec_env);
     return -1;

core/iwasm/libraries/thread-mgr/thread_manager.h

@@ -64,7 +64,7 @@ wasm_exec_env_get_cluster(WASMExecEnv *exec_env);
 
 int32
 wasm_cluster_create_thread(WASMExecEnv *exec_env,
-                           wasm_module_inst_t module_inst,
+                           wasm_module_inst_t module_inst, bool alloc_aux_stack,
                            void *(*thread_routine)(void *), void *arg);
 
 int32

samples/wasi-threads/wasm-apps/CMakeLists.txt

@@ -13,12 +13,13 @@ endif ()
 
 set (CMAKE_SYSROOT                  "${WASI_SYSROOT}")
 set (CMAKE_C_COMPILER               "${WASI_SDK_DIR}/bin/clang")
+set (CMAKE_ASM_COMPILER               "${WASI_SDK_DIR}/bin/clang")
 set (CMAKE_C_COMPILER_TARGET        "wasm32-wasi")
 
 function (compile_sample SOURCE_FILE)
   get_filename_component (FILE_NAME ${SOURCE_FILE} NAME_WLE)
   set (WASM_MODULE ${FILE_NAME}.wasm)
-  add_executable (${WASM_MODULE} ${SOURCE_FILE})
+  add_executable (${WASM_MODULE} ${SOURCE_FILE} ${ARGN})
 
   target_compile_options (${WASM_MODULE} PRIVATE
     -pthread -ftls-model=local-exec)
@@ -34,5 +35,5 @@ function (compile_sample SOURCE_FILE)
   )
 endfunction ()
 
-compile_sample(no_pthread.c)
-compile_sample(exception_propagation.c)
+compile_sample(no_pthread.c wasi_thread_start.S)
+compile_sample(exception_propagation.c wasi_thread_start.S)

samples/wasi-threads/wasm-apps/exception_propagation.c

@@ -9,15 +9,21 @@
 #include <stdlib.h>
 #include <stdio.h>
 #include <assert.h>
-#include <wasi/api.h>
 #include <semaphore.h>
 #include <stdbool.h>
 #include <unistd.h>
 
+#include "wasi_thread_start.h"
+
 #define TIMEOUT_SECONDS 10
 #define NUM_THREADS 3
 static sem_t sem;
 
+typedef struct {
+    start_args_t base;
+    bool throw_exception;
+} shared_t;
+
 void
 run_long_task()
 {
@@ -26,12 +32,12 @@ run_long_task()
         sleep(1);
 }
 
-__attribute__((export_name("wasi_thread_start"))) void
-wasi_thread_start(int thread_id, int *start_arg)
+void
+__wasi_thread_start_C(int thread_id, int *start_arg)
 {
-    bool has_to_throw_exception = (bool)start_arg;
+    shared_t *data = (shared_t *)start_arg;
 
-    if (has_to_throw_exception) {
+    if (data->throw_exception) {
         // Wait for all other threads (including main thread) to be ready
         printf("Waiting before throwing exception\n");
         for (int i = 0; i < NUM_THREADS; i++)
@@ -52,26 +58,36 @@ wasi_thread_start(int thread_id, int *start_arg)
 int
 main(int argc, char **argv)
 {
-    int thread_id = -1;
+    int thread_id = -1, i;
+    shared_t data[NUM_THREADS] = { 0 };
+
     if (sem_init(&sem, 0, 0) != 0) {
         printf("Failed to init semaphore\n");
         return EXIT_FAILURE;
     }
 
+    for (i = 0; i < NUM_THREADS; i++) {
+        // No graceful memory free to simplify the example
+        if (!start_args_init(&data[i].base)) {
+            printf("Failed to allocate thread's stack\n");
+            return EXIT_FAILURE;
+        }
+    }
+
     // Create a thread that throws an exception
-    thread_id = __wasi_thread_spawn((void *)true);
+    data[0].throw_exception = true;
+    thread_id = __wasi_thread_spawn(&data[0]);
     if (thread_id < 0) {
         printf("Failed to create thread: %d\n", thread_id);
         return EXIT_FAILURE;
     }
-
     // Create two additional threads to test exception propagation
-    thread_id = __wasi_thread_spawn((void *)false);
+    thread_id = __wasi_thread_spawn(&data[1]);
     if (thread_id < 0) {
         printf("Failed to create thread: %d\n", thread_id);
         return EXIT_FAILURE;
     }
-    thread_id = __wasi_thread_spawn((void *)false);
+    thread_id = __wasi_thread_spawn(&data[2]);
     if (thread_id < 0) {
         printf("Failed to create thread: %d\n", thread_id);
         return EXIT_FAILURE;

samples/wasi-threads/wasm-apps/no_pthread.c

@@ -9,18 +9,20 @@
 #include <stdlib.h>
 #include <stdio.h>
 #include <assert.h>
-#include <wasi/api.h>
+
+#include "wasi_thread_start.h"
 
 static const int64_t SECOND = 1000 * 1000 * 1000;
 
 typedef struct {
+    start_args_t base;
     int th_ready;
     int value;
     int thread_id;
 } shared_t;
 
-__attribute__((export_name("wasi_thread_start"))) void
-wasi_thread_start(int thread_id, int *start_arg)
+void
+__wasi_thread_start_C(int thread_id, int *start_arg)
 {
     shared_t *data = (shared_t *)start_arg;
 
@@ -38,24 +40,35 @@ wasi_thread_start(int thread_id, int *start_arg)
 int
 main(int argc, char **argv)
 {
-    shared_t data = { 0, 52, -1 };
+    shared_t data = { { NULL }, 0, 52, -1 };
     int thread_id;
+    int ret = EXIT_SUCCESS;
+
+    if (!start_args_init(&data.base)) {
+        printf("Stack allocation for thread failed\n");
+        return EXIT_FAILURE;
+    }
 
     thread_id = __wasi_thread_spawn(&data);
     if (thread_id < 0) {
         printf("Failed to create thread: %d\n", thread_id);
-        return EXIT_FAILURE;
+        ret = EXIT_FAILURE;
+        goto final;
     }
 
     if (__builtin_wasm_memory_atomic_wait32(&data.th_ready, 0, SECOND) == 2) {
         printf("Timeout\n");
-        return EXIT_FAILURE;
+        ret = EXIT_FAILURE;
+        goto final;
     }
 
     printf("Thread completed, new value: %d, thread id: %d\n", data.value,
            data.thread_id);
 
     assert(thread_id == data.thread_id);
 
-    return EXIT_SUCCESS;
+final:
+    start_args_deinit(&data.base);
+
+    return ret;
 }

samples/wasi-threads/wasm-apps/wasi_thread_start.S

@@ -0,0 +1,22 @@
+# A slightly modified copy of the wasi-libc implementation
+# https://github.com/WebAssembly/wasi-libc/pull/376/
+	.globaltype	__stack_pointer, i32
+	.functype	__wasi_thread_start_C (i32, i32) -> ()
+
+	.globl	wasi_thread_start
+
+wasi_thread_start:
+	.functype	wasi_thread_start (i32, i32) -> ()
+
+	# Set up the minimum C environment.
+	# Note: offsetof(start_arg, stack) == 0
+	local.get   1  # start_arg
+	i32.load    0  # stack
+	global.set  __stack_pointer
+
+	# Make the C function do the rest of work.
+	local.get   0  # tid
+	local.get   1  # start_arg
+	call __wasi_thread_start_C
+
+	end_function

samples/wasi-threads/wasm-apps/wasi_thread_start.h

@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2022 Amazon.com Inc. or its affiliates. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ */
+#ifndef WASI_THREAD_START_H
+#define WASI_THREAD_START_H
+
+#define STACK_SIZE 1024
+
+typedef struct {
+    void *stack;
+} start_args_t;
+
+static inline int
+start_args_init(start_args_t *start_args)
+{
+    start_args->stack = malloc(STACK_SIZE);
+    if (!start_args->stack) {
+        return 0;
+    }
+
+    start_args->stack += STACK_SIZE;
+    return 1;
+}
+
+static inline void
+start_args_deinit(start_args_t *start_args)
+{
+    free(start_args->stack - STACK_SIZE);
+}
+
+#endif