Shrink the size of the anyfunc table in `VMContext` (#3850)

* Shrink the size of the anyfunc table in `VMContext`

This commit shrinks the size of the `VMCallerCheckedAnyfunc` table
allocated into a `VMContext` to be the size of the number of "escaped"
functions in a module rather than the number of functions in a module.
Escaped functions include exports, table elements, etc, and are
typically an order of magnitude smaller than the number of functions in
general. This should greatly shrink the `VMContext` for some modules
which while we aren't necessarily having any problems with that today
shouldn't cause any problems in the future.

The original motivation for this was that this came up during the recent
lazy-table-initialization work and while it no longer has a direct
performance benefit since tables aren't initialized at all on
instantiation it should still improve long-running instances
theoretically with smaller `VMContext` allocations as well as better
locality between anyfuncs.

* Fix some tests

* Remove redundant hash set

* Use a helper for pushing function type information

* Use a more descriptive `is_escaping` method

* Clarify a comment

* Fix condition

crates/cranelift/src/lib.rs

@@ -256,12 +256,13 @@ fn func_signature(
     types: &TypeTables,
     index: FuncIndex,
 ) -> ir::Signature {
+    let func = &translation.module.functions[index];
     let call_conv = match translation.module.defined_func_index(index) {
-        // If this is a defined function in the module and it's never possibly
-        // exported, then we can optimize this function to use the fastest
-        // calling convention since it's purely an internal implementation
-        // detail of the module itself.
-        Some(idx) if !translation.escaped_funcs.contains(&idx) => CallConv::Fast,
+        // If this is a defined function in the module and it doesn't escape
+        // then we can optimize this function to use the fastest calling
+        // convention since it's purely an internal implementation detail of
+        // the module itself.
+        Some(_idx) if !func.is_escaping() => CallConv::Fast,
 
         // ... otherwise if it's an imported function or if it's a possibly
         // exported function then we use the default ABI wasmtime would
@@ -269,11 +270,7 @@ fn func_signature(
         _ => wasmtime_call_conv(isa),
     };
     let mut sig = blank_sig(isa, call_conv);
-    push_types(
-        isa,
-        &mut sig,
-        &types.wasm_signatures[translation.module.functions[index]],
-    );
+    push_types(isa, &mut sig, &types.wasm_signatures[func.signature]);
     return sig;
 }
 

crates/environ/src/module.rs

@@ -897,8 +897,15 @@ pub struct Module {
     /// Number of imported or aliased globals in the module.
     pub num_imported_globals: usize,
 
+    /// Number of functions that "escape" from this module may need to have a
+    /// `VMCallerCheckedAnyfunc` constructed for them.
+    ///
+    /// This is also the number of functions in the `functions` array below with
+    /// an `anyfunc` index (and is the maximum anyfunc index).
+    pub num_escaped_funcs: usize,
+
     /// Types of functions, imported and local.
-    pub functions: PrimaryMap<FuncIndex, SignatureIndex>,
+    pub functions: PrimaryMap<FuncIndex, FunctionType>,
 
     /// WebAssembly tables.
     pub table_plans: PrimaryMap<TableIndex, TablePlan>,
@@ -1108,11 +1115,30 @@ impl Module {
             EntityIndex::Global(i) => EntityType::Global(self.globals[i]),
             EntityIndex::Table(i) => EntityType::Table(self.table_plans[i].table),
             EntityIndex::Memory(i) => EntityType::Memory(self.memory_plans[i].memory),
-            EntityIndex::Function(i) => EntityType::Function(self.functions[i]),
+            EntityIndex::Function(i) => EntityType::Function(self.functions[i].signature),
             EntityIndex::Instance(i) => EntityType::Instance(self.instances[i]),
             EntityIndex::Module(i) => EntityType::Module(self.modules[i]),
         }
     }
+
+    /// Appends a new function to this module with the given type information,
+    /// used for functions that either don't escape or aren't certain whether
+    /// they escape yet.
+    pub fn push_function(&mut self, signature: SignatureIndex) -> FuncIndex {
+        self.functions.push(FunctionType {
+            signature,
+            anyfunc: AnyfuncIndex::reserved_value(),
+        })
+    }
+
+    /// Appends a new function to this module with the given type information.
+    pub fn push_escaped_function(
+        &mut self,
+        signature: SignatureIndex,
+        anyfunc: AnyfuncIndex,
+    ) -> FuncIndex {
+        self.functions.push(FunctionType { signature, anyfunc })
+    }
 }
 
 /// All types which are recorded for the entirety of a translation.
@@ -1144,3 +1170,28 @@ pub struct InstanceSignature {
     /// The name of what's being exported as well as its type signature.
     pub exports: IndexMap<String, EntityType>,
 }
+
+/// Type information about functions in a wasm module.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct FunctionType {
+    /// The type of this function, indexed into the module-wide type tables for
+    /// a module compilation.
+    pub signature: SignatureIndex,
+    /// The index into the anyfunc table, if present. Note that this is
+    /// `reserved_value()` if the function does not escape from a module.
+    pub anyfunc: AnyfuncIndex,
+}
+
+impl FunctionType {
+    /// Returns whether this function's type is one that "escapes" the current
+    /// module, meaning that the function is exported, used in `ref.func`, used
+    /// in a table, etc.
+    pub fn is_escaping(&self) -> bool {
+        !self.anyfunc.is_reserved_value()
+    }
+}
+
+/// Index into the anyfunc table within a VMContext for a function.
+#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
+pub struct AnyfuncIndex(u32);
+cranelift_entity::entity_impl!(AnyfuncIndex);

crates/environ/src/module_environ.rs

@@ -1,6 +1,7 @@
 use crate::module::{
-    Initializer, InstanceSignature, MemoryInitialization, MemoryInitializer, MemoryPlan, Module,
-    ModuleSignature, ModuleType, ModuleUpvar, TableInitializer, TablePlan, TypeTables,
+    AnyfuncIndex, Initializer, InstanceSignature, MemoryInitialization, MemoryInitializer,
+    MemoryPlan, Module, ModuleSignature, ModuleType, ModuleUpvar, TableInitializer, TablePlan,
+    TypeTables,
 };
 use crate::{
     DataIndex, DefinedFuncIndex, ElemIndex, EntityIndex, EntityType, FuncIndex, Global,
@@ -10,7 +11,7 @@ use crate::{
 };
 use cranelift_entity::packed_option::ReservedValue;
 use std::borrow::Cow;
-use std::collections::{hash_map::Entry, HashMap, HashSet};
+use std::collections::{hash_map::Entry, HashMap};
 use std::convert::{TryFrom, TryInto};
 use std::mem;
 use std::path::PathBuf;
@@ -61,18 +62,6 @@ pub struct ModuleTranslation<'data> {
     /// References to the function bodies.
     pub function_body_inputs: PrimaryMap<DefinedFuncIndex, FunctionBodyData<'data>>,
 
-    /// The set of defined functions within this module which are "possibly
-    /// exported" which means that the host can possibly call them. This
-    /// includes functions such as:
-    ///
-    /// * Exported functions
-    /// * Functions in element segments
-    /// * Functions via `ref.func` instructions
-    ///
-    /// This set is used to determine the set of type signatures that need
-    /// trampolines for the host to call into.
-    pub escaped_funcs: HashSet<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.
@@ -254,9 +243,12 @@ impl<'data> ModuleEnvironment<'data> {
                     .module
                     .functions
                     .iter()
-                    .filter_map(|(i, sig)| match self.result.module.defined_func_index(i) {
-                        Some(i) if !self.result.escaped_funcs.contains(&i) => None,
-                        _ => Some(*sig),
+                    .filter_map(|(_, func)| {
+                        if func.is_escaping() {
+                            Some(func.signature)
+                        } else {
+                            None
+                        }
                     })
                     .collect();
                 self.result.exported_signatures.sort_unstable();
@@ -398,7 +390,7 @@ impl<'data> ModuleEnvironment<'data> {
                     let sigindex = entry?;
                     let ty = TypeIndex::from_u32(sigindex);
                     let sig_index = self.result.module.types[ty].unwrap_function();
-                    self.result.module.functions.push(sig_index);
+                    self.result.module.push_function(sig_index);
                 }
             }
 
@@ -624,7 +616,7 @@ impl<'data> ModuleEnvironment<'data> {
                 let func_index = FuncIndex::from_u32(func_index);
 
                 if self.tunables.generate_native_debuginfo {
-                    let sig_index = self.result.module.functions[func_index];
+                    let sig_index = self.result.module.functions[func_index].signature;
                     let sig = &self.types.wasm_signatures[sig_index];
                     let mut locals = Vec::new();
                     for pair in body.get_locals_reader()? {
@@ -882,7 +874,7 @@ impl<'data> ModuleEnvironment<'data> {
                                     self.result.module.num_imported_tables += 1;
                                 }
                                 EntityType::Function(sig) => {
-                                    self.result.module.functions.push(*sig);
+                                    self.result.module.push_function(*sig);
                                     self.result.module.num_imported_funcs += 1;
                                     self.result.debuginfo.wasm_file.imported_func_count += 1;
                                 }
@@ -1133,9 +1125,7 @@ and for re-adding support for interface types you can see this issue:
 
     fn push_type(&mut self, ty: EntityType) -> EntityIndex {
         match ty {
-            EntityType::Function(ty) => {
-                EntityIndex::Function(self.result.module.functions.push(ty))
-            }
+            EntityType::Function(ty) => EntityIndex::Function(self.result.module.push_function(ty)),
             EntityType::Table(ty) => {
                 let plan = TablePlan::for_table(ty, &self.tunables);
                 EntityIndex::Table(self.result.module.table_plans.push(plan))
@@ -1179,9 +1169,14 @@ and for re-adding support for interface types you can see this issue:
     }
 
     fn flag_func_escaped(&mut self, func: FuncIndex) {
-        if let Some(idx) = self.result.module.defined_func_index(func) {
-            self.result.escaped_funcs.insert(idx);
+        let ty = &mut self.result.module.functions[func];
+        // If this was already assigned an anyfunc index no need to re-assign it.
+        if ty.is_escaping() {
+            return;
         }
+        let index = self.result.module.num_escaped_funcs as u32;
+        ty.anyfunc = AnyfuncIndex::from_u32(index);
+        self.result.module.num_escaped_funcs += 1;
     }
 
     fn declare_type_func(&mut self, wasm: WasmFuncType) -> WasmResult<()> {

crates/environ/src/vmoffsets.rs

@@ -16,13 +16,14 @@
 //      tables: [VMTableDefinition; module.num_defined_tables],
 //      memories: [VMMemoryDefinition; module.num_defined_memories],
 //      globals: [VMGlobalDefinition; module.num_defined_globals],
-//      anyfuncs: [VMCallerCheckedAnyfunc; module.num_imported_functions + module.num_defined_functions],
+//      anyfuncs: [VMCallerCheckedAnyfunc; module.num_escaped_funcs],
 // }
 
 use crate::{
-    DefinedGlobalIndex, DefinedMemoryIndex, DefinedTableIndex, FuncIndex, GlobalIndex, MemoryIndex,
-    Module, TableIndex,
+    AnyfuncIndex, DefinedGlobalIndex, DefinedMemoryIndex, DefinedTableIndex, FuncIndex,
+    GlobalIndex, MemoryIndex, Module, TableIndex,
 };
+use cranelift_entity::packed_option::ReservedValue;
 use more_asserts::assert_lt;
 use std::convert::TryFrom;
 
@@ -68,6 +69,9 @@ pub struct VMOffsets<P> {
     pub num_defined_memories: u32,
     /// The number of defined globals in the module.
     pub num_defined_globals: u32,
+    /// The number of escaped functions in the module, the size of the anyfuncs
+    /// array.
+    pub num_escaped_funcs: u32,
 
     // precalculated offsets of various member fields
     interrupts: u32,
@@ -131,6 +135,9 @@ pub struct VMOffsetsFields<P> {
     pub num_defined_memories: u32,
     /// The number of defined globals in the module.
     pub num_defined_globals: u32,
+    /// The numbe of escaped functions in the module, the size of the anyfunc
+    /// array.
+    pub num_escaped_funcs: u32,
 }
 
 impl<P: PtrSize> VMOffsets<P> {
@@ -146,6 +153,7 @@ impl<P: PtrSize> VMOffsets<P> {
             num_defined_tables: cast_to_u32(module.table_plans.len()),
             num_defined_memories: cast_to_u32(module.memory_plans.len()),
             num_defined_globals: cast_to_u32(module.globals.len()),
+            num_escaped_funcs: cast_to_u32(module.num_escaped_funcs),
         })
     }
 
@@ -174,6 +182,7 @@ impl<P: PtrSize> VMOffsets<P> {
                     num_defined_globals: _,
                     num_defined_memories: _,
                     num_defined_functions: _,
+                    num_escaped_funcs: _,
 
                     // used as the initial size below
                     size,
@@ -229,6 +238,7 @@ impl<P: PtrSize> From<VMOffsetsFields<P>> for VMOffsets<P> {
             num_defined_tables: fields.num_defined_tables,
             num_defined_memories: fields.num_defined_memories,
             num_defined_globals: fields.num_defined_globals,
+            num_escaped_funcs: fields.num_escaped_funcs,
             interrupts: 0,
             epoch_ptr: 0,
             externref_activations_table: 0,
@@ -298,7 +308,7 @@ impl<P: PtrSize> From<VMOffsetsFields<P>> for VMOffsets<P> {
             size(defined_globals)
                 = cmul(ret.num_defined_globals, ret.size_of_vmglobal_definition()),
             size(defined_anyfuncs) = cmul(
-                cadd(ret.num_imported_functions, ret.num_defined_functions),
+                ret.num_escaped_funcs,
                 ret.size_of_vmcaller_checked_anyfunc(),
             ),
         }
@@ -674,11 +684,9 @@ impl<P: PtrSize> VMOffsets<P> {
     /// Return the offset to the `VMCallerCheckedAnyfunc` for the given function
     /// index (either imported or defined).
     #[inline]
-    pub fn vmctx_anyfunc(&self, index: FuncIndex) -> u32 {
-        assert_lt!(
-            index.as_u32(),
-            self.num_imported_functions + self.num_defined_functions
-        );
+    pub fn vmctx_anyfunc(&self, index: AnyfuncIndex) -> u32 {
+        assert!(!index.is_reserved_value());
+        assert_lt!(index.as_u32(), self.num_escaped_funcs);
         self.vmctx_anyfuncs_begin()
             + index.as_u32() * u32::from(self.size_of_vmcaller_checked_anyfunc())
     }

crates/runtime/src/externref.rs

@@ -1045,6 +1045,7 @@ mod tests {
             num_defined_tables: 0,
             num_defined_memories: 0,
             num_defined_globals: 0,
+            num_escaped_funcs: 0,
         });
         assert_eq!(
             offsets.vm_extern_data_ref_count(),
@@ -1071,6 +1072,7 @@ mod tests {
             num_defined_tables: 0,
             num_defined_memories: 0,
             num_defined_globals: 0,
+            num_escaped_funcs: 0,
         });
         assert_eq!(
             offsets.vm_extern_ref_activation_table_next() as usize,
@@ -1097,6 +1099,7 @@ mod tests {
             num_defined_tables: 0,
             num_defined_memories: 0,
             num_defined_globals: 0,
+            num_escaped_funcs: 0,
         });
         assert_eq!(
             offsets.vm_extern_ref_activation_table_end() as usize,

crates/runtime/src/instance.rs

@@ -30,9 +30,9 @@ use std::{mem, ptr, slice};
 use wasmtime_environ::{
     packed_option::ReservedValue, DataIndex, DefinedGlobalIndex, DefinedMemoryIndex,
     DefinedTableIndex, ElemIndex, EntityIndex, EntityRef, EntitySet, FuncIndex, GlobalIndex,
-    HostPtr, MemoryIndex, Module, PrimaryMap, TableIndex, TrapCode, VMOffsets, WasmType,
+    GlobalInit, HostPtr, MemoryIndex, Module, PrimaryMap, SignatureIndex, TableIndex,
+    TableInitialization, TrapCode, VMOffsets, WasmType,
 };
-use wasmtime_environ::{GlobalInit, TableInitialization};
 
 mod allocator;
 
@@ -483,8 +483,12 @@ impl Instance {
     /// than tracking state related to whether it's been initialized
     /// before, because resetting that state on (re)instantiation is
     /// very expensive if there are many anyfuncs.
-    fn construct_anyfunc(&mut self, index: FuncIndex, into: *mut VMCallerCheckedAnyfunc) {
-        let sig = self.module().functions[index];
+    fn construct_anyfunc(
+        &mut self,
+        index: FuncIndex,
+        sig: SignatureIndex,
+        into: *mut VMCallerCheckedAnyfunc,
+    ) {
         let type_index = self.runtime_info.signature(sig);
 
         let (func_ptr, vmctx) = if let Some(def_index) = self.module().defined_func_index(index) {
@@ -551,9 +555,13 @@ impl Instance {
             // expensive, so it's better for instantiation performance
             // if we don't have to track "is-initialized" state at
             // all!
-            let anyfunc: *mut VMCallerCheckedAnyfunc =
-                self.vmctx_plus_offset::<VMCallerCheckedAnyfunc>(self.offsets.vmctx_anyfunc(index));
-            self.construct_anyfunc(index, anyfunc);
+            let func = &self.module().functions[index];
+            let sig = func.signature;
+            let anyfunc: *mut VMCallerCheckedAnyfunc = self
+                .vmctx_plus_offset::<VMCallerCheckedAnyfunc>(
+                    self.offsets.vmctx_anyfunc(func.anyfunc),
+                );
+            self.construct_anyfunc(index, sig, anyfunc);
 
             Some(anyfunc)
         }

crates/wasmtime/src/trampoline/func.rs

@@ -6,7 +6,9 @@ use anyhow::Result;
 use std::any::Any;
 use std::panic::{self, AssertUnwindSafe};
 use std::sync::Arc;
-use wasmtime_environ::{EntityIndex, FunctionInfo, Module, ModuleType, SignatureIndex};
+use wasmtime_environ::{
+    AnyfuncIndex, EntityIndex, FunctionInfo, Module, ModuleType, SignatureIndex,
+};
 use wasmtime_jit::{CodeMemory, ProfilingAgent};
 use wasmtime_runtime::{
     Imports, InstanceAllocationRequest, InstanceAllocator, InstanceHandle,
@@ -152,7 +154,8 @@ pub unsafe fn create_raw_function(
 
     let sig_id = SignatureIndex::from_u32(u32::max_value() - 1);
     module.types.push(ModuleType::Function(sig_id));
-    let func_id = module.functions.push(sig_id);
+    let func_id = module.push_escaped_function(sig_id, AnyfuncIndex::from_u32(0));
+    module.num_escaped_funcs = 1;
     module
         .exports
         .insert(String::new(), EntityIndex::Function(func_id));

crates/wasmtime/src/trampoline/global.rs

@@ -2,7 +2,9 @@ use crate::store::{InstanceId, StoreOpaque};
 use crate::trampoline::create_handle;
 use crate::{GlobalType, Mutability, Val};
 use anyhow::Result;
-use wasmtime_environ::{EntityIndex, Global, GlobalInit, Module, ModuleType, SignatureIndex};
+use wasmtime_environ::{
+    AnyfuncIndex, EntityIndex, Global, GlobalInit, Module, ModuleType, SignatureIndex,
+};
 use wasmtime_runtime::VMFunctionImport;
 
 pub fn create_global(store: &mut StoreOpaque, gt: &GlobalType, val: Val) -> Result<InstanceId> {
@@ -40,8 +42,9 @@ pub fn create_global(store: &mut StoreOpaque, gt: &GlobalType, val: Val) -> Resu
                 let sig_id = SignatureIndex::from_u32(u32::max_value() - 1);
                 one_signature = Some((sig_id, f.type_index));
                 module.types.push(ModuleType::Function(sig_id));
-                let func_index = module.functions.push(sig_id);
+                let func_index = module.push_escaped_function(sig_id, AnyfuncIndex::from_u32(0));
                 module.num_imported_funcs = 1;
+                module.num_escaped_funcs = 1;
                 module
                     .initializers
                     .push(wasmtime_environ::Initializer::Import {