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

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};
 
@@ -110,7 +111,7 @@ pub(crate) trait ABI {
 }
 
 /// ABI-specific representation of a function argument.
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub enum ABIArg {
     /// A register argument.
     Reg {
@@ -229,37 +230,38 @@ impl ABIResult {
 pub(crate) type ABIParams = SmallVec<[ABIArg; 6]>;
 
 /// An ABI-specific representation of a function signature.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct ABISig {
     /// Function parameters.
     pub params: ABIParams,
     /// Function result.
     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.