cranelift: stack-switching support (#11003)

* cranelift: stack-switching support

This initial commit represents the "pr2" base commit with
minimal merge conflicts resolved.  Due to OOB conflicts, this
commit is not functional as-is, but using it as a base in
order to allow for easier reviewing of the delta from
this commit to what will be used for the PR against upstream.

Co-authored-by: Daniel Hillerström <[email protected]>
Co-authored-by: Paul Osborne <[email protected]>

* cranelift: stack-switching updates pass 1

This first set of changes updates the base pr in order to
compiled and pass basic checks (compile, clippy, fmt) with
the biggest part of the change being to eliminate injection
of tracing/assertions in JIT'ed code.

* cranelift: stack-switching: restore original visibility for a few func_environ members

* cranelift: stack-switching conditional compilation

At this point, the only bit we really branch on is what we
do in order to avoid problems tying into wasmtime_environ.
This is basd on the approach and macro used by the gc code for
converting presence/absence of the cranelift feature flag to
cranelift compile time.  This is a bit of a half-measure for now
as we still compile most stack-switching code in cranelift, but
this does enough to avoid causing problems with missing definitions
in wasmtime_environ.

* cranelift: avoid "as" casts in stack-switching

Replace either with infallible From or fallible, panicing
TryFrom alternatives where required.

* cranelift: cleanup stack-switching control_effect signatures

After removing emission of runtime trace logging and assertions,
there were several unused parameters.  Remove those from the
ControlEffect signatures completely.

* cranelift: rename stack-switching VMArray to VMHostArray

This matches a change to the mirrored runtime type in
the upstream changes.

* stack-switching: fix typo

Co-authored-by: Daniel Hillerström <[email protected]>

* stack-switching: used Index impl for get_stack_slot_data

* stack-switching: use smallvec over vec in several cases

* stack-switching: avoid resumetable naming confusion

* stack-switching: cleanup unused params from unchecked_get_continuation

The extra parameters here used to be used for emitting runtime
assertions, but with those gone we just had unused params
and lifetimes, clean those out.

* stack_switching: simplify store_data_entries assertion

* stack-switching: simplify translate_table_{grow,fill} control flow

* stack-switching: remove translate_resume_throw stub

There's already a stub elsewhere and this is not called, when
exceptions are added and it is time to revisit, this method
can be restored.

* stack-switching: compute control_context_size based on target triple

* stack-switching: VMHostArrayRef updates

Rename VMHostArray -> VMHostArrayRef
Change impl to compute address with offset upfront rather than
on each load.

* stack-switching: move cranelift code to live under func_environ

This matches the directory structure for gc and aids in visibility
for a few members required by stack-switching code in cranelift.

* stack-switching: formatting fix

* stack-switching: reduce visibility on a few additional items

* stack-switching: simplify contobj fatptr con/de-struction

* stack-switching: add disas tests to cover new instructions

* stack-switching: fix layout of VMContObj

In the course of the various runtime updates, the layout of the
runtime VMContObj got switched around.  This resulted in failures
when doing certain table operations on continuations.

This change fixes that layout problem and adds some tests with
offsets to avoid the problem.  Due to the way that we interact
with the VMContObj in cranelift, we don't use these offsets outside
of the tests.

* Fix formatting of merge conflict resolution

* cranelift: remove ir::function::get_stack_slot_data

This method isn't required as sized_stack_slots is already pub.

* stack-switching: reduce visibility of a couple func_environ methods

* stack-switching: define VMContObj as two words

This change migrates VMContObj and its usages in cranelift and runtime
to work with the VMContObj fat pointer as two words in order to better
target different architectures (still gated to x86_64 for now).

To support this, a size type was plumbed into the builtins function
signature types (as is done for component types) that maps to
usize.

* fixup! stack-switching: define VMContObj as two words

* stack-switching: add stub Val::ContRef

This type is not fully complete until continuation/gc integration
is revisited (#10248) but without these changes, test cases are
now failing on panics as we need some representation of
continuation references in the runtime Val enumeration.

Runtime errors with TODO notes are added for the stubbed
code paths to revisit later.

* fixup! stack-switching: add stub Val::ContRef

* fixup! stack-switching: add stub Val::ContRef

* fixup! stack-switching: define VMContObj as two words

prtest:full

* stack-switching: don't conflate host and target pointer sizes

Disas tests were failing on i686 targeting x86_64 as the size of the
host pointer was leaking into what we were using to do codegen
in a few paths.  This patch is a bit of a hack as it seems like
using a generic <T> for T: *mut u8 (as an example) is a bit
questionable.  To keep things small, I do a hacky typecheck to map
pointers to the target pointer size here.

* stack-switching: VMHostArray entry sizes based off env PtrSize

Revisiting the previous commit with an approach that should be
less brittle.

---------

Co-authored-by: Frank Emrich <[email protected]>
Co-authored-by: Daniel Hillerström <[email protected]>

Author: 3 people

cranelift/codegen/src/isa/x64/abi.rs

@@ -101,6 +101,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
         mut args: ArgsAccumulator,
     ) -> CodegenResult<(u32, Option<usize>)> {
         let is_fastcall = call_conv == CallConv::WindowsFastcall;
+        let is_tail = call_conv == CallConv::Tail;
 
         let mut next_gpr = 0;
         let mut next_vreg = 0;
@@ -181,6 +182,11 @@ impl ABIMachineSpec for X64ABIMachineSpec {
             //   This is consistent with LLVM's behavior, and is needed for
             //   some uses of Cranelift (e.g., the rustc backend).
             //
+            // - Otherwise, if the calling convention is Tail, we behave as in
+            //   the previous case, even if `enable_llvm_abi_extensions` is not
+            //   set in the flags: This is a custom calling convention defined
+            //   by Cranelift, LLVM doesn't know about it.
+            //
             // - Otherwise, both SysV and Fastcall specify behavior (use of
             //   vector register, a register pair, or passing by reference
             //   depending on the case), but for simplicity, we will just panic if
@@ -194,6 +200,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
             if param.value_type.bits() > 64
                 && !(param.value_type.is_vector() || param.value_type.is_float())
                 && !flags.enable_llvm_abi_extensions()
+                && !is_tail
             {
                 panic!(
                     "i128 args/return values not supported unless LLVM ABI extensions are enabled"

crates/c-api/src/val.rs

@@ -98,6 +98,7 @@ impl wasm_val_t {
             Val::ExternRef(_) => crate::abort("creating a wasm_val_t from an externref"),
             Val::ExnRef(_) => crate::abort("creating a wasm_val_t from  an exnref"),
             Val::V128(_) => crate::abort("creating a wasm_val_t from a v128"),
+            Val::ContRef(_) => crate::abort("creating a wasm_val_t from a contref"),
         }
     }
 
@@ -258,6 +259,7 @@ impl wasmtime_val_t {
                     v128: val.as_u128().to_le_bytes(),
                 },
             },
+            Val::ContRef(_) => crate::abort("contrefs not yet supported in C API (#10248)"),
         }
     }
 

crates/cranelift/src/func_environ.rs

@@ -1,4 +1,5 @@
 mod gc;
+pub(crate) mod stack_switching;
 
 use crate::compiler::Compiler;
 use crate::translate::{
@@ -171,6 +172,16 @@ pub struct FuncEnvironment<'module_environment> {
     /// always present even if this is a "leaf" function, as we have to call
     /// into the host to trap when signal handlers are disabled.
     pub(crate) stack_limit_at_function_entry: Option<ir::GlobalValue>,
+
+    /// Used by the stack switching feature. If set, we have a allocated a
+    /// slot on this function's stack to be used for the
+    /// current stack's `handler_list` field.
+    stack_switching_handler_list_buffer: Option<ir::StackSlot>,
+
+    /// Used by the stack switching feature. If set, we have a allocated a
+    /// slot on this function's stack to be used for the
+    /// current continuation's `values` field.
+    stack_switching_values_buffer: Option<ir::StackSlot>,
 }
 
 impl<'module_environment> FuncEnvironment<'module_environment> {
@@ -224,6 +235,9 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             translation,
 
             stack_limit_at_function_entry: None,
+
+            stack_switching_handler_list_buffer: None,
+            stack_switching_values_buffer: None,
         }
     }
 
@@ -1958,8 +1972,6 @@ impl<'a, 'func, 'module_env> Call<'a, 'func, 'module_env> {
                 return CheckIndirectCallTypeSignature::StaticTrap;
             }
 
-            WasmHeapType::Cont | WasmHeapType::ConcreteCont(_) | WasmHeapType::NoCont => todo!(), // FIXME: #10248 stack switching support.
-
             // Engine-indexed types don't show up until runtime and it's a Wasm
             // validation error to perform a call through a non-function table,
             // so these cases are dynamically not reachable.
@@ -1977,6 +1989,9 @@ impl<'a, 'func, 'module_env> Call<'a, 'func, 'module_env> {
             | WasmHeapType::Exn
             | WasmHeapType::ConcreteExn(_)
             | WasmHeapType::NoExn
+            | WasmHeapType::Cont
+            | WasmHeapType::ConcreteCont(_)
+            | WasmHeapType::NoCont
             | WasmHeapType::None => {
                 unreachable!()
             }
@@ -2267,7 +2282,9 @@ impl<'module_environment> TargetEnvironment for FuncEnvironment<'module_environm
         let needs_stack_map = match wasm_ty.top() {
             WasmHeapTopType::Extern | WasmHeapTopType::Any | WasmHeapTopType::Exn => true,
             WasmHeapTopType::Func => false,
-            WasmHeapTopType::Cont => todo!(), // FIXME: #10248 stack switching support.
+            // TODO(#10248) Once continuations can be stored on the GC heap, we
+            // will need stack maps for continuation objects.
+            WasmHeapTopType::Cont => false,
         };
         (ty, needs_stack_map)
     }
@@ -2320,22 +2337,32 @@ impl FuncEnvironment<'_> {
         let mut pos = builder.cursor();
         let table = self.table(table_index);
         let ty = table.ref_type.heap_type;
-        let grow = if ty.is_vmgcref_type() {
-            gc::builtins::table_grow_gc_ref(self, &mut pos.func)?
-        } else {
-            debug_assert_eq!(ty.top(), WasmHeapTopType::Func);
-            self.builtin_functions.table_grow_func_ref(&mut pos.func)
-        };
-
         let (table_vmctx, defined_table_index) =
             self.table_vmctx_and_defined_index(&mut pos, table_index);
-
         let index_type = table.idx_type;
         let delta = self.cast_index_to_i64(&mut pos, delta, index_type);
-        let call_inst = pos
-            .ins()
-            .call(grow, &[table_vmctx, defined_table_index, delta, init_value]);
-        let result = pos.func.dfg.first_result(call_inst);
+
+        let mut args: SmallVec<[_; 6]> = smallvec![table_vmctx, defined_table_index, delta];
+        let grow = match ty.top() {
+            WasmHeapTopType::Extern | WasmHeapTopType::Any | WasmHeapTopType::Exn => {
+                args.push(init_value);
+                gc::builtins::table_grow_gc_ref(self, pos.func)?
+            }
+            WasmHeapTopType::Func => {
+                args.push(init_value);
+                self.builtin_functions.table_grow_func_ref(pos.func)
+            }
+            WasmHeapTopType::Cont => {
+                let (revision, contref) =
+                    stack_switching::fatpointer::deconstruct(self, &mut pos, init_value);
+                args.extend_from_slice(&[contref, revision]);
+                stack_switching::builtins::table_grow_cont_obj(self, pos.func)?
+            }
+        };
+
+        let call_inst = pos.ins().call(grow, &args);
+        let result = builder.func.dfg.first_result(call_inst);
+
         Ok(self.convert_pointer_to_index_type(builder.cursor(), result, index_type, false))
     }
 
@@ -2367,7 +2394,15 @@ impl FuncEnvironment<'_> {
             }
 
             // Continuation types.
-            WasmHeapTopType::Cont => todo!(), // FIXME: #10248 stack switching support.
+            WasmHeapTopType::Cont => {
+                let (elem_addr, flags) = table_data.prepare_table_addr(self, builder, index);
+                Ok(builder.ins().load(
+                    stack_switching::fatpointer::fatpointer_type(self),
+                    flags,
+                    elem_addr,
+                    0,
+                ))
+            }
         }
     }
 
@@ -2415,7 +2450,11 @@ impl FuncEnvironment<'_> {
             }
 
             // Continuation types.
-            WasmHeapTopType::Cont => todo!(), // FIXME: #10248 stack switching support.
+            WasmHeapTopType::Cont => {
+                let (elem_addr, flags) = table_data.prepare_table_addr(self, builder, index);
+                builder.ins().store(flags, value, elem_addr, 0);
+                Ok(())
+            }
         }
     }
 
@@ -2429,21 +2468,31 @@ impl FuncEnvironment<'_> {
     ) -> WasmResult<()> {
         let mut pos = builder.cursor();
         let table = self.table(table_index);
-        let index_type = table.idx_type;
-        let dst = self.cast_index_to_i64(&mut pos, dst, index_type);
-        let len = self.cast_index_to_i64(&mut pos, len, index_type);
         let ty = table.ref_type.heap_type;
-        let libcall = if ty.is_vmgcref_type() {
-            gc::builtins::table_fill_gc_ref(self, &mut pos.func)?
-        } else {
-            debug_assert_eq!(ty.top(), WasmHeapTopType::Func);
-            self.builtin_functions.table_fill_func_ref(&mut pos.func)
-        };
-
+        let dst = self.cast_index_to_i64(&mut pos, dst, table.idx_type);
+        let len = self.cast_index_to_i64(&mut pos, len, table.idx_type);
         let (table_vmctx, table_index) = self.table_vmctx_and_defined_index(&mut pos, table_index);
 
-        pos.ins()
-            .call(libcall, &[table_vmctx, table_index, dst, val, len]);
+        let mut args: SmallVec<[_; 6]> = smallvec![table_vmctx, table_index, dst];
+        let libcall = match ty.top() {
+            WasmHeapTopType::Any | WasmHeapTopType::Extern | WasmHeapTopType::Exn => {
+                args.push(val);
+                gc::builtins::table_fill_gc_ref(self, &mut pos.func)?
+            }
+            WasmHeapTopType::Func => {
+                args.push(val);
+                self.builtin_functions.table_fill_func_ref(&mut pos.func)
+            }
+            WasmHeapTopType::Cont => {
+                let (revision, contref) =
+                    stack_switching::fatpointer::deconstruct(self, &mut pos, val);
+                args.extend_from_slice(&[contref, revision]);
+                stack_switching::builtins::table_fill_cont_obj(self, &mut pos.func)?
+            }
+        };
+
+        args.push(len);
+        builder.ins().call(libcall, &args);
 
         Ok(())
     }
@@ -2795,7 +2844,10 @@ impl FuncEnvironment<'_> {
             WasmHeapTopType::Any | WasmHeapTopType::Extern | WasmHeapTopType::Exn => {
                 pos.ins().iconst(types::I32, 0)
             }
-            WasmHeapTopType::Cont => todo!(), // FIXME: #10248 stack switching support.
+            WasmHeapTopType::Cont => {
+                let zero = pos.ins().iconst(self.pointer_type(), 0);
+                stack_switching::fatpointer::construct(self, &mut pos, zero, zero)
+            }
         })
     }
 
@@ -2811,9 +2863,18 @@ impl FuncEnvironment<'_> {
             return Ok(pos.ins().iconst(ir::types::I32, 0));
         }
 
-        let byte_is_null =
-            pos.ins()
-                .icmp_imm(cranelift_codegen::ir::condcodes::IntCC::Equal, value, 0);
+        let byte_is_null = match ty.heap_type.top() {
+            WasmHeapTopType::Cont => {
+                let (_revision, contref) =
+                    stack_switching::fatpointer::deconstruct(self, &mut pos, value);
+                pos.ins()
+                    .icmp_imm(cranelift_codegen::ir::condcodes::IntCC::Equal, contref, 0)
+            }
+            _ => pos
+                .ins()
+                .icmp_imm(cranelift_codegen::ir::condcodes::IntCC::Equal, value, 0),
+        };
+
         Ok(pos.ins().uextend(ir::types::I32, byte_is_null))
     }
 
@@ -3556,6 +3617,118 @@ impl FuncEnvironment<'_> {
         self.isa.triple().architecture == target_lexicon::Architecture::X86_64
     }
 
+    pub fn translate_cont_bind(
+        &mut self,
+        builder: &mut FunctionBuilder<'_>,
+        contobj: ir::Value,
+        args: &[ir::Value],
+    ) -> ir::Value {
+        stack_switching::instructions::translate_cont_bind(self, builder, contobj, args)
+    }
+
+    pub fn translate_cont_new(
+        &mut self,
+        builder: &mut FunctionBuilder<'_>,
+        func: ir::Value,
+        arg_types: &[WasmValType],
+        return_types: &[WasmValType],
+    ) -> WasmResult<ir::Value> {
+        stack_switching::instructions::translate_cont_new(
+            self,
+            builder,
+            func,
+            arg_types,
+            return_types,
+        )
+    }
+
+    pub fn translate_resume(
+        &mut self,
+        builder: &mut FunctionBuilder<'_>,
+        type_index: u32,
+        contobj: ir::Value,
+        resume_args: &[ir::Value],
+        resumetable: &[(u32, Option<ir::Block>)],
+    ) -> WasmResult<Vec<ir::Value>> {
+        stack_switching::instructions::translate_resume(
+            self,
+            builder,
+            type_index,
+            contobj,
+            resume_args,
+            resumetable,
+        )
+    }
+
+    pub fn translate_suspend(
+        &mut self,
+        builder: &mut FunctionBuilder<'_>,
+        tag_index: u32,
+        suspend_args: &[ir::Value],
+        tag_return_types: &[ir::Type],
+    ) -> Vec<ir::Value> {
+        stack_switching::instructions::translate_suspend(
+            self,
+            builder,
+            tag_index,
+            suspend_args,
+            tag_return_types,
+        )
+    }
+
+    /// Translates switch instructions.
+    pub fn translate_switch(
+        &mut self,
+        builder: &mut FunctionBuilder,
+        tag_index: u32,
+        contobj: ir::Value,
+        switch_args: &[ir::Value],
+        return_types: &[ir::Type],
+    ) -> WasmResult<Vec<ir::Value>> {
+        stack_switching::instructions::translate_switch(
+            self,
+            builder,
+            tag_index,
+            contobj,
+            switch_args,
+            return_types,
+        )
+    }
+
+    pub fn continuation_arguments(&self, index: TypeIndex) -> &[WasmValType] {
+        let idx = self.module.types[index].unwrap_module_type_index();
+        self.types[self.types[idx]
+            .unwrap_cont()
+            .clone()
+            .unwrap_module_type_index()]
+        .unwrap_func()
+        .params()
+    }
+
+    pub fn continuation_returns(&self, index: TypeIndex) -> &[WasmValType] {
+        let idx = self.module.types[index].unwrap_module_type_index();
+        self.types[self.types[idx]
+            .unwrap_cont()
+            .clone()
+            .unwrap_module_type_index()]
+        .unwrap_func()
+        .returns()
+    }
+
+    pub fn tag_params(&self, tag_index: TagIndex) -> &[WasmValType] {
+        let idx = self.module.tags[tag_index].signature;
+        self.types[idx.unwrap_module_type_index()]
+            .unwrap_func()
+            .params()
+    }
+
+    pub fn tag_returns(&self, tag_index: TagIndex) -> &[WasmValType] {
+        let idx = self.module.tags[tag_index].signature;
+        self.types[idx.unwrap_module_type_index()]
+            .unwrap_func()
+            .returns()
+    }
+
     pub fn use_x86_blendv_for_relaxed_laneselect(&self, ty: Type) -> bool {
         self.isa.has_x86_blendv_lowering(ty)
     }
@@ -4138,17 +4311,3 @@ 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 suppress 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;
-}