Stack switching: Infrastructure and runtime support

.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

@@ -432,6 +432,7 @@ default = [
   "gc",
   "gc-drc",
   "gc-null",
+  "stack-switching",
   "winch",
   "pulley",
 
@@ -490,6 +491,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

@@ -141,6 +141,10 @@ pub extern "C" fn wasmtime_config_wasm_memory64_set(c: &mut wasm_config_t, enabl
 }
 
 #[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);
+}
+
 #[cfg(any(feature = "cranelift", feature = "winch"))]
 pub extern "C" fn wasmtime_config_strategy_set(
     c: &mut wasm_config_t,

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,6 +942,8 @@ impl CommonOptions {
             );
         }
 
+        #[cfg(any(feature = "async", feature = "stack-switching"))]
+
         match_feature! {
             ["async" : self.wasm.async_stack_size]
             size => config.async_stack_size(size),
@@ -940,7 +961,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 +1004,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 +1047,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

@@ -45,4 +45,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

@@ -3809,3 +3809,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")
+        }
     }
 }
 

crates/cranelift/src/translate/code_translator.rs

@@ -2898,6 +2898,56 @@ pub fn translate_operator(
             // representation, so we don't actually need to do anything.
         }
 
+        Operator::ContNew { cont_type_index: _ } => {
+            // TODO(10248) This is added in a follow-up PR
+            return Err(wasmtime_environ::WasmError::Unsupported(
+                "codegen for stack switching instructions not implemented, yet".to_string(),
+            ));
+        }
+        Operator::ContBind {
+            argument_index: _,
+            result_index: _,
+        } => {
+            // TODO(10248) This is added in a follow-up PR
+            return Err(wasmtime_environ::WasmError::Unsupported(
+                "codegen for stack switching instructions not implemented, yet".to_string(),
+            ));
+        }
+        Operator::Suspend { tag_index: _ } => {
+            // TODO(10248) This is added in a follow-up PR
+            return Err(wasmtime_environ::WasmError::Unsupported(
+                "codegen for stack switching instructions not implemented, yet".to_string(),
+            ));
+        }
+        Operator::Resume {
+            cont_type_index: _,
+            resume_table: _,
+        } => {
+            // TODO(10248) This is added in a follow-up PR
+            return Err(wasmtime_environ::WasmError::Unsupported(
+                "codegen for stack switching instructions not implemented, yet".to_string(),
+            ));
+        }
+        Operator::ResumeThrow {
+            cont_type_index: _,
+            tag_index: _,
+            resume_table: _,
+        } => {
+            // TODO(10248) This depends on exception handling
+            return Err(wasmtime_environ::WasmError::Unsupported(
+                "resume.throw instructions not supported, yet".to_string(),
+            ));
+        }
+        Operator::Switch {
+            cont_type_index: _,
+            tag_index: _,
+        } => {
+            // TODO(10248) This is added in a follow-up PR
+            return Err(wasmtime_environ::WasmError::Unsupported(
+                "codegen for stack switching instructions not implemented, yet".to_string(),
+            ));
+        }
+
         Operator::GlobalAtomicGet { .. }
         | Operator::GlobalAtomicSet { .. }
         | Operator::GlobalAtomicRmwAdd { .. }
@@ -2939,17 +2989,6 @@ pub fn translate_operator(
             ));
         }
 
-        Operator::ContNew { .. }
-        | Operator::ContBind { .. }
-        | Operator::Suspend { .. }
-        | Operator::Resume { .. }
-        | Operator::ResumeThrow { .. }
-        | Operator::Switch { .. } => {
-            return Err(wasm_unsupported!(
-                "stack-switching operators are not yet implemented"
-            ));
-        }
-
         Operator::I64MulWideS => {
             let (arg1, arg2) = state.pop2();
             let arg1 = builder.ins().sextend(I128, arg1);

crates/environ/Cargo.toml

@@ -69,6 +69,7 @@ compile = [
   "dep:wasm-encoder",
   "dep:wasmprinter",
 ]
+stack-switching = []
 threads = ['std']
 wmemcheck = ['std']
 std = [

crates/environ/src/builtin.rs

@@ -226,6 +226,25 @@ macro_rules! foreach_builtin_function {
             // Raises an unconditional trap where the trap information must have
             // been previously filled in.
             raise(vmctx: vmctx);
+
+            // Creates a new continuation from a funcref.
+            #[cfg(feature = "stack-switching")]
+            cont_new(vmctx: vmctx, r: pointer, param_count: u32, result_count: u32) -> pointer;
+
+            // Returns an index for Wasm's `table.grow` instruction
+            // for `contobj`s.  Note that the initial
+            // Option<VMContObj> (i.e., the value to fill the new
+            // slots with) is split into two arguments: The underlying
+            // continuation reference and the revision count.  To
+            // denote the continuation being `None`, `init_contref`
+            // may be 0.
+            #[cfg(feature = "stack-switching")]
+            table_grow_cont_obj(vmctx: vmctx, table: u32, delta: u64, init_contref: pointer, init_revision: u64) -> pointer;
+
+            // `value_contref` and `value_revision` together encode
+            // the Option<VMContObj>, as in previous libcall.
+            #[cfg(feature = "stack-switching")]
+            table_fill_cont_obj(vmctx: vmctx, table: u32, dst: u64, value_contref: pointer, value_revision: u64, len: u64) -> bool;
         }
     };
 }
@@ -367,6 +386,7 @@ impl BuiltinFunctionIndex {
             (@get memory32_grow pointer) => (TrapSentinel::NegativeTwo);
             (@get table_grow_func_ref pointer) => (TrapSentinel::NegativeTwo);
             (@get table_grow_gc_ref pointer) => (TrapSentinel::NegativeTwo);
+            (@get table_grow_cont_obj pointer) => (TrapSentinel::NegativeTwo);
 
             // Atomics-related functions return a negative value indicating trap
             // indicate a trap.
@@ -406,6 +426,8 @@ impl BuiltinFunctionIndex {
             (@get fma_f32x4 f32x4) => (return None);
             (@get fma_f64x2 f64x2) => (return None);
 
+            (@get cont_new pointer) => (TrapSentinel::Negative);
+
             // Bool-returning functions use `false` as an indicator of a trap.
             (@get $name:ident bool) => (TrapSentinel::Falsy);
 

crates/environ/src/gc.rs

@@ -40,10 +40,15 @@ pub const VM_GC_HEADER_TYPE_INDEX_OFFSET: u32 = 4;
 /// Get the byte size of the given Wasm type when it is stored inside the GC
 /// heap.
 pub fn byte_size_of_wasm_ty_in_gc_heap(ty: &WasmStorageType) -> u32 {
+    use crate::{WasmHeapType::*, WasmRefType};
     match ty {
         WasmStorageType::I8 => 1,
         WasmStorageType::I16 => 2,
         WasmStorageType::Val(ty) => match ty {
+            WasmValType::Ref(WasmRefType {
+                nullable: _,
+                heap_type: ConcreteCont(_) | Cont,
+            }) => unimplemented!("Stack switching feature not compatbile with GC, yet"),
             WasmValType::I32 | WasmValType::F32 | WasmValType::Ref(_) => 4,
             WasmValType::I64 | WasmValType::F64 => 8,
             WasmValType::V128 => 16,
@@ -178,7 +183,9 @@ pub trait GcTypeLayouts {
             WasmCompositeInnerType::Array(ty) => Some(self.array_layout(ty).into()),
             WasmCompositeInnerType::Struct(ty) => Some(self.struct_layout(ty).into()),
             WasmCompositeInnerType::Func(_) => None,
-            WasmCompositeInnerType::Cont(_) => None,
+            WasmCompositeInnerType::Cont(_) => {
+                unimplemented!("Stack switching feature not compatbile with GC, yet")
+            }
         }
     }
 

crates/environ/src/lib.rs

@@ -33,6 +33,7 @@ pub mod obj;
 mod ref_bits;
 mod scopevec;
 mod stack_map;
+mod stack_switching;
 mod trap_encoding;
 mod tunables;
 mod types;
@@ -51,6 +52,7 @@ pub use crate::module_types::*;
 pub use crate::ref_bits::*;
 pub use crate::scopevec::ScopeVec;
 pub use crate::stack_map::*;
+pub use crate::stack_switching::*;
 pub use crate::trap_encoding::*;
 pub use crate::tunables::*;
 pub use crate::types::*;