Fully support multiple returns in Wasmtime

cranelift/codegen/src/context.rs

@@ -267,13 +267,7 @@ impl Context {
         isa: &dyn TargetIsa,
     ) -> CodegenResult<Option<crate::isa::unwind::UnwindInfo>> {
         if let Some(backend) = isa.get_mach_backend() {
-            use crate::isa::CallConv;
-            use crate::machinst::UnwindInfoKind;
-            let unwind_info_kind = match self.func.signature.call_conv {
-                CallConv::Fast | CallConv::Cold | CallConv::SystemV => UnwindInfoKind::SystemV,
-                CallConv::WindowsFastcall => UnwindInfoKind::Windows,
-                _ => UnwindInfoKind::None,
-            };
+            let unwind_info_kind = isa.unwind_info_kind();
             let result = self.mach_compile_result.as_ref().unwrap();
             return backend.emit_unwind_info(result, unwind_info_kind);
         }

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

@@ -197,18 +197,26 @@ impl ABIMachineSpec for AArch64MachineDeps {
             next_stack = 16;
         }
 
-        // Note on return values: on the regular non-baldrdash ABI, we may return values in 8
-        // registers for V128 and I64 registers independently of the number of register values
-        // returned in the other class. That is, we can return values in up to 8 integer and 8
-        // vector registers at once.
-        // In Baldrdash, we can only use one register for return value for all the register
-        // classes. That is, we can't return values in both one integer and one vector register;
-        // only one return value may be in a register.
-
-        let (max_per_class_reg_vals, mut remaining_reg_vals) = match (args_or_rets, is_baldrdash) {
-            (ArgsOrRets::Args, _) => (8, 16),     // x0-x7 and v0-v7
-            (ArgsOrRets::Rets, false) => (8, 16), // x0-x7 and v0-v7
-            (ArgsOrRets::Rets, true) => (1, 1),   // x0 or v0, but not both
+        let (max_per_class_reg_vals, mut remaining_reg_vals) = match args_or_rets {
+            ArgsOrRets::Args => (8, 16), // x0-x7 and v0-v7
+
+            // Note on return values: on the regular ABI, we may return values
+            // in 8 registers for V128 and I64 registers independently of the
+            // number of register values returned in the other class. That is,
+            // we can return values in up to 8 integer and
+            // 8 vector registers at once.
+            //
+            // In Baldrdash and Wasmtime, we can only use one register for
+            // return value for all the register classes. That is, we can't
+            // return values in both one integer and one vector register; only
+            // one return value may be in a register.
+            ArgsOrRets::Rets => {
+                if is_baldrdash || call_conv.extends_wasmtime() {
+                    (1, 1) // x0 or v0, but not both
+                } else {
+                    (8, 16) // x0-x7 and v0-v7
+                }
+            }
         };
 
         for i in 0..params.len() {
@@ -282,15 +290,18 @@ impl ABIMachineSpec for AArch64MachineDeps {
                 // Compute the stack slot's size.
                 let size = (ty_bits(param.value_type) / 8) as u64;
 
-                let size = if call_conv != isa::CallConv::AppleAarch64 {
+                let size = if call_conv == isa::CallConv::AppleAarch64
+                    || (call_conv.extends_wasmtime() && args_or_rets == ArgsOrRets::Rets)
+                {
+                    // MacOS aarch64 and Wasmtime allow stack slots with
+                    // sizes less than 8 bytes. They still need to be
+                    // properly aligned on their natural data alignment,
+                    // though.
+                    size
+                } else {
                     // Every arg takes a minimum slot of 8 bytes. (16-byte stack
                     // alignment happens separately after all args.)
                     std::cmp::max(size, 8)
-                } else {
-                    // MacOS aarch64 allows stack slots with sizes less than 8
-                    // bytes. They still need to be properly aligned on their
-                    // natural data alignment, though.
-                    size
                 };
 
                 // Align the stack slot.

cranelift/codegen/src/isa/call_conv.rs

@@ -29,6 +29,16 @@ pub enum CallConv {
     Baldrdash2020,
     /// Specialized convention for the probestack function.
     Probestack,
+    /// Wasmtime equivalent of SystemV, not ABI-stable.
+    ///
+    /// Currently only differs in how multiple return values are handled,
+    /// returning the first return value in a register and everything else
+    /// through a return-pointer.
+    WasmtimeSystemV,
+    /// Wasmtime equivalent of WindowsFastcall, not ABI-stable.
+    ///
+    /// Differs from fastcall in the same way as `WasmtimeSystemV`.
+    WasmtimeFastcall,
 }
 
 impl CallConv {
@@ -63,7 +73,7 @@ impl CallConv {
     /// Is the calling convention extending the Windows Fastcall ABI?
     pub fn extends_windows_fastcall(self) -> bool {
         match self {
-            Self::WindowsFastcall | Self::BaldrdashWindows => true,
+            Self::WindowsFastcall | Self::BaldrdashWindows | Self::WasmtimeFastcall => true,
             _ => false,
         }
     }
@@ -75,6 +85,14 @@ impl CallConv {
             _ => false,
         }
     }
+
+    /// Is the calling convention extending the Wasmtime ABI?
+    pub fn extends_wasmtime(self) -> bool {
+        match self {
+            Self::WasmtimeSystemV | Self::WasmtimeFastcall => true,
+            _ => false,
+        }
+    }
 }
 
 impl fmt::Display for CallConv {
@@ -89,6 +107,8 @@ impl fmt::Display for CallConv {
             Self::BaldrdashWindows => "baldrdash_windows",
             Self::Baldrdash2020 => "baldrdash_2020",
             Self::Probestack => "probestack",
+            Self::WasmtimeSystemV => "wasmtime_system_v",
+            Self::WasmtimeFastcall => "wasmtime_fastcall",
         })
     }
 }
@@ -106,6 +126,8 @@ impl str::FromStr for CallConv {
             "baldrdash_windows" => Ok(Self::BaldrdashWindows),
             "baldrdash_2020" => Ok(Self::Baldrdash2020),
             "probestack" => Ok(Self::Probestack),
+            "wasmtime_system_v" => Ok(Self::WasmtimeSystemV),
+            "wasmtime_fastcall" => Ok(Self::WasmtimeFastcall),
             _ => Err(()),
         }
     }

cranelift/codegen/src/isa/mod.rs

@@ -57,7 +57,7 @@ use crate::flowgraph;
 use crate::ir;
 #[cfg(feature = "unwind")]
 use crate::isa::unwind::systemv::RegisterMappingError;
-use crate::machinst::MachBackend;
+use crate::machinst::{MachBackend, UnwindInfoKind};
 use crate::regalloc;
 use crate::result::CodegenResult;
 use crate::settings;
@@ -68,7 +68,7 @@ use core::any::Any;
 use core::fmt;
 use core::fmt::{Debug, Formatter};
 use core::hash::Hasher;
-use target_lexicon::{triple, Architecture, PointerWidth, Triple};
+use target_lexicon::{triple, Architecture, OperatingSystem, PointerWidth, Triple};
 use thiserror::Error;
 
 #[cfg(feature = "riscv")]
@@ -476,6 +476,18 @@ pub trait TargetIsa: fmt::Display + Send + Sync {
     /// IntCC condition for Unsigned Subtraction Overflow (Borrow/Carry).
     fn unsigned_sub_overflow_condition(&self) -> ir::condcodes::IntCC;
 
+    /// Returns the flavor of unwind information emitted for this target.
+    fn unwind_info_kind(&self) -> UnwindInfoKind {
+        match self.triple().operating_system {
+            #[cfg(feature = "unwind")]
+            OperatingSystem::Windows => UnwindInfoKind::Windows,
+            #[cfg(feature = "unwind")]
+            _ => UnwindInfoKind::SystemV,
+            #[cfg(not(feature = "unwind"))]
+            _ => UnwindInfoKind::None,
+        }
+    }
+
     /// Creates unwind information for the function.
     ///
     /// Returns `None` if there is no unwind information for the function.

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

@@ -237,10 +237,20 @@ impl ABIMachineSpec for X64ABIMachineSpec {
                         extension: param.extension,
                     });
                 } else {
-                    // Compute size. Every arg takes a minimum slot of 8 bytes. (16-byte
-                    // stack alignment happens separately after all args.)
+                    // Compute size. For the wasmtime ABI it differs from native
+                    // ABIs in how multiple values are returned, so we take a
+                    // leaf out of arm64's book by not rounding everything up to
+                    // 8 bytes. For all ABI arguments, and other ABI returns,
+                    // though, each slot takes a minimum of 8 bytes.
+                    //
+                    // Note that in all cases 16-byte stack alignment happens
+                    // separately after all args.
                     let size = (reg_ty.bits() / 8) as u64;
-                    let size = std::cmp::max(size, 8);
+                    let size = if args_or_rets == ArgsOrRets::Rets && call_conv.extends_wasmtime() {
+                        size
+                    } else {
+                        std::cmp::max(size, 8)
+                    };
                     // Align.
                     debug_assert!(size.is_power_of_two());
                     next_stack = align_to(next_stack, size);
@@ -824,15 +834,7 @@ impl From<StackAMode> for SyntheticAmode {
 }
 
 fn get_intreg_for_arg(call_conv: &CallConv, idx: usize, arg_idx: usize) -> Option<Reg> {
-    let is_fastcall = match call_conv {
-        CallConv::Fast
-        | CallConv::Cold
-        | CallConv::SystemV
-        | CallConv::BaldrdashSystemV
-        | CallConv::Baldrdash2020 => false,
-        CallConv::WindowsFastcall => true,
-        _ => panic!("int args only supported for SysV or Fastcall calling convention"),
-    };
+    let is_fastcall = call_conv.extends_windows_fastcall();
 
     // Fastcall counts by absolute argument number; SysV counts by argument of
     // this (integer) class.
@@ -853,15 +855,7 @@ fn get_intreg_for_arg(call_conv: &CallConv, idx: usize, arg_idx: usize) -> Optio
 }
 
 fn get_fltreg_for_arg(call_conv: &CallConv, idx: usize, arg_idx: usize) -> Option<Reg> {
-    let is_fastcall = match call_conv {
-        CallConv::Fast
-        | CallConv::Cold
-        | CallConv::SystemV
-        | CallConv::BaldrdashSystemV
-        | CallConv::Baldrdash2020 => false,
-        CallConv::WindowsFastcall => true,
-        _ => panic!("float args only supported for SysV or Fastcall calling convention"),
-    };
+    let is_fastcall = call_conv.extends_windows_fastcall();
 
     // Fastcall counts by absolute argument number; SysV counts by argument of
     // this (floating-point) class.
@@ -894,7 +888,10 @@ fn get_intreg_for_retval(
             1 => Some(regs::rdx()),
             _ => None,
         },
-        CallConv::BaldrdashSystemV | CallConv::Baldrdash2020 => {
+        CallConv::BaldrdashSystemV
+        | CallConv::Baldrdash2020
+        | CallConv::WasmtimeSystemV
+        | CallConv::WasmtimeFastcall => {
             if intreg_idx == 0 && retval_idx == 0 {
                 Some(regs::rax())
             } else {
@@ -922,7 +919,10 @@ fn get_fltreg_for_retval(
             1 => Some(regs::xmm1()),
             _ => None,
         },
-        CallConv::BaldrdashSystemV | CallConv::Baldrdash2020 => {
+        CallConv::BaldrdashSystemV
+        | CallConv::Baldrdash2020
+        | CallConv::WasmtimeFastcall
+        | CallConv::WasmtimeSystemV => {
             if fltreg_idx == 0 && retval_idx == 0 {
                 Some(regs::xmm0())
             } else {
@@ -992,12 +992,12 @@ fn get_callee_saves(call_conv: &CallConv, regs: &Set<Writable<RealReg>>) -> Vec<
         CallConv::BaldrdashWindows => {
             todo!("baldrdash windows");
         }
-        CallConv::Fast | CallConv::Cold | CallConv::SystemV => regs
+        CallConv::Fast | CallConv::Cold | CallConv::SystemV | CallConv::WasmtimeSystemV => regs
             .iter()
             .cloned()
             .filter(|r| is_callee_save_systemv(r.to_reg()))
             .collect(),
-        CallConv::WindowsFastcall => regs
+        CallConv::WindowsFastcall | CallConv::WasmtimeFastcall => regs
             .iter()
             .cloned()
             .filter(|r| is_callee_save_fastcall(r.to_reg()))

cranelift/codegen/src/isa/x64/inst/mod.rs

@@ -1122,11 +1122,16 @@ impl Inst {
     pub(crate) fn store(ty: Type, from_reg: Reg, to_addr: impl Into<SyntheticAmode>) -> Inst {
         let rc = from_reg.get_class();
         match rc {
-            RegClass::I64 => {
-                // Always store the full register, to ensure that the high bits are properly set
-                // when doing a full reload.
-                Inst::mov_r_m(OperandSize::Size64, from_reg, to_addr)
-            }
+            RegClass::I64 => Inst::mov_r_m(
+                match ty {
+                    types::B1 => OperandSize::Size8,
+                    types::I32 | types::R32 => OperandSize::Size32,
+                    types::I64 | types::R64 => OperandSize::Size64,
+                    _ => unimplemented!("integer store of type: {}", ty),
+                },
+                from_reg,
+                to_addr,
+            ),
             RegClass::V128 => {
                 let opcode = match ty {
                     types::F32 => SseOpcode::Movss,

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

@@ -503,10 +503,12 @@ fn callee_saved_regs_used(isa: &dyn TargetIsa, func: &ir::Function) -> RegisterS
 pub fn prologue_epilogue(func: &mut ir::Function, isa: &dyn TargetIsa) -> CodegenResult<()> {
     match func.signature.call_conv {
         // For now, just translate fast and cold as system_v.
-        CallConv::Fast | CallConv::Cold | CallConv::SystemV => {
+        CallConv::Fast | CallConv::Cold | CallConv::SystemV | CallConv::WasmtimeSystemV => {
             system_v_prologue_epilogue(func, isa)
         }
-        CallConv::WindowsFastcall => fastcall_prologue_epilogue(func, isa),
+        CallConv::WindowsFastcall | CallConv::WasmtimeFastcall => {
+            fastcall_prologue_epilogue(func, isa)
+        }
         CallConv::BaldrdashSystemV | CallConv::BaldrdashWindows => {
             baldrdash_prologue_epilogue(func, isa)
         }
@@ -1084,16 +1086,17 @@ pub fn create_unwind_info(
     isa: &dyn TargetIsa,
 ) -> CodegenResult<Option<crate::isa::unwind::UnwindInfo>> {
     use crate::isa::unwind::UnwindInfo;
+    use crate::machinst::UnwindInfoKind;
 
     // Assumption: RBP is being used as the frame pointer for both calling conventions
     // In the future, we should be omitting frame pointer as an optimization, so this will change
-    Ok(match func.signature.call_conv {
-        CallConv::Fast | CallConv::Cold | CallConv::SystemV => {
+    Ok(match isa.unwind_info_kind() {
+        UnwindInfoKind::SystemV => {
             super::unwind::systemv::create_unwind_info(func, isa)?.map(|u| UnwindInfo::SystemV(u))
         }
-        CallConv::WindowsFastcall => {
+        UnwindInfoKind::Windows => {
             super::unwind::winx64::create_unwind_info(func, isa)?.map(|u| UnwindInfo::WindowsX64(u))
         }
-        _ => None,
+        UnwindInfoKind::None => None,
     })
 }

cranelift/codegen/src/isa/x86/unwind/systemv.rs

@@ -3,7 +3,7 @@
 use crate::ir::Function;
 use crate::isa::{
     unwind::systemv::{RegisterMappingError, UnwindInfo},
-    CallConv, RegUnit, TargetIsa,
+    RegUnit, TargetIsa,
 };
 use crate::result::CodegenResult;
 use gimli::{write::CommonInformationEntry, Encoding, Format, Register, X86_64};
@@ -97,8 +97,8 @@ pub(crate) fn create_unwind_info(
     isa: &dyn TargetIsa,
 ) -> CodegenResult<Option<UnwindInfo>> {
     // Only System V-like calling conventions are supported
-    match func.signature.call_conv {
-        CallConv::Fast | CallConv::Cold | CallConv::SystemV => {}
+    match isa.unwind_info_kind() {
+        crate::machinst::UnwindInfoKind::SystemV => {}
         _ => return Ok(None),
     }
 

cranelift/codegen/src/machinst/abi.rs

@@ -189,9 +189,6 @@ pub trait ABICallee {
         from_slot: SpillSlot,
         ty: Option<Type>,
     ) -> Self::I;
-
-    /// Desired unwind info type.
-    fn unwind_info_kind(&self) -> UnwindInfoKind;
 }
 
 /// Trait implemented by an object that tracks ABI-related state and can

cranelift/codegen/src/machinst/abi_impl.rs

@@ -647,7 +647,8 @@ impl<M: ABIMachineSpec> ABICalleeImpl<M> {
                 || call_conv == isa::CallConv::Cold
                 || call_conv.extends_baldrdash()
                 || call_conv.extends_windows_fastcall()
-                || call_conv == isa::CallConv::AppleAarch64,
+                || call_conv == isa::CallConv::AppleAarch64
+                || call_conv == isa::CallConv::WasmtimeSystemV,
             "Unsupported calling convention: {:?}",
             call_conv
         );
@@ -1370,18 +1371,6 @@ impl<M: ABIMachineSpec> ABICallee for ABICalleeImpl<M> {
         .next()
         .unwrap()
     }
-
-    fn unwind_info_kind(&self) -> UnwindInfoKind {
-        match self.sig.call_conv {
-            #[cfg(feature = "unwind")]
-            isa::CallConv::Fast | isa::CallConv::Cold | isa::CallConv::SystemV => {
-                UnwindInfoKind::SystemV
-            }
-            #[cfg(feature = "unwind")]
-            isa::CallConv::WindowsFastcall => UnwindInfoKind::Windows,
-            _ => UnwindInfoKind::None,
-        }
-    }
 }
 
 fn abisig_to_uses_and_defs<M: ABIMachineSpec>(sig: &ABISig) -> (Vec<Reg>, Vec<Writable<Reg>>) {