winch: Add known a subset of known libcalls and improve call emission

This change is a follow up to:
- bytecodealliance#7155
- bytecodealliance#7035

One of the objectives of this change is to make it easy to emit
function calls at the MacroAssembler layer, for cases in which it's
challenging to know ahead-of-time if a particular functionality can be
achieved natively (e.g. rounding and SSE4.2).  The original implementation
of function call emission, made this objective difficult to achieve and
it was also difficult to reason about.
I decided to simplify the overall approach to function calls as part of
this PR; in essence, the `call` module now exposes a single function
`FnCall::emit` which is reponsible of gathtering the dependencies and
orchestrating the emission of the call. This new approach deliberately
avoids holding any state regarding the function call for simplicity.

This change also standardizes the usage of `Callee` as the main
entrypoint for function call emission, as of this change 4 `Callee`
types exist (`Local`, `Builtin`, `Import`, `FuncRef`), each callee kind
is mappable to a `CalleeKind` which is the materialized version of
a callee which Cranelift understands.

This change also moves the creation of the `BuiltinFunctions` to the
`ISA` level given that they can be safely used accross multiple function
compilations.

Finally, this change also introduces support for some of the
"well-known" libcalls and hooks those libcalls at the
`MacroAssembler::float_round` callsite.

--

prtest:full

crates/winch/src/compiler.rs

@@ -1,18 +1,29 @@
 use anyhow::Result;
 use object::write::{Object, SymbolId};
 use std::any::Any;
+use std::mem;
 use std::sync::Mutex;
 use wasmparser::FuncValidatorAllocations;
 use wasmtime_cranelift_shared::{CompiledFunction, ModuleTextBuilder};
 use wasmtime_environ::{
     CompileError, DefinedFuncIndex, FilePos, FuncIndex, FunctionBodyData, FunctionLoc,
-    ModuleTranslation, ModuleTypes, PrimaryMap, TrapEncodingBuilder, WasmFunctionInfo,
+    ModuleTranslation, ModuleTypes, PrimaryMap, TrapEncodingBuilder, VMOffsets, WasmFunctionInfo,
 };
-use winch_codegen::{TargetIsa, TrampolineKind};
+use winch_codegen::{BuiltinFunctions, TargetIsa, TrampolineKind};
+
+/// Function compilation context.
+/// This struct holds information that can be shared globally across
+/// all function compilations.
+struct CompilationContext {
+    /// Validator allocations.
+    allocations: FuncValidatorAllocations,
+    /// Builtin functions available to JIT code.
+    builtins: BuiltinFunctions,
+}
 
 pub(crate) struct Compiler {
     isa: Box<dyn TargetIsa>,
-    allocations: Mutex<Vec<FuncValidatorAllocations>>,
+    contexts: Mutex<Vec<CompilationContext>>,
 }
 
 /// The compiled function environment.
@@ -30,20 +41,26 @@ impl Compiler {
     pub fn new(isa: Box<dyn TargetIsa>) -> Self {
         Self {
             isa,
-            allocations: Mutex::new(Vec::new()),
+            contexts: Mutex::new(Vec::new()),
         }
     }
 
-    fn take_allocations(&self) -> FuncValidatorAllocations {
-        self.allocations
-            .lock()
-            .unwrap()
-            .pop()
-            .unwrap_or_else(Default::default)
+    /// Get a compilation context or create a new one if none available.
+    fn get_context(&self, translation: &ModuleTranslation) -> CompilationContext {
+        self.contexts.lock().unwrap().pop().unwrap_or_else(|| {
+            let pointer_size = self.isa.pointer_bytes();
+            let vmoffsets = VMOffsets::new(pointer_size, &translation.module);
+            CompilationContext {
+                allocations: Default::default(),
+                builtins: BuiltinFunctions::new(&vmoffsets, self.isa.wasmtime_call_conv()),
+            }
+        })
     }
 
-    fn save_allocations(&self, allocs: FuncValidatorAllocations) {
-        self.allocations.lock().unwrap().push(allocs)
+    /// Save a compilation context.
+    fn save_context(&self, mut context: CompilationContext, allocs: FuncValidatorAllocations) {
+        context.allocations = allocs;
+        self.contexts.lock().unwrap().push(context);
     }
 }
 
@@ -65,12 +82,20 @@ impl wasmtime_environ::Compiler for Compiler {
                 .try_into()
                 .unwrap(),
         );
-        let mut validator = validator.into_validator(self.take_allocations());
+        let mut context = self.get_context(translation);
+        let mut validator = validator.into_validator(mem::take(&mut context.allocations));
         let buffer = self
             .isa
-            .compile_function(ty, types, &body, &translation, &mut validator)
+            .compile_function(
+                ty,
+                &body,
+                translation,
+                types,
+                &mut context.builtins,
+                &mut validator,
+            )
             .map_err(|e| CompileError::Codegen(format!("{e:?}")));
-        self.save_allocations(validator.into_allocations());
+        self.save_context(context, validator.into_allocations());
         let buffer = buffer?;
         let compiled_function =
             CompiledFunction::new(buffer, CompiledFuncEnv {}, self.isa.function_alignment());

winch/codegen/src/abi/mod.rs

@@ -46,6 +46,7 @@
 use crate::isa::{reg::Reg, CallingConvention};
 use crate::masm::OperandSize;
 use smallvec::SmallVec;
+use std::collections::HashSet;
 use std::ops::{Add, BitAnd, Not, Sub};
 use wasmtime_environ::{WasmFuncType, WasmHeapType, WasmType};
 
@@ -237,29 +238,30 @@ pub(crate) struct ABISig {
     pub result: ABIResult,
     /// Stack space needed for stack arguments.
     pub stack_bytes: u32,
+    /// All the registers used in the [`ABISig`].
+    /// Note that this collection is guaranteed to
+    /// be unique: in some cases some registers might
+    /// be used as params as a well as returns (e.g. xmm0 in x64).
+    pub regs: HashSet<Reg>,
 }
 
 impl ABISig {
     /// Create a new ABI signature.
     pub fn new(params: ABIParams, result: ABIResult, stack_bytes: u32) -> Self {
+        let regs = params
+            .iter()
+            .filter_map(|r| r.get_reg())
+            .collect::<SmallVec<[Reg; 6]>>();
+        let result_regs = result.regs();
+        let chained = regs.into_iter().chain(result_regs);
+
         Self {
             params,
             result,
             stack_bytes,
+            regs: HashSet::from_iter(chained),
         }
     }
-
-    /// Returns an iterator over all the registers used as params.
-    pub fn param_regs(&self) -> impl Iterator<Item = Reg> + '_ {
-        self.params.iter().filter_map(|r| r.get_reg())
-    }
-
-    /// Returns an iterator over all the registers used in the signature.
-    pub fn regs(&self) -> impl Iterator<Item = Reg> + '_ {
-        let params_iter = self.param_regs();
-        let result_iter = self.result.regs();
-        params_iter.chain(result_iter)
-    }
 }
 
 /// Returns the size in bytes of a given WebAssembly type.

winch/codegen/src/codegen/builtin.rs

@@ -5,16 +5,57 @@ use crate::{
     codegen::env::ptr_type_from_ptr_size,
     CallingConvention,
 };
-use wasmtime_environ::{BuiltinFunctionIndex, PtrSize, WasmType};
+use cranelift_codegen::ir::LibCall;
+use std::sync::Arc;
+use wasmtime_environ::{BuiltinFunctionIndex, PtrSize, VMOffsets, WasmType};
+
+#[derive(Copy, Clone)]
+pub(crate) enum BuiltinType {
+    /// Dynamic built-in function, derived from the VMContext.
+    Dynamic {
+        /// The offset of the built-in function.
+        offset: u32,
+        /// The built-in function base, relative to the VMContext.
+        base: u32,
+    },
+    /// A known libcall.
+    /// See [`cranelift_codegen::ir::LibCall`] for more details.
+    Known(LibCall),
+}
+
+impl BuiltinType {
+    /// Create a new dynamic built-in function type.
+    pub fn dynamic(offset: u32, base: u32) -> Self {
+        Self::Dynamic { offset, base }
+    }
+
+    /// Create a new known built-in function type.
+    pub fn known(libcall: LibCall) -> Self {
+        Self::Known(libcall)
+    }
+}
+
+#[derive(Clone)]
+pub struct BuiltinFunction {
+    inner: Arc<BuiltinFunctionInner>,
+}
+
+impl BuiltinFunction {
+    pub(crate) fn sig(&self) -> &ABISig {
+        &self.inner.sig
+    }
+
+    pub(crate) fn ty(&self) -> BuiltinType {
+        self.inner.ty
+    }
+}
 
 /// Metadata about a builtin function.
-pub(crate) struct BuiltinFunction {
+pub struct BuiltinFunctionInner {
     /// The ABI specific signature of the function.
-    pub sig: ABISig,
-    /// The offset of the builtin function
-    pub offset: u32,
-    /// The builtin function base, relative to the VMContext.
-    pub base: u32,
+    sig: ABISig,
+    /// The built-in function type.
+    ty: BuiltinType,
 }
 
 macro_rules! declare_function_sig {
@@ -35,6 +76,22 @@ macro_rules! declare_function_sig {
             ptr_type: WasmType,
             /// The builtin functions base relative to the VMContext.
             base: u32,
+            /// F32 Ceil.
+            ceil_f32: Option<BuiltinFunction>,
+            /// F64 Ceil.
+            ceil_f64: Option<BuiltinFunction>,
+            /// F32 Floor.
+            floor_f32: Option<BuiltinFunction>,
+            /// F64 Floor.
+            floor_f64: Option<BuiltinFunction>,
+            /// F32 Trunc.
+            trunc_f32: Option<BuiltinFunction>,
+            /// F64 Trunc.
+            trunc_f64: Option<BuiltinFunction>,
+            /// F32 Nearest.
+            nearest_f32: Option<BuiltinFunction>,
+            /// F64 Nearest.
+            nearest_f64: Option<BuiltinFunction>,
             $(
                 $name: Option<BuiltinFunction>,
             )*
@@ -43,13 +100,21 @@ macro_rules! declare_function_sig {
         // Until all the builtin functions are used.
         #[allow(dead_code)]
         impl BuiltinFunctions {
-            pub fn new(ptr: impl PtrSize, call_conv: CallingConvention, base: u32) -> Self {
-                let size = ptr.size();
+            pub fn new<P: PtrSize>(vmoffsets: &VMOffsets<P>, call_conv: CallingConvention) -> Self {
+                let size = vmoffsets.ptr.size();
                 Self {
                     ptr_size: size,
                     call_conv,
-                    base,
+                    base: vmoffsets.vmctx_builtin_functions(),
                     ptr_type: ptr_type_from_ptr_size(size),
+                    ceil_f32: None,
+                    ceil_f64: None,
+                    floor_f32: None,
+                    floor_f64: None,
+                    trunc_f32: None,
+                    trunc_f64: None,
+                    nearest_f32: None,
+                    nearest_f64: None,
                     $(
                         $name: None,
                     )*
@@ -68,6 +133,14 @@ macro_rules! declare_function_sig {
                 WasmType::I32
             }
 
+            fn f32(&self) -> WasmType {
+                WasmType::F32
+            }
+
+            fn f64(&self) -> WasmType {
+                WasmType::F64
+            }
+
             fn i64(&self) -> WasmType {
                 WasmType::I64
             }
@@ -76,21 +149,116 @@ macro_rules! declare_function_sig {
                 self.pointer()
             }
 
+            fn over_f64<A: ABI>(&self) -> ABISig {
+                A::sig_from(&[self.f64()], &[self.f64()], &self.call_conv)
+            }
+
+            fn over_f32<A: ABI>(&self) -> ABISig {
+                A::sig_from(&[self.f64()], &[self.f64()], &self.call_conv)
+            }
+
+            pub(crate) fn ceil_f32<A: ABI>(&mut self) -> BuiltinFunction {
+                if self.ceil_f32.is_none() {
+                    let sig = self.over_f32::<A>();
+                    let inner = Arc::new(BuiltinFunctionInner { sig, ty: BuiltinType::known(LibCall::CeilF32) });
+                    self.ceil_f32 = Some(BuiltinFunction {
+                        inner,
+                    });
+                }
+                self.ceil_f32.as_ref().unwrap().clone()
+            }
+
+            pub(crate) fn ceil_f64<A: ABI>(&mut self) -> BuiltinFunction {
+                if self.ceil_f64.is_none() {
+                    let sig = self.over_f64::<A>();
+                    let inner = Arc::new(BuiltinFunctionInner { sig, ty: BuiltinType::known(LibCall::CeilF64) });
+                    self.ceil_f64 = Some(BuiltinFunction {
+                        inner,
+                    });
+                }
+                self.ceil_f64.as_ref().unwrap().clone()
+            }
+
+            pub(crate) fn floor_f32<A: ABI>(&mut self) -> BuiltinFunction {
+                if self.floor_f32.is_none() {
+                    let sig = self.over_f32::<A>();
+                    let inner = Arc::new(BuiltinFunctionInner { sig, ty: BuiltinType::known(LibCall::FloorF32) });
+                    self.floor_f32 = Some(BuiltinFunction {
+                        inner,
+                    });
+                }
+                self.floor_f32.as_ref().unwrap().clone()
+            }
+
+            pub(crate) fn floor_f64<A: ABI>(&mut self) -> BuiltinFunction {
+                if self.floor_f64.is_none() {
+                    let sig = self.over_f64::<A>();
+                    let inner = Arc::new(BuiltinFunctionInner { sig, ty: BuiltinType::known(LibCall::FloorF64) });
+                    self.floor_f64 = Some(BuiltinFunction {
+                        inner,
+                    });
+                }
+                self.floor_f64.as_ref().unwrap().clone()
+            }
+
+            pub(crate) fn trunc_f32<A: ABI>(&mut self) -> BuiltinFunction {
+                if self.trunc_f32.is_none() {
+                    let sig = self.over_f32::<A>();
+                    let inner = Arc::new(BuiltinFunctionInner { sig, ty: BuiltinType::known(LibCall::TruncF32) });
+                    self.trunc_f32 = Some(BuiltinFunction {
+                        inner,
+                    });
+                }
+                self.trunc_f32.as_ref().unwrap().clone()
+            }
+
+            pub(crate) fn trunc_f64<A: ABI>(&mut self) -> BuiltinFunction {
+                if self.trunc_f64.is_none() {
+                    let sig = self.over_f64::<A>();
+                    let inner = Arc::new(BuiltinFunctionInner { sig, ty: BuiltinType::known(LibCall::TruncF64) });
+                    self.trunc_f64 = Some(BuiltinFunction {
+                        inner,
+                    });
+                }
+                self.trunc_f64.as_ref().unwrap().clone()
+            }
+
+            pub(crate) fn nearest_f32<A: ABI>(&mut self) -> BuiltinFunction {
+                if self.nearest_f32.is_none() {
+                    let sig = self.over_f32::<A>();
+                    let inner = Arc::new(BuiltinFunctionInner { sig, ty: BuiltinType::known(LibCall::NearestF32) });
+                    self.nearest_f32 = Some(BuiltinFunction {
+                        inner,
+                    });
+                }
+                self.nearest_f32.as_ref().unwrap().clone()
+            }
+
+            pub(crate) fn nearest_f64<A: ABI>(&mut self) -> BuiltinFunction {
+                if self.nearest_f64.is_none() {
+                    let sig = self.over_f64::<A>();
+                    let inner = Arc::new(BuiltinFunctionInner { sig, ty: BuiltinType::known(LibCall::NearestF64) });
+                    self.nearest_f64 = Some(BuiltinFunction {
+                        inner,
+                    });
+                }
+                self.nearest_f64.as_ref().unwrap().clone()
+            }
+
             $(
-                pub(crate) fn $name<A: ABI, P: PtrSize>(&mut self) -> &BuiltinFunction {
+                pub(crate) fn $name<A: ABI, P: PtrSize>(&mut self) -> BuiltinFunction {
                     if self.$name.is_none() {
                         let params = vec![ $(self.$param() ),* ];
                         let result = vec![ $(self.$result() )?];
                         let sig = A::sig_from(&params, &result, &self.call_conv);
                         let index = BuiltinFunctionIndex::$name();
+                        let inner = Arc::new(BuiltinFunctionInner { sig, ty: BuiltinType::dynamic(index.index() * (self.ptr_size as u32), self.base) });
                         self.$name = Some(BuiltinFunction {
-                            sig,
-                            offset: index.index() * (self.ptr_size as u32),
-                            base: self.base,
+                            inner,
                         });
                     }
 
-                    self.$name.as_ref().unwrap()
+                    self.$name.as_ref().unwrap().clone()
                 }
              )*
         }