Implement lowered-then-lifted functions (bytecodealliance#4327)

* Implement lowered-then-lifted functions

This commit is a few features bundled into one, culminating in the
implementation of lowered-then-lifted functions for the component model.
It's probably not going to be used all that often but this is possible
within a valid component so Wasmtime needs to do something relatively
reasonable. The main things implemented in this commit are:

* Component instances are now assigned a `RuntimeComponentInstanceIndex`
  to differentiate each one. This will be used in the future to detect
  fusion (one instance lowering a function from another instance). For
  now it's used to allocate separate `VMComponentFlags` for each
  internal component instance.

* The `CoreExport<FuncIndex>` of lowered functions was changed to a
  `CoreDef` since technically a lowered function can use another lowered
  function as the callee. This ended up being not too difficult to plumb
  through as everything else was already in place.

* A need arose to compile host-to-wasm trampolines which weren't already
  present. Currently wasm in a component is always entered through a
  host-to-wasm trampoline but core wasm modules are the source of all
  the trampolines. In the case of a lowered-then-lifted function there
  may not actually be any core wasm modules, so component objects now
  contain necessary trampolines not otherwise provided by the core wasm
  objects. This feature required splitting a new function into the
  `Compiler` trait for creating a host-to-wasm trampoline. After doing
  this core wasm compilation was also updated to leverage this which
  further enabled compiling trampolines in parallel as opposed to the
  previous synchronous compilation.

* Review comments

Author: alexcrichton

crates/cranelift/src/compiler.rs

@@ -288,30 +288,37 @@ impl wasmtime_environ::Compiler for Compiler {
         }))
     }
 
+    fn compile_host_to_wasm_trampoline(
+        &self,
+        ty: &WasmFuncType,
+    ) -> Result<Box<dyn Any + Send>, CompileError> {
+        self.host_to_wasm_trampoline(ty)
+            .map(|x| Box::new(x) as Box<_>)
+    }
+
     fn emit_obj(
         &self,
         translation: &ModuleTranslation,
-        types: &ModuleTypes,
         funcs: PrimaryMap<DefinedFuncIndex, Box<dyn Any + Send>>,
+        compiled_trampolines: Vec<Box<dyn Any + Send>>,
         tunables: &Tunables,
         obj: &mut Object<'static>,
     ) -> Result<(PrimaryMap<DefinedFuncIndex, FunctionInfo>, Vec<Trampoline>)> {
         let funcs: CompiledFunctions = funcs
             .into_iter()
             .map(|(_i, f)| *f.downcast().unwrap())
             .collect();
+        let compiled_trampolines: Vec<CompiledFunction> = compiled_trampolines
+            .into_iter()
+            .map(|f| *f.downcast().unwrap())
+            .collect();
 
         let mut builder = ModuleTextBuilder::new(obj, &translation.module, &*self.isa);
         if self.linkopts.force_jump_veneers {
             builder.force_veneers();
         }
         let mut addrs = AddressMapSection::default();
         let mut traps = TrapEncodingBuilder::default();
-        let compiled_trampolines = translation
-            .exported_signatures
-            .iter()
-            .map(|i| self.host_to_wasm_trampoline(&types[*i]))
-            .collect::<Result<Vec<_>, _>>()?;
 
         let mut func_starts = Vec::with_capacity(funcs.len());
         for (i, func) in funcs.iter() {
@@ -325,6 +332,10 @@ impl wasmtime_environ::Compiler for Compiler {
         }
 
         // Build trampolines for every signature that can be used by this module.
+        assert_eq!(
+            translation.exported_signatures.len(),
+            compiled_trampolines.len()
+        );
         let mut trampolines = Vec::with_capacity(translation.exported_signatures.len());
         for (i, func) in translation
             .exported_signatures

crates/cranelift/src/compiler/component.rs

@@ -10,9 +10,9 @@ use object::write::Object;
 use std::any::Any;
 use wasmtime_environ::component::{
     CanonicalOptions, Component, ComponentCompiler, ComponentTypes, LowerImport, LoweredIndex,
-    TrampolineInfo, VMComponentOffsets,
+    LoweringInfo, VMComponentOffsets,
 };
-use wasmtime_environ::PrimaryMap;
+use wasmtime_environ::{PrimaryMap, SignatureIndex, Trampoline};
 
 impl ComponentCompiler for Compiler {
     fn compile_lowered_trampoline(
@@ -52,6 +52,7 @@ impl ComponentCompiler for Compiler {
         let mut host_sig = ir::Signature::new(crate::wasmtime_call_conv(isa));
 
         let CanonicalOptions {
+            instance,
             memory,
             realloc,
             post_return,
@@ -71,6 +72,14 @@ impl ComponentCompiler for Compiler {
             i32::try_from(offsets.lowering_data(lowering.index)).unwrap(),
         ));
 
+        // flags: *mut VMComponentFlags
+        host_sig.params.push(ir::AbiParam::new(pointer_type));
+        callee_args.push(
+            builder
+                .ins()
+                .iadd_imm(vmctx, i64::from(offsets.flags(instance))),
+        );
+
         // memory: *mut VMMemoryDefinition
         host_sig.params.push(ir::AbiParam::new(pointer_type));
         callee_args.push(match memory {
@@ -145,32 +154,42 @@ impl ComponentCompiler for Compiler {
 
     fn emit_obj(
         &self,
-        trampolines: PrimaryMap<LoweredIndex, Box<dyn Any + Send>>,
+        lowerings: PrimaryMap<LoweredIndex, Box<dyn Any + Send>>,
+        trampolines: Vec<(SignatureIndex, Box<dyn Any + Send>)>,
         obj: &mut Object<'static>,
-    ) -> Result<PrimaryMap<LoweredIndex, TrampolineInfo>> {
-        let trampolines: PrimaryMap<LoweredIndex, CompiledFunction> = trampolines
+    ) -> Result<(PrimaryMap<LoweredIndex, LoweringInfo>, Vec<Trampoline>)> {
+        let lowerings: PrimaryMap<LoweredIndex, CompiledFunction> = lowerings
             .into_iter()
             .map(|(_, f)| *f.downcast().unwrap())
             .collect();
+        let trampolines: Vec<(SignatureIndex, CompiledFunction)> = trampolines
+            .into_iter()
+            .map(|(i, f)| (i, *f.downcast().unwrap()))
+            .collect();
+
         let module = Default::default();
         let mut text = ModuleTextBuilder::new(obj, &module, &*self.isa);
         let mut ret = PrimaryMap::new();
-        for (idx, trampoline) in trampolines.iter() {
+        for (idx, lowering) in lowerings.iter() {
             let (_symbol, range) = text.append_func(
                 false,
-                format!("_wasm_component_host_trampoline{}", idx.as_u32()).into_bytes(),
-                &trampoline,
+                format!("_wasm_component_lowering_trampoline{}", idx.as_u32()).into_bytes(),
+                &lowering,
             );
 
-            let i = ret.push(TrampolineInfo {
+            let i = ret.push(LoweringInfo {
                 start: u32::try_from(range.start).unwrap(),
                 length: u32::try_from(range.end - range.start).unwrap(),
             });
             assert_eq!(i, idx);
         }
+        let ret_trampolines = trampolines
+            .iter()
+            .map(|(i, func)| text.trampoline(*i, func))
+            .collect();
 
         text.finish()?;
 
-        Ok(ret)
+        Ok((ret, ret_trampolines))
     }
 }

crates/environ/src/compilation.rs

@@ -148,12 +148,21 @@ pub trait Compiler: Send + Sync {
         types: &ModuleTypes,
     ) -> Result<Box<dyn Any + Send>, CompileError>;
 
+    /// Creates a function of type `VMTrampoline` which will then call the
+    /// function pointer argument which has the `ty` type provided.
+    fn compile_host_to_wasm_trampoline(
+        &self,
+        ty: &WasmFuncType,
+    ) -> Result<Box<dyn Any + Send>, CompileError>;
+
     /// Collects the results of compilation into an in-memory object.
     ///
     /// This function will receive the same `Box<dyn Ayn>` produced as part of
     /// `compile_function`, as well as the general compilation environment with
-    /// the translation/types. This method is expected to populate information
-    /// in the object file such as:
+    /// the translation. THe `trampolines` argument is generated by
+    /// `compile_host_to_wasm_trampoline` for each of
+    /// `module.exported_signatures`. This method is expected to populate
+    /// information in the object file such as:
     ///
     /// * Compiled code in a `.text` section
     /// * Unwind information in Wasmtime-specific sections
@@ -163,11 +172,14 @@ pub trait Compiler: Send + Sync {
     ///
     /// The final result of compilation will contain more sections inserted by
     /// the compiler-agnostic runtime.
+    ///
+    /// This function returns information about the compiled functions (where
+    /// they are in the text section) along with where trampolines are located.
     fn emit_obj(
         &self,
         module: &ModuleTranslation,
-        types: &ModuleTypes,
         funcs: PrimaryMap<DefinedFuncIndex, Box<dyn Any + Send>>,
+        trampolines: Vec<Box<dyn Any + Send>>,
         tunables: &Tunables,
         obj: &mut Object<'static>,
     ) -> Result<(PrimaryMap<DefinedFuncIndex, FunctionInfo>, Vec<Trampoline>)>;

crates/environ/src/component/compiler.rs

@@ -1,5 +1,5 @@
 use crate::component::{Component, ComponentTypes, LowerImport, LoweredIndex};
-use crate::PrimaryMap;
+use crate::{PrimaryMap, SignatureIndex, Trampoline};
 use anyhow::Result;
 use object::write::Object;
 use serde::{Deserialize, Serialize};
@@ -8,7 +8,7 @@ use std::any::Any;
 /// Description of where a trampoline is located in the text section of a
 /// compiled image.
 #[derive(Serialize, Deserialize)]
-pub struct TrampolineInfo {
+pub struct LoweringInfo {
     /// The byte offset from the start of the text section where this trampoline
     /// starts.
     pub start: u32,
@@ -42,8 +42,8 @@ pub trait ComponentCompiler: Send + Sync {
         types: &ComponentTypes,
     ) -> Result<Box<dyn Any + Send>>;
 
-    /// Emits the `trampolines` specified into the in-progress ELF object
-    /// specified by `obj`.
+    /// Emits the `lowerings` and `trampolines` specified into the in-progress
+    /// ELF object specified by `obj`.
     ///
     /// Returns a map of trampoline information for where to find them all in
     /// the text section.
@@ -52,7 +52,8 @@ pub trait ComponentCompiler: Send + Sync {
     /// trampolines as necessary.
     fn emit_obj(
         &self,
-        trampolines: PrimaryMap<LoweredIndex, Box<dyn Any + Send>>,
+        lowerings: PrimaryMap<LoweredIndex, Box<dyn Any + Send>>,
+        tramplines: Vec<(SignatureIndex, Box<dyn Any + Send>)>,
         obj: &mut Object<'static>,
-    ) -> Result<PrimaryMap<LoweredIndex, TrampolineInfo>>;
+    ) -> Result<(PrimaryMap<LoweredIndex, LoweringInfo>, Vec<Trampoline>)>;
 }

crates/environ/src/component/info.rs

@@ -114,6 +114,10 @@ pub struct Component {
     /// when instantiating this component.
     pub num_runtime_instances: u32,
 
+    /// Same as `num_runtime_instances`, but for `RuntimeComponentInstanceIndex`
+    /// instead.
+    pub num_runtime_component_instances: u32,
+
     /// The number of runtime memories (maximum `RuntimeMemoryIndex`) needed to
     /// instantiate this component.
     ///
@@ -355,7 +359,7 @@ pub enum Export {
         /// The component function type of the function being created.
         ty: TypeFuncIndex,
         /// Which core WebAssembly export is being lifted.
-        func: CoreExport<FuncIndex>,
+        func: CoreDef,
         /// Any options, if present, associated with this lifting.
         options: CanonicalOptions,
     },
@@ -369,6 +373,9 @@ pub enum Export {
 /// Canonical ABI options associated with a lifted or lowered function.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CanonicalOptions {
+    /// The component instance that this bundle was associated with.
+    pub instance: RuntimeComponentInstanceIndex,
+
     /// The encoding used for strings.
     pub string_encoding: StringEncoding,
 
@@ -382,17 +389,6 @@ pub struct CanonicalOptions {
     pub post_return: Option<RuntimePostReturnIndex>,
 }
 
-impl Default for CanonicalOptions {
-    fn default() -> CanonicalOptions {
-        CanonicalOptions {
-            string_encoding: StringEncoding::Utf8,
-            memory: None,
-            realloc: None,
-            post_return: None,
-        }
-    }
-}
-
 /// Possible encodings of strings within the component model.
 //
 // Note that the `repr(u8)` is load-bearing here since this is used in an

crates/environ/src/component/translate/inline.rs

@@ -104,7 +104,14 @@ pub(super) fn run(
     // initial frame. When the inliner finishes it will return the exports of
     // the root frame which are then used for recording the exports of the
     // component.
-    let mut frames = vec![InlinerFrame::new(result, ComponentClosure::default(), args)];
+    let index = RuntimeComponentInstanceIndex::from_u32(0);
+    inliner.result.num_runtime_component_instances += 1;
+    let mut frames = vec![InlinerFrame::new(
+        index,
+        result,
+        ComponentClosure::default(),
+        args,
+    )];
     let exports = inliner.run(&mut frames)?;
     assert!(frames.is_empty());
 
@@ -195,6 +202,8 @@ struct Inliner<'a> {
 /// inliner frames are stored on the heap to avoid recursion based on user
 /// input.
 struct InlinerFrame<'a> {
+    instance: RuntimeComponentInstanceIndex,
+
     /// The remaining initializers to process when instantiating this component.
     initializers: std::slice::Iter<'a, LocalInitializer<'a>>,
 
@@ -312,7 +321,7 @@ enum ComponentFuncDef<'a> {
     /// A core wasm function was lifted into a component function.
     Lifted {
         ty: TypeFuncIndex,
-        func: CoreExport<FuncIndex>,
+        func: CoreDef,
         options: CanonicalOptions,
     },
 }
@@ -509,19 +518,7 @@ impl<'a> Inliner<'a> {
                 let options = self.canonical_options(frame, options);
                 frame.component_funcs.push(ComponentFuncDef::Lifted {
                     ty: *ty,
-                    func: match frame.funcs[*func].clone() {
-                        CoreDef::Export(e) => e.map_index(|i| match i {
-                            EntityIndex::Function(i) => i,
-                            _ => unreachable!("not possible in valid components"),
-                        }),
-
-                        // TODO: lifting a lowered function only happens within
-                        // one component so this runs afoul of "someone needs to
-                        // really closely interpret the may_{enter,leave} flags"
-                        // in the component model spec. That has not currently
-                        // been done so this is left to panic.
-                        CoreDef::Lowered(_) => unimplemented!("lifting a lowered function"),
-                    },
+                    func: frame.funcs[*func].clone(),
                     options,
                 });
             }
@@ -618,7 +615,12 @@ impl<'a> Inliner<'a> {
             // stack.
             ComponentInstantiate(component, args) => {
                 let component: &ComponentDef<'a> = &frame.components[*component];
+                let index = RuntimeComponentInstanceIndex::from_u32(
+                    self.result.num_runtime_component_instances,
+                );
+                self.result.num_runtime_component_instances += 1;
                 let frame = InlinerFrame::new(
+                    index,
                     &self.nested_components[component.index],
                     component.closure.clone(),
                     args.iter()
@@ -872,6 +874,7 @@ impl<'a> Inliner<'a> {
                 })
         });
         CanonicalOptions {
+            instance: frame.instance,
             string_encoding: options.string_encoding,
             memory,
             realloc,
@@ -882,6 +885,7 @@ impl<'a> Inliner<'a> {
 
 impl<'a> InlinerFrame<'a> {
     fn new(
+        instance: RuntimeComponentInstanceIndex,
         translation: &'a Translation<'a>,
         closure: ComponentClosure<'a>,
         args: HashMap<&'a str, ComponentItemDef<'a>>,
@@ -891,6 +895,7 @@ impl<'a> InlinerFrame<'a> {
         // all the maps below. Given that doing such would be wordy and compile
         // time is otherwise not super crucial it's not done at this time.
         InlinerFrame {
+            instance,
             translation,
             closure,
             args,

crates/environ/src/component/types.rs

@@ -122,6 +122,9 @@ indices! {
     /// refer back to previously created instances for exports and such.
     pub struct RuntimeInstanceIndex(u32);
 
+    /// Same as `RuntimeInstanceIndex` but tracks component instances instead.
+    pub struct RuntimeComponentInstanceIndex(u32);
+
     /// Used to index imports into a `Component`
     ///
     /// This does not correspond to anything in the binary format for the