Make direct calls to statically-known imported functions (#11228)

* Make direct calls to statically-known imported functions

This will allow for their eventual inlining.

* address review feedback and fix up some disas tests

* Add native code disassembly test

* disas: sort `read_dir` entries for deterministic order

Because we always use namespace 0 for wasm functions when creating CLIF
functions, sorting the parsed CLIF functions by (namespace, index) does not
differentiate between the `index`th Wasm function in two different modules in
the same component, meaning we would pass through whatever ordering reading from
the filesystem gave us, which can be non-deterministic across
platforms. Instead, sort the directory entries based on file paths.

* fix clippy

Author: fitzgen

crates/cranelift/src/func_environ.rs

@@ -118,7 +118,6 @@ pub struct FuncEnvironment<'module_environment> {
     #[cfg(feature = "gc")]
     gc_heap_bound: Option<ir::GlobalValue>,
 
-    #[cfg(feature = "wmemcheck")]
     translation: &'module_environment ModuleTranslation<'module_environment>,
 
     /// Heaps implementing WebAssembly linear memories.
@@ -224,7 +223,6 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             // functions should consume at least some fuel.
             fuel_consumed: 1,
 
-            #[cfg(feature = "wmemcheck")]
             translation,
 
             stack_limit_at_function_entry: None,
@@ -1321,26 +1319,34 @@ impl<'a, 'func, 'module_env> Call<'a, 'func, 'module_env> {
                 .vmctx_vmfunction_import_wasm_call(callee_index),
         )
         .unwrap();
-        let func_addr = self
-            .builder
-            .ins()
-            .load(pointer_type, mem_flags, base, body_offset);
 
         // First append the callee vmctx address.
         let vmctx_offset =
             i32::try_from(self.env.offsets.vmctx_vmfunction_import_vmctx(callee_index)).unwrap();
-        let vmctx = self
+        let callee_vmctx = self
             .builder
             .ins()
             .load(pointer_type, mem_flags, base, vmctx_offset);
-        real_call_args.push(vmctx);
+        real_call_args.push(callee_vmctx);
         real_call_args.push(caller_vmctx);
 
         // Then append the regular call arguments.
         real_call_args.extend_from_slice(call_args);
 
-        // Finally, make the indirect call!
-        Ok(self.indirect_call_inst(sig_ref, func_addr, &real_call_args))
+        // If we statically know the imported function (e.g. this is a
+        // component-to-component call where we statically know both components)
+        // then we can actually still make a direct call (although we do have to
+        // pass the callee's vmctx that we just loaded, not our own). Otherwise,
+        // we really do an indirect call.
+        if self.env.translation.known_imported_functions[callee_index].is_some() {
+            Ok(self.direct_call_inst(callee, &real_call_args))
+        } else {
+            let func_addr = self
+                .builder
+                .ins()
+                .load(pointer_type, mem_flags, base, body_offset);
+            Ok(self.indirect_call_inst(sig_ref, func_addr, &real_call_args))
+        }
     }
 
     /// Do an indirect call through the given funcref table.
@@ -2651,7 +2657,8 @@ impl FuncEnvironment<'_> {
             // wasm module (e.g. imports or libcalls) are either encoded through
             // the `vmcontext` as relative jumps (hence no relocations) or
             // they're libcalls with absolute relocations.
-            colocated: self.module.defined_func_index(index).is_some(),
+            colocated: self.module.defined_func_index(index).is_some()
+                || self.translation.known_imported_functions[index].is_some(),
         }))
     }
 

crates/environ/src/compile/module_environ.rs

@@ -2,15 +2,16 @@ use crate::module::{
     FuncRefIndex, Initializer, MemoryInitialization, MemoryInitializer, Module, TableSegment,
     TableSegmentElements,
 };
-use crate::prelude::*;
 use crate::{
     ConstExpr, ConstOp, DataIndex, DefinedFuncIndex, ElemIndex, EngineOrModuleTypeIndex,
     EntityIndex, EntityType, FuncIndex, GlobalIndex, IndexType, InitMemory, MemoryIndex,
     ModuleInternedTypeIndex, ModuleTypesBuilder, PrimaryMap, SizeOverflow, StaticMemoryInitializer,
     TableIndex, TableInitialValue, Tag, TagIndex, Tunables, TypeConvert, TypeIndex, Unsigned,
     WasmError, WasmHeapTopType, WasmHeapType, WasmResult, WasmValType, WasmparserTypeConverter,
 };
+use crate::{StaticModuleIndex, prelude::*};
 use anyhow::{Result, bail};
+use cranelift_entity::SecondaryMap;
 use cranelift_entity::packed_option::ReservedValue;
 use std::borrow::Cow;
 use std::collections::HashMap;
@@ -54,6 +55,13 @@ pub struct ModuleTranslation<'data> {
     /// References to the function bodies.
     pub function_body_inputs: PrimaryMap<DefinedFuncIndex, FunctionBodyData<'data>>,
 
+    /// For each imported function, the single statically-known defined function
+    /// that satisfies that import, if any. This is used to turn what would
+    /// otherwise be indirect calls through the imports table into direct calls,
+    /// when possible.
+    pub known_imported_functions:
+        SecondaryMap<FuncIndex, Option<(StaticModuleIndex, DefinedFuncIndex)>>,
+
     /// A list of type signatures which are considered exported from this
     /// module, or those that can possibly be called. This list is sorted, and
     /// trampolines for each of these signatures are required.

crates/environ/src/component/dfg.rs

@@ -159,6 +159,40 @@ pub enum SideEffect {
     Resource(DefinedResourceIndex),
 }
 
+/// A sound approximation of a particular module's set of instantiations.
+///
+/// This type forms a simple lattice that we can use in static analyses that in
+/// turn let us specialize a module's compilation to exactly the imports it is
+/// given.
+#[derive(Clone, Copy, Default)]
+pub enum AbstractInstantiations<'a> {
+    /// The associated module is instantiated many times.
+    Many,
+
+    /// The module is instantiated exactly once, with the given definitions as
+    /// arguments to that instantiation.
+    One(&'a [info::CoreDef]),
+
+    /// The module is never instantiated.
+    #[default]
+    None,
+}
+
+impl AbstractInstantiations<'_> {
+    /// Join two facts about a particular module's instantiation together.
+    ///
+    /// This is the least-upper-bound operation on the lattice.
+    pub fn join(&mut self, other: Self) {
+        *self = match (*self, other) {
+            (Self::Many, _) | (_, Self::Many) => Self::Many,
+            (Self::One(a), Self::One(b)) if a == b => Self::One(a),
+            (Self::One(_), Self::One(_)) => Self::Many,
+            (Self::One(a), Self::None) | (Self::None, Self::One(a)) => Self::One(a),
+            (Self::None, Self::None) => Self::None,
+        }
+    }
+}
+
 macro_rules! id {
     ($(pub struct $name:ident(u32);)*) => ($(
         #[derive(Debug, Copy, Clone, Hash, Eq, PartialEq)]

crates/environ/src/component/translate.rs

@@ -1,4 +1,5 @@
 use crate::ScopeVec;
+use crate::component::dfg::AbstractInstantiations;
 use crate::component::*;
 use crate::prelude::*;
 use crate::{
@@ -8,6 +9,8 @@ use crate::{
 };
 use anyhow::anyhow;
 use anyhow::{Result, bail};
+use cranelift_entity::SecondaryMap;
+use cranelift_entity::packed_option::PackedOption;
 use indexmap::IndexMap;
 use std::collections::HashMap;
 use std::mem;
@@ -492,12 +495,133 @@ impl<'a, 'data> Translator<'a, 'data> {
             &self.static_modules,
             &self.static_components,
         )?;
+
         self.partition_adapter_modules(&mut component);
+
         let translation =
             component.finish(self.types.types_mut_for_inlining(), self.result.types_ref())?;
+
+        self.analyze_function_imports(&translation);
+
         Ok((translation, self.static_modules))
     }
 
+    fn analyze_function_imports(&mut self, translation: &ComponentTranslation) {
+        // First, abstract interpret the initializers to create a map from each
+        // static module to its abstract set of instantiations.
+        let mut instantiations = SecondaryMap::<StaticModuleIndex, AbstractInstantiations>::new();
+        let mut instance_to_module =
+            PrimaryMap::<RuntimeInstanceIndex, PackedOption<StaticModuleIndex>>::new();
+        for init in &translation.component.initializers {
+            match init {
+                GlobalInitializer::InstantiateModule(instantiation) => match instantiation {
+                    InstantiateModule::Static(module, args) => {
+                        instantiations[*module].join(AbstractInstantiations::One(&*args));
+                        instance_to_module.push(Some(*module).into());
+                    }
+                    _ => {
+                        instance_to_module.push(None.into());
+                    }
+                },
+                _ => continue,
+            }
+        }
+
+        // Second, make sure to mark exported modules as instantiated many
+        // times, since they could be linked with who-knows-what at runtime.
+        for item in translation.component.export_items.values() {
+            if let Export::ModuleStatic { index, .. } = item {
+                instantiations[*index].join(AbstractInstantiations::Many)
+            }
+        }
+
+        // Finally, iterate over our instantiations and record statically-known
+        // function imports so that they can get translated into direct calls
+        // (and eventually get inlined) rather than indirect calls through the
+        // imports table.
+        for (module, instantiations) in instantiations.iter() {
+            let args = match instantiations {
+                dfg::AbstractInstantiations::Many | dfg::AbstractInstantiations::None => continue,
+                dfg::AbstractInstantiations::One(args) => args,
+            };
+
+            let mut imported_func_counter = 0_u32;
+            for (i, arg) in args.iter().enumerate() {
+                // Only consider function imports.
+                let (_, _, crate::types::EntityType::Function(_)) =
+                    self.static_modules[module].module.import(i).unwrap()
+                else {
+                    continue;
+                };
+
+                let imported_func = FuncIndex::from_u32(imported_func_counter);
+                imported_func_counter += 1;
+                debug_assert!(
+                    self.static_modules[module]
+                        .module
+                        .defined_func_index(imported_func)
+                        .is_none()
+                );
+
+                match arg {
+                    CoreDef::InstanceFlags(_) => unreachable!("instance flags are not a function"),
+
+                    // We could in theory inline these trampolines, so it could
+                    // potentially make sense to record that we know this
+                    // imported function is this particular trampoline. However,
+                    // everything else is based around (module,
+                    // defined-function) pairs and these trampolines don't fit
+                    // that paradigm. Also, inlining trampolines gets really
+                    // tricky when we consider the stack pointer, frame pointer,
+                    // and return address note-taking that they do for the
+                    // purposes of stack walking. We could, with enough effort,
+                    // turn them into direct calls even though we probably
+                    // wouldn't ever inline them, but it just doesn't seem worth
+                    // the effort.
+                    CoreDef::Trampoline(_) => continue,
+
+                    // This imported function is an export from another
+                    // instance, a perfect candidate for becoming an inlinable
+                    // direct call!
+                    CoreDef::Export(export) => {
+                        let Some(arg_module) = &instance_to_module[export.instance].expand() else {
+                            // Instance of a dynamic module that is not part of
+                            // this component, not a statically-known module
+                            // inside this component. We have to do an indirect
+                            // call.
+                            continue;
+                        };
+
+                        let ExportItem::Index(EntityIndex::Function(arg_func)) = &export.item
+                        else {
+                            unreachable!("function imports must be functions")
+                        };
+
+                        let Some(arg_module_def_func) = self.static_modules[*arg_module]
+                            .module
+                            .defined_func_index(*arg_func)
+                        else {
+                            // TODO: we should ideally follow re-export chains
+                            // to bottom out the instantiation argument in
+                            // either a definition or an import at the root
+                            // component boundary. In practice, this pattern is
+                            // rare, so following these chains is left for the
+                            // Future.
+                            continue;
+                        };
+
+                        assert!(
+                            self.static_modules[module].known_imported_functions[imported_func]
+                                .is_none()
+                        );
+                        self.static_modules[module].known_imported_functions[imported_func] =
+                            Some((*arg_module, arg_module_def_func));
+                    }
+                }
+            }
+        }
+    }
+
     fn translate_payload(
         &mut self,
         payload: Payload<'data>,

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

@@ -1045,8 +1045,7 @@ impl<'a> Inliner<'a> {
             // and an initializer is recorded to indicate that it's being
             // instantiated.
             ModuleInstantiate(module, args) => {
-                let instance_module;
-                let init = match &frame.modules[*module] {
+                let (instance_module, init) = match &frame.modules[*module] {
                     ModuleDef::Static(idx, _ty) => {
                         let mut defs = Vec::new();
                         for (module, name, _ty) in self.nested_modules[*idx].module.imports() {
@@ -1055,8 +1054,10 @@ impl<'a> Inliner<'a> {
                                 self.core_def_of_module_instance_export(frame, instance, name),
                             );
                         }
-                        instance_module = InstanceModule::Static(*idx);
-                        dfg::Instance::Static(*idx, defs.into())
+                        (
+                            InstanceModule::Static(*idx),
+                            dfg::Instance::Static(*idx, defs.into()),
+                        )
                     }
                     ModuleDef::Import(path, ty) => {
                         let mut defs = IndexMap::new();
@@ -1069,20 +1070,24 @@ impl<'a> Inliner<'a> {
                                 .insert(name.to_string(), def);
                         }
                         let index = self.runtime_import(path);
-                        instance_module = InstanceModule::Import(*ty);
-                        dfg::Instance::Import(index, defs)
+                        (
+                            InstanceModule::Import(*ty),
+                            dfg::Instance::Import(index, defs),
+                        )
                     }
                 };
 
-                let idx = self.result.instances.push(init);
+                let instance = self.result.instances.push(init);
+                let instance2 = self.runtime_instances.push(instance_module);
+                assert_eq!(instance, instance2);
+
                 self.result
                     .side_effects
-                    .push(dfg::SideEffect::Instance(idx));
-                let idx2 = self.runtime_instances.push(instance_module);
-                assert_eq!(idx, idx2);
+                    .push(dfg::SideEffect::Instance(instance));
+
                 frame
                     .module_instances
-                    .push(ModuleInstanceDef::Instantiated(idx, *module));
+                    .push(ModuleInstanceDef::Instantiated(instance, *module));
             }
 
             ModuleSynthetic(map) => {
@@ -1772,6 +1777,7 @@ impl<'a> ComponentItemDef<'a> {
     }
 }
 
+#[derive(Clone, Copy)]
 enum InstanceModule {
     Static(StaticModuleIndex),
     Import(TypeModuleIndex),

crates/environ/src/module.rs

@@ -542,6 +542,13 @@ impl Module {
         })
     }
 
+    /// Get this module's `i`th import.
+    pub fn import(&self, i: usize) -> Option<(&str, &str, EntityType)> {
+        match self.initializers.get(i)? {
+            Initializer::Import { name, field, index } => Some((name, field, self.type_of(*index))),
+        }
+    }
+
     /// Returns the type of an item based on its index
     pub fn type_of(&self, index: EntityIndex) -> EntityType {
         match index {