Add trampoline compilation support for lowered imports (#4206)

* Add trampoline compilation support for lowered imports

This commit adds support to the component model implementation for
compiling trampolines suitable for calling host imports. Currently this
is purely just the compilation side of things, modifying the
wasmtime-cranelift crate and additionally filling out a new
`VMComponentOffsets` type (similar to `VMOffsets`). The actual creation
of a `VMComponentContext` is still not performed and will be a
subsequent PR.

Internally though some tests are actually possible with this where we at
least assert that compilation of a component and creation of everything
in-memory doesn't panic or trip any assertions, so some tests are added
here for that as well.

* Fix some test errors

crates/wasmtime/Cargo.toml

@@ -108,4 +108,7 @@ memory-init-cow = ["wasmtime-runtime/memory-init-cow"]
 # Enables in-progress support for the component model. Note that this feature is
 # in-progress, buggy, and incomplete. This is primarily here for internal
 # testing purposes.
-component-model = ["wasmtime-environ/component-model"]
+component-model = [
+  "wasmtime-environ/component-model",
+  "wasmtime-cranelift?/component-model",
+]

crates/wasmtime/src/component/component.rs

@@ -1,10 +1,18 @@
+use crate::signatures::SignatureCollection;
 use crate::{Engine, Module};
 use anyhow::{bail, Context, Result};
 use std::fs;
+use std::ops::Range;
 use std::path::Path;
+use std::ptr::NonNull;
 use std::sync::Arc;
-use wasmtime_environ::component::{ComponentTypes, ModuleUpvarIndex, Translation, Translator};
+use wasmtime_environ::component::{
+    ComponentTypes, Initializer, LoweredIndex, ModuleUpvarIndex, TrampolineInfo, Translation,
+    Translator,
+};
 use wasmtime_environ::PrimaryMap;
+use wasmtime_jit::CodeMemory;
+use wasmtime_runtime::VMFunctionBody;
 
 /// A compiled WebAssembly Component.
 //
@@ -15,9 +23,36 @@ pub struct Component {
 }
 
 struct ComponentInner {
+    /// Type information calculated during translation about this component.
     component: wasmtime_environ::component::Component,
+
+    /// Core wasm modules that the component defined internally, indexed by the
+    /// compile-time-assigned `ModuleUpvarIndex`.
     upvars: PrimaryMap<ModuleUpvarIndex, Module>,
+
+    /// Registered core wasm signatures of this component, or otherwise the
+    /// mapping of the component-local `SignatureIndex` to the engine-local
+    /// `VMSharedSignatureIndex`.
+    signatures: SignatureCollection,
+
+    /// Type information about this component and all the various types it
+    /// defines internally. All type indices for `component` will point into
+    /// this field.
     types: Arc<ComponentTypes>,
+
+    /// The in-memory ELF image of the compiled trampolines for this component.
+    ///
+    /// This is currently only used for wasm-to-host trampolines when
+    /// `canon.lower` is encountered.
+    trampoline_obj: CodeMemory,
+
+    /// The index ranges within `trampoline_obj`'s mmap memory for the entire
+    /// text section.
+    text: Range<usize>,
+
+    /// Where trampolines are located within the `text` section of
+    /// `trampoline_obj`.
+    trampolines: PrimaryMap<LoweredIndex, TrampolineInfo>,
 }
 
 impl Component {
@@ -84,26 +119,73 @@ impl Component {
         let Translation {
             component, upvars, ..
         } = translation;
-        let upvars = upvars.into_iter().map(|(_, t)| t).collect::<Vec<_>>();
-        let upvars = engine
-            .run_maybe_parallel(upvars, |module| {
-                let (mmap, info) = Module::compile_functions(engine, module, types.module_types())?;
-                // FIXME: the `SignatureCollection` here is re-registering the
-                // entire list of wasm types within `types` on each invocation.
-                // That's ok semantically but is quite slow to do so. This
-                // should build up a mapping from `SignatureIndex` to
-                // `VMSharedSignatureIndex` once and then reuse that for each
-                // module somehow.
-                Module::from_parts(engine, mmap, info, types.clone())
-            })?
-            .into_iter()
-            .collect();
+        let (upvars, trampolines) = engine.join_maybe_parallel(
+            // In one (possibly) parallel task all the modules found within this
+            // component are compiled. Note that this will further parallelize
+            // function compilation internally too.
+            || -> Result<_> {
+                let upvars = upvars.into_iter().map(|(_, t)| t).collect::<Vec<_>>();
+                let modules = engine.run_maybe_parallel(upvars, |module| {
+                    let (mmap, info) =
+                        Module::compile_functions(engine, module, types.module_types())?;
+                    // FIXME: the `SignatureCollection` here is re-registering the
+                    // entire list of wasm types within `types` on each invocation.
+                    // That's ok semantically but is quite slow to do so. This
+                    // should build up a mapping from `SignatureIndex` to
+                    // `VMSharedSignatureIndex` once and then reuse that for each
+                    // module somehow.
+                    Module::from_parts(engine, mmap, info, types.clone())
+                })?;
+
+                Ok(modules.into_iter().collect::<PrimaryMap<_, _>>())
+            },
+            // In another (possibly) parallel task we compile lowering
+            // trampolines necessary found in the component.
+            || -> Result<_> {
+                let lowerings = component
+                    .initializers
+                    .iter()
+                    .filter_map(|init| match init {
+                        Initializer::LowerImport(i) => Some(i),
+                        _ => None,
+                    })
+                    .collect::<Vec<_>>();
+                let compiler = engine.compiler().component_compiler();
+                let trampolines = engine
+                    .run_maybe_parallel(lowerings, |lowering| {
+                        compiler.compile_lowered_trampoline(&component, lowering, &types)
+                    })?
+                    .into_iter()
+                    .collect();
+                let mut obj = engine.compiler().object()?;
+                let trampolines = compiler.emit_obj(trampolines, &mut obj)?;
+                Ok((trampolines, wasmtime_jit::mmap_vec_from_obj(obj)?))
+            },
+        );
+        let upvars = upvars?;
+        let (trampolines, trampoline_obj) = trampolines?;
+        let mut trampoline_obj = CodeMemory::new(trampoline_obj);
+        let code = trampoline_obj.publish()?;
+        let text = wasmtime_jit::subslice_range(code.text, code.mmap);
+
+        // FIXME: for the same reason as above where each module is
+        // re-registering everything this should only be registered once. This
+        // is benign for now but could do with refactorings later on.
+        let signatures = SignatureCollection::new_for_module(
+            engine.signatures(),
+            types.module_types(),
+            [].into_iter(),
+        );
 
         Ok(Component {
             inner: Arc::new(ComponentInner {
                 component,
                 upvars,
                 types,
+                trampolines,
+                trampoline_obj,
+                text,
+                signatures,
             }),
         })
     }
@@ -119,4 +201,15 @@ impl Component {
     pub(crate) fn types(&self) -> &Arc<ComponentTypes> {
         &self.inner.types
     }
+
+    pub(crate) fn signatures(&self) -> &SignatureCollection {
+        &self.inner.signatures
+    }
+
+    pub(crate) fn trampoline_ptr(&self, index: LoweredIndex) -> NonNull<VMFunctionBody> {
+        let info = &self.inner.trampolines[index];
+        let text = &self.inner.trampoline_obj.mmap()[self.inner.text.clone()];
+        let trampoline = &text[info.start as usize..][..info.length as usize];
+        NonNull::new(trampoline.as_ptr() as *mut VMFunctionBody).unwrap()
+    }
 }

crates/wasmtime/src/component/instance.rs

@@ -246,7 +246,16 @@ impl<'a> Instantiator<'a> {
                         unsafe { crate::Instance::new_started(store, module, imports.as_ref())? };
                     self.data.instances.push(i);
                 }
-                Initializer::LowerImport(_) => unimplemented!(),
+                Initializer::LowerImport(i) => {
+                    drop(self.component.trampoline_ptr(i.index));
+                    drop(
+                        self.component
+                            .signatures()
+                            .shared_signature(i.canonical_abi)
+                            .unwrap(),
+                    );
+                    unimplemented!()
+                }
 
                 Initializer::ExtractMemory(export) => {
                     let memory = match self.data.lookup_export(store.0, export) {

crates/wasmtime/src/engine.rs

@@ -224,6 +224,25 @@ impl Engine {
             .collect::<Result<Vec<B>, E>>()
     }
 
+    /// Executes `f1` and `f2` in parallel if parallel compilation is enabled at
+    /// both runtime and compile time, otherwise runs them synchronously.
+    #[allow(dead_code)] // only used for the component-model feature right now
+    pub(crate) fn join_maybe_parallel<T, U>(
+        &self,
+        f1: impl FnOnce() -> T + Send,
+        f2: impl FnOnce() -> U + Send,
+    ) -> (T, U)
+    where
+        T: Send,
+        U: Send,
+    {
+        if self.config().parallel_compilation {
+            #[cfg(feature = "parallel-compilation")]
+            return rayon::join(f1, f2);
+        }
+        (f1(), f2())
+    }
+
     /// Returns the target triple which this engine is compiling code for
     /// and/or running code for.
     pub(crate) fn target(&self) -> target_lexicon::Triple {