Stack switching: Infrastructure and runtime support (#10388)

* [pr1] base

* prtest:full

* make sure to use ControlFlow result in trace_suspended_continuation

* stack-switching: cleanup: remove stray c-api changes

These are remnants of unrelated wasmfx wasmtime experiments, possibly
suitable for later submission against upstream.

* stack-switching: reuse async_stack_size

* stack-switching: delete delete_me debugging

* stack-switching: address feedback in environ::types

* stack-switching: remove unused code from vmoffsets

* stack-switching: drop dependency on std

* stack-switching: add compilation checks to ci matrix

* stack-switching: remove debug_println cruft

* stack-switching: export environ consts consistently

* stack-switching: export vm pub items consistently

* table_pool: reduced capacity for large elements

VMContRef elements which takes up two words and we don't want to
double the size of all tables in order to support storing these.
This change changes the table to target storing the requested
max number of elements if they are "nominally" sized with
(potentially) reduced capacity for non-nominally sized types when
encountered.

Continuations are the only type of element which may result in
fewer table slots being available than requested.

* stack-switching: extend conditional compilation

A fair bit of the definitions for stack switching are still
enabled, but this patch takes things a bit further to avoid
compilation problems; notably, cont_new is now not compiled
in unless the feature is enabled.

* stack-switching: formatting fixes

* stack-switching: address new clippy checks

In addition, to get clippy to fully pass, plumbed in
additional config to make winch paths happy; there's no
impl for winch yet but plumbing through the feature is
required to make paths incorporating macros at various
layers satisfied (and it is expected we'll use the
features in the future).

* stack-switching: more conditional compilation fixes

* stack-switching: additional conditional compile on table builtins for continuations

* stack-switching: additional conditional compile fixes

* stack-switching: additional conditional compile in store

* stack-switching: remove overly strict assertion

* stack-switching: remove errantly dropped no_mangle in config c-api

* stack-switching: VMContObj::from_raw_parts

* stack-switching: remove duplicate async_stack_size feature check

* stack-switching: VMArray -> VMHostArray

* stack-switching: remove unnecessary clippy exception

* stack-switching: fix docs referenced VMRuntimeLimits

* stack-switching: fix doc typo

* stack-switching: follow recommendations for type casts

* stack-switching: use usize::next_multiple_of

* stack-switching: update outdated comment

* stack-switching: use feature gate instead of allow(dead_code)

* stack-switching: rework backtrace using chunks/zip

* stack-switching: move tests to footer module

This is a bit more consistent with the prevailing style
in tree and (subjectively) makes finding the tests
as a reader more straightforward.

Tests left unchanged sans some import cleanup.

* stack-swictchding: verify stack_chain offsets at runtime

* fixup! stack-switching: use feature gate instead of allow(dead_code)

* stack-switching: document continuation roots tracing using match arms

---------

Co-authored-by: Paul Osborne <paul.osborne@fastly.com>

Author: frank-emrich

.github/workflows/main.yml

@@ -371,6 +371,9 @@ jobs:
               -p wasmtime --no-default-features --features threads
               -p wasmtime --no-default-features --features runtime,threads
               -p wasmtime --no-default-features --features cranelift,threads
+              -p wasmtime --no-default-features --features stack-switching
+              -p wasmtime --no-default-features --features cranelift,stack-switching
+              -p wasmtime --no-default-features --features runtime,stack-switching
               -p wasmtime --features incremental-cache
               -p wasmtime --features profile-pulley
               -p wasmtime --all-features

Cargo.toml

@@ -431,6 +431,7 @@ default = [
   "gc",
   "gc-drc",
   "gc-null",
+  "stack-switching",
   "winch",
   "pulley",
 
@@ -489,6 +490,7 @@ gc = ["wasmtime-cli-flags/gc", "wasmtime/gc"]
 gc-drc = ["gc", "wasmtime/gc-drc", "wasmtime-cli-flags/gc-drc"]
 gc-null = ["gc", "wasmtime/gc-null", "wasmtime-cli-flags/gc-null"]
 pulley = ["wasmtime-cli-flags/pulley"]
+stack-switching = ["wasmtime/stack-switching", "wasmtime-cli-flags/stack-switching"]
 
 # CLI subcommands for the `wasmtime` executable. See `wasmtime $cmd --help`
 # for more information on each subcommand.

crates/c-api/include/wasmtime/config.h

@@ -250,6 +250,14 @@ WASMTIME_CONFIG_PROP(void, wasm_wide_arithmetic, bool)
 
 #ifdef WASMTIME_FEATURE_COMPILER
 
+/**
+ * \brief Configures whether the WebAssembly stack switching
+ * proposal is enabled.
+ *
+ * This setting is `false` by default.
+ */
+WASMTIME_CONFIG_PROP(void, wasm_stack_switching, bool)
+
 /**
  * \brief Configures how JIT code will be compiled.
  *

crates/c-api/src/config.rs

@@ -140,6 +140,11 @@ pub extern "C" fn wasmtime_config_wasm_memory64_set(c: &mut wasm_config_t, enabl
     c.config.wasm_memory64(enable);
 }
 
+#[unsafe(no_mangle)]
+pub extern "C" fn wasmtime_config_wasm_stack_switching_set(c: &mut wasm_config_t, enable: bool) {
+    c.config.wasm_stack_switching(enable);
+}
+
 #[unsafe(no_mangle)]
 #[cfg(any(feature = "cranelift", feature = "winch"))]
 pub extern "C" fn wasmtime_config_strategy_set(

crates/cli-flags/Cargo.toml

@@ -39,3 +39,4 @@ gc-null = ["gc", "wasmtime/gc-null"]
 threads = ["wasmtime/threads"]
 memory-protection-keys = ["wasmtime/memory-protection-keys"]
 pulley = ["wasmtime/pulley"]
+stack-switching = ["wasmtime/stack-switching"]

crates/cli-flags/src/lib.rs

@@ -377,6 +377,8 @@ wasmtime_option_group! {
         pub component_model_error_context: Option<bool>,
         /// Configure support for the function-references proposal.
         pub function_references: Option<bool>,
+        /// Configure support for the stack-switching proposal.
+        pub stack_switching: Option<bool>,
         /// Configure support for the GC proposal.
         pub gc: Option<bool>,
         /// Configure support for the custom-page-sizes proposal.
@@ -818,6 +820,23 @@ impl CommonOptions {
             config.native_unwind_info(enable);
         }
 
+        // async_stack_size enabled by either async or stack-switching, so
+        // cannot directly use match_feature!
+        #[cfg(any(feature = "async", feature = "stack-switching"))]
+        {
+            if let Some(size) = self.wasm.async_stack_size {
+                config.async_stack_size(size);
+            }
+        }
+        #[cfg(not(any(feature = "async", feature = "stack-switching")))]
+        {
+            if let Some(_size) = self.wasm.async_stack_size {
+                anyhow::bail!(concat!(
+                    "support for async/stack-switching disabled at compile time"
+                ));
+            }
+        }
+
         match_feature! {
             ["pooling-allocator" : self.opts.pooling_allocator.or(pooling_allocator_default)]
             enable => {
@@ -923,11 +942,6 @@ impl CommonOptions {
             );
         }
 
-        match_feature! {
-            ["async" : self.wasm.async_stack_size]
-            size => config.async_stack_size(size),
-            _ => err,
-        }
         match_feature! {
             ["async" : self.wasm.async_stack_zeroing]
             enable => config.async_stack_zeroing(enable),
@@ -940,7 +954,7 @@ impl CommonOptions {
             // If `-Wasync-stack-size` isn't passed then automatically adjust it
             // to the wasm stack size provided here too. That prevents the need
             // to pass both when one can generally be inferred from the other.
-            #[cfg(feature = "async")]
+            #[cfg(any(feature = "async", feature = "stack-switching"))]
             if self.wasm.async_stack_size.is_none() {
                 const DEFAULT_HOST_STACK: usize = 512 << 10;
                 config.async_stack_size(max + DEFAULT_HOST_STACK);
@@ -983,6 +997,9 @@ impl CommonOptions {
         if let Some(enable) = self.wasm.memory64.or(all) {
             config.wasm_memory64(enable);
         }
+        if let Some(enable) = self.wasm.stack_switching {
+            config.wasm_stack_switching(enable);
+        }
         if let Some(enable) = self.wasm.custom_page_sizes.or(all) {
             config.wasm_custom_page_sizes(enable);
         }
@@ -1023,6 +1040,7 @@ impl CommonOptions {
             ("gc", gc, wasm_gc)
             ("gc", reference_types, wasm_reference_types)
             ("gc", function_references, wasm_function_references)
+            ("stack-switching", stack_switching, wasm_stack_switching)
         }
         Ok(())
     }

crates/cranelift/Cargo.toml

@@ -46,4 +46,5 @@ wmemcheck = ["wasmtime-environ/wmemcheck"]
 gc = ["wasmtime-environ/gc"]
 gc-drc = ["gc", "wasmtime-environ/gc-drc"]
 gc-null = ["gc", "wasmtime-environ/gc-null"]
+stack-switching = []
 threads = ["wasmtime-environ/threads"]

crates/cranelift/src/func_environ.rs

@@ -3793,3 +3793,17 @@ fn index_type_to_ir_type(index_type: IndexType) -> ir::Type {
         IndexType::I64 => I64,
     }
 }
+
+/// TODO(10248) This is removed in the next stack switching PR. It stops the
+/// compiler from complaining about the stack switching libcalls being dead
+/// code.
+#[cfg(feature = "stack-switching")]
+#[allow(
+    dead_code,
+    reason = "Dummy function to supress more dead code warnings"
+)]
+pub fn use_stack_switching_libcalls() {
+    let _ = BuiltinFunctions::cont_new;
+    let _ = BuiltinFunctions::table_grow_cont_obj;
+    let _ = BuiltinFunctions::table_fill_cont_obj;
+}

crates/cranelift/src/func_environ/gc/enabled.rs

@@ -153,7 +153,12 @@ fn read_field_at_addr(
                         .call(get_interned_func_ref, &[vmctx, func_ref_id, expected_ty]);
                     builder.func.dfg.first_result(call_inst)
                 }
-                WasmHeapTopType::Cont => todo!(), // FIXME: #10248 stack switching support.
+                WasmHeapTopType::Cont => {
+                    // TODO(#10248) GC integration for stack switching
+                    return Err(wasmtime_environ::WasmError::Unsupported(
+                        "Stack switching feature not compatbile with GC, yet".to_string(),
+                    ));
+                }
             },
         },
     };
@@ -1032,6 +1037,8 @@ pub fn translate_ref_test(
         | WasmHeapType::NoExtern
         | WasmHeapType::Func
         | WasmHeapType::NoFunc
+        | WasmHeapType::Cont
+        | WasmHeapType::NoCont
         | WasmHeapType::I31 => unreachable!("handled top, bottom, and i31 types above"),
 
         // For these abstract but non-top and non-bottom types, we check the
@@ -1086,8 +1093,12 @@ pub fn translate_ref_test(
 
             func_env.is_subtype(builder, actual_shared_ty, expected_shared_ty)
         }
-
-        WasmHeapType::Cont | WasmHeapType::ConcreteCont(_) | WasmHeapType::NoCont => todo!(), // FIXME: #10248 stack switching support.
+        WasmHeapType::ConcreteCont(_) => {
+            // TODO(#10248) GC integration for stack switching
+            return Err(wasmtime_environ::WasmError::Unsupported(
+                "Stack switching feature not compatbile with GC, yet".to_string(),
+            ));
+        }
     };
     builder.ins().jump(continue_block, &[result.into()]);
 
@@ -1409,8 +1420,9 @@ impl FuncEnvironment<'_> {
             WasmHeapType::Func | WasmHeapType::ConcreteFunc(_) | WasmHeapType::NoFunc => {
                 unreachable!()
             }
-
-            WasmHeapType::Cont | WasmHeapType::ConcreteCont(_) | WasmHeapType::NoCont => todo!(), // FIXME: #10248 stack switching support.
+            WasmHeapType::Cont | WasmHeapType::ConcreteCont(_) | WasmHeapType::NoCont => {
+                unreachable!()
+            }
         };
 
         match (ty.nullable, might_be_i31) {

crates/cranelift/src/lib.rs

@@ -61,6 +61,10 @@ pub const TRAP_HEAP_MISALIGNED: TrapCode =
     TrapCode::unwrap_user(Trap::HeapMisaligned as u8 + TRAP_OFFSET);
 pub const TRAP_TABLE_OUT_OF_BOUNDS: TrapCode =
     TrapCode::unwrap_user(Trap::TableOutOfBounds as u8 + TRAP_OFFSET);
+pub const TRAP_UNHANDLED_TAG: TrapCode =
+    TrapCode::unwrap_user(Trap::UnhandledTag as u8 + TRAP_OFFSET);
+pub const TRAP_CONTINUATION_ALREADY_CONSUMED: TrapCode =
+    TrapCode::unwrap_user(Trap::ContinuationAlreadyConsumed as u8 + TRAP_OFFSET);
 pub const TRAP_CAST_FAILURE: TrapCode =
     TrapCode::unwrap_user(Trap::CastFailure as u8 + TRAP_OFFSET);
 
@@ -202,7 +206,11 @@ fn reference_type(wasm_ht: WasmHeapType, pointer_type: ir::Type) -> ir::Type {
     match wasm_ht.top() {
         WasmHeapTopType::Func => pointer_type,
         WasmHeapTopType::Any | WasmHeapTopType::Extern => ir::types::I32,
-        WasmHeapTopType::Cont => todo!(), // FIXME: #10248 stack switching support.
+        WasmHeapTopType::Cont =>
+        // TODO(10248) This is added in a follow-up PR
+        {
+            unimplemented!("codegen for stack switching types not implemented, yet")
+        }
     }
 }