Wasmtime: Add (optional) bottom-up function inlining to Wasm compilation

This commit plumbs together two pieces of recently-added infrastructure:

1. function inlining in Cranelift, and
2. the parallel bottom-up inlining scheduler in Wasmtime.

Sprinkle some very simple inlining heuristics on top, and this gives us function
inlining in Wasm compilation.

The default Wasmtime configuration does not enable inlining, and when we do
enable it, we only enable it for cross-component calls by default (since
presumably the toolchain that produced a particular core Wasm module, like LLVM,
already performed any inlining that was beneficial within that module, but that
toolchain couldn't know how that Wasm module would be getting linked together
with other modules via component composition, and so it could not have done any
cross-component inlining). For what it is worth, there is a config knob to
enable intra-module function inlining, but this is primarily for use by our
fuzzers, so that they can easily excercise and explore this new inlining
functionality.

All this plumbing required some changes to the `wasmtime_environ::Compiler`
trait, since Winch cannot do inlining but Cranelift can. This is mostly
encapsulated in the new `wasmtime_environ::InliningCompiler` trait, for the most
part. Additionally, we take care not to construct the call graph, or any other
data structures required only by the inliner and not regular compilation, both
when using Winch and when using Cranelift with inlining disabled.

Finally, we add a `disas` test to verify that we successfully inline a series of
calls from a function in one component, to a cross-component adapter function,
to a function in another component. Most test coverage is expected to come from
our fuzzing, however.

Author: fitzgen

cranelift/codegen/src/inline.rs

@@ -73,7 +73,7 @@ pub trait Inline {
     /// Failure to uphold these invariants may result in panics during
     /// compilation or incorrect runtime behavior in the generated code.
     fn inline(
-        &self,
+        &mut self,
         caller: &ir::Function,
         call_inst: ir::Inst,
         call_opcode: ir::Opcode,
@@ -82,12 +82,12 @@ pub trait Inline {
     ) -> InlineCommand<'_>;
 }
 
-impl<'a, T> Inline for &'a T
+impl<'a, T> Inline for &'a mut T
 where
     T: Inline,
 {
     fn inline(
-        &self,
+        &mut self,
         caller: &ir::Function,
         inst: ir::Inst,
         opcode: ir::Opcode,
@@ -102,7 +102,12 @@ where
 /// instruction, and inline the callee when directed to do so.
 ///
 /// Returns whether any call was inlined.
-pub(crate) fn do_inlining(func: &mut ir::Function, inliner: impl Inline) -> CodegenResult<bool> {
+pub(crate) fn do_inlining(
+    func: &mut ir::Function,
+    mut inliner: impl Inline,
+) -> CodegenResult<bool> {
+    trace!("function {} before inlining: {}", func.name, func);
+
     let mut inlined_any = false;
     let mut allocs = InliningAllocs::default();
 
@@ -175,6 +180,12 @@ pub(crate) fn do_inlining(func: &mut ir::Function, inliner: impl Inline) -> Code
         }
     }
 
+    if inlined_any {
+        trace!("function {} after inlining: {}", func.name, func);
+    } else {
+        trace!("function {} did not have any callees inlined", func.name);
+    }
+
     Ok(inlined_any)
 }
 

cranelift/codegen/src/ir/dfg.rs

@@ -333,6 +333,16 @@ impl<'a> Iterator for Values<'a> {
             .find(|kv| valid_valuedata(*kv.1))
             .map(|kv| kv.0)
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.inner.size_hint()
+    }
+}
+
+impl ExactSizeIterator for Values<'_> {
+    fn len(&self) -> usize {
+        self.inner.len()
+    }
 }
 
 /// Handling values.

cranelift/entity/src/primary.rs

@@ -306,6 +306,15 @@ where
     }
 }
 
+impl<K, V> From<PrimaryMap<K, V>> for Vec<V>
+where
+    K: EntityRef,
+{
+    fn from(map: PrimaryMap<K, V>) -> Self {
+        map.elems
+    }
+}
+
 impl<K: EntityRef + fmt::Debug, V: fmt::Debug> fmt::Debug for PrimaryMap<K, V> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut struct_ = f.debug_struct("PrimaryMap");

cranelift/filetests/src/test_inline.rs

@@ -121,7 +121,7 @@ struct Inliner<'a>(Ref<'a, HashMap<ir::UserFuncName, ir::Function>>);
 
 impl<'a> Inline for Inliner<'a> {
     fn inline(
-        &self,
+        &mut self,
         caller: &ir::Function,
         _inst: ir::Inst,
         _opcode: ir::Opcode,

crates/cli-flags/src/lib.rs

@@ -111,6 +111,27 @@ wasmtime_option_group! {
         /// dense.
         pub memory_guaranteed_dense_image_size: Option<u64>,
 
+        /// Whether to perform function inlining during compilation.
+        pub compiler_inlining: Option<bool>,
+
+        /// Whether to perform function inlining within the same core Wasm
+        /// module or only for calls from one module into a different
+        /// module. Only exposed for fuzzing.
+        #[doc(hidden)]
+        #[serde(default)]
+        #[serde(deserialize_with = "crate::opt::cli_parse_wrapper")]
+        pub compiler_inlining_intra_module: Option<wasmtime::IntraModuleInlining>,
+
+        /// Configure the "small-callee" size for function inlining
+        /// heuristics. Only exposed for fuzzing.
+        #[doc(hidden)]
+        pub compiler_inlining_small_callee_size: Option<u32>,
+
+        /// Configure the "sum size threshold" for function inlining
+        /// heuristics. Only exposed for fuzzing.
+        #[doc(hidden)]
+        pub compiler_inlining_sum_size_threshold: Option<u32>,
+
         /// The maximum number of WebAssembly instances which can be created
         /// with the pooling allocator.
         pub pooling_total_core_instances: Option<u32>,
@@ -816,6 +837,18 @@ impl CommonOptions {
         if let Some(size) = self.opts.memory_guaranteed_dense_image_size {
             config.memory_guaranteed_dense_image_size(size);
         }
+        if let Some(enable) = self.opts.compiler_inlining {
+            config.compiler_inlining(enable);
+        }
+        if let Some(cfg) = self.opts.compiler_inlining_intra_module {
+            config.compiler_inlining_intra_module(cfg);
+        }
+        if let Some(size) = self.opts.compiler_inlining_small_callee_size {
+            config.compiler_inlining_small_callee_size(size);
+        }
+        if let Some(size) = self.opts.compiler_inlining_sum_size_threshold {
+            config.compiler_inlining_sum_size_threshold(size);
+        }
         if let Some(enable) = self.opts.signals_based_traps {
             config.signals_based_traps(enable);
         }

crates/cli-flags/src/opt.rs

@@ -23,6 +23,7 @@ macro_rules! wasmtime_option_group {
         pub struct $opts:ident {
             $(
                 $(#[doc = $doc:tt])*
+                $(#[doc($doc_attr:meta)])?
                 $(#[serde($serde_attr:meta)])*
                 pub $opt:ident: $container:ident<$payload:ty>,
             )+
@@ -31,6 +32,7 @@ macro_rules! wasmtime_option_group {
                 #[prefixed = $prefix:tt]
                 $(#[serde($serde_attr2:meta)])*
                 $(#[doc = $prefixed_doc:tt])*
+                $(#[doc($prefixed_doc_attr:meta)])?
                 pub $prefixed:ident: Vec<(String, Option<String>)>,
             )?
         }
@@ -43,6 +45,7 @@ macro_rules! wasmtime_option_group {
         pub struct $opts {
             $(
                 $(#[serde($serde_attr)])*
+                $(#[doc($doc_attr)])?
                 pub $opt: $container<$payload>,
             )+
             $(
@@ -555,6 +558,30 @@ impl WasmtimeOptionValue for wasmtime::MpkEnabled {
     }
 }
 
+impl WasmtimeOptionValue for wasmtime::IntraModuleInlining {
+    const VAL_HELP: &'static str = "[=y|n|gc]";
+
+    fn parse(val: Option<&str>) -> Result<Self> {
+        match val {
+            None | Some("n") | Some("no") | Some("false") => Ok(wasmtime::IntraModuleInlining::Yes),
+            Some("y") | Some("yes") | Some("true") => Ok(wasmtime::IntraModuleInlining::No),
+            Some("gc") => Ok(wasmtime::IntraModuleInlining::WhenUsingGc),
+            Some(s) => bail!(
+                "unknown compiler intra-module inlining flag `{s}`, \
+                 only yes,no,gc,<nothing> accepted"
+            ),
+        }
+    }
+
+    fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            wasmtime::IntraModuleInlining::Yes => f.write_str("y"),
+            wasmtime::IntraModuleInlining::No => f.write_str("n"),
+            wasmtime::IntraModuleInlining::WhenUsingGc => f.write_str("gc"),
+        }
+    }
+}
+
 impl WasmtimeOptionValue for WasiNnGraph {
     const VAL_HELP: &'static str = "=<format>::<dir>";
     fn parse(val: Option<&str>) -> Result<Self> {