Improve signature lookup happening during instantiation (#2818)

This commit is intended to be a perf improvement for instantiation of
modules with lots of functions. Previously the `lookup_shared_signature`
callback was showing up quite high in profiles as part of instantiation.

As some background, this callback is used to translate from a module's
`SignatureIndex` to a `VMSharedSignatureIndex` which the instance
stores. This callback is called for two reasons, one is to translate all
of the module's own types into `VMSharedSignatureIndex` for the purposes
of `call_indirect` (the translation of that loads from this table to
compare indices). The second reason is that a `VMCallerCheckedAnyfunc`
is prepared for all functions and this embeds a `VMSharedSignatureIndex`
inside of it.

The slow part today is that the lookup callback was called
once-per-function and each lookup involved hashing a full
`WasmFuncType`. Albeit our hash algorithm is still Rust's default
SipHash algorithm which is quite slow, but we also shouldn't need to
re-hash each signature if we see it multiple times anyway.

The fix applied in this commit is to change this lookup callback to an
`enum` where one variant is that there's a table to lookup from. This
table is a `PrimaryMap` which means that lookup is quite fast. The only
thing we need to do is to prepare the table ahead of time. Currently
this happens on the instantiation path because in my measurments the
creation of the table is quite fast compared to the rest of
instantiation. If this becomes an issue, though, we can look into
creating the table as part of `SigRegistry::register_module` and caching
it somewhere (I'm not entirely sure where but I'm sure we can figure it
out).

There's in generally not a ton of efficiency around the `SigRegistry`
type. I'm hoping though that this fixes the next-lowest-hanging-fruit in
terms of performance without complicating the implementation too much. I
tried a few variants and this change seemed like the best balance
between simplicity and still a nice performance gain.

Locally I measured an improvement in instantiation time for a large-ish
module by reducing the time from ~3ms to ~2.6ms per instance.

crates/runtime/src/externref.rs

@@ -828,16 +828,6 @@ impl StackMapRegistry {
 
         let mut inner = self.inner.borrow_mut();
 
-        // Check if we've already registered this module.
-        if let Some(existing_module) = inner.ranges.get(&max) {
-            assert_eq!(existing_module.range, module_stack_maps.range);
-            debug_assert_eq!(
-                existing_module.pc_to_stack_map,
-                module_stack_maps.pc_to_stack_map,
-            );
-            return;
-        }
-
         // Assert that this chunk of ranges doesn't collide with any other known
         // chunks.
         if let Some((_, prev)) = inner.ranges.range(max..).next() {

crates/runtime/src/instance/allocator.rs

@@ -45,8 +45,8 @@ pub struct InstanceAllocationRequest<'a> {
     /// The imports to use for the instantiation.
     pub imports: Imports<'a>,
 
-    /// A callback for looking up shared signature indexes.
-    pub lookup_shared_signature: &'a dyn Fn(SignatureIndex) -> VMSharedSignatureIndex,
+    /// Translation from `SignatureIndex` to `VMSharedSignatureIndex`
+    pub shared_signatures: SharedSignatures<'a>,
 
     /// The host state to associate with the instance.
     pub host_state: Box<dyn Any>,
@@ -165,6 +165,46 @@ pub unsafe trait InstanceAllocator: Send + Sync {
     unsafe fn deallocate_fiber_stack(&self, stack: &wasmtime_fiber::FiberStack);
 }
 
+pub enum SharedSignatures<'a> {
+    /// Used for instantiating user-defined modules
+    Table(&'a PrimaryMap<SignatureIndex, VMSharedSignatureIndex>),
+    /// Used for instance creation that has only a single function
+    Always(VMSharedSignatureIndex),
+    /// Used for instance creation that has no functions
+    None,
+}
+
+impl SharedSignatures<'_> {
+    fn lookup(&self, index: SignatureIndex) -> VMSharedSignatureIndex {
+        match self {
+            SharedSignatures::Table(table) => table[index],
+            SharedSignatures::Always(index) => *index,
+            SharedSignatures::None => unreachable!(),
+        }
+    }
+}
+
+impl<'a> From<VMSharedSignatureIndex> for SharedSignatures<'a> {
+    fn from(val: VMSharedSignatureIndex) -> SharedSignatures<'a> {
+        SharedSignatures::Always(val)
+    }
+}
+
+impl<'a> From<Option<VMSharedSignatureIndex>> for SharedSignatures<'a> {
+    fn from(val: Option<VMSharedSignatureIndex>) -> SharedSignatures<'a> {
+        match val {
+            Some(idx) => SharedSignatures::Always(idx),
+            None => SharedSignatures::None,
+        }
+    }
+}
+
+impl<'a> From<&'a PrimaryMap<SignatureIndex, VMSharedSignatureIndex>> for SharedSignatures<'a> {
+    fn from(val: &'a PrimaryMap<SignatureIndex, VMSharedSignatureIndex>) -> SharedSignatures<'a> {
+        SharedSignatures::Table(val)
+    }
+}
+
 fn get_table_init_start(
     init: &TableInitializer,
     instance: &Instance,
@@ -413,7 +453,7 @@ unsafe fn initialize_vmcontext(instance: &Instance, req: InstanceAllocationReque
     let mut ptr = instance.signature_ids_ptr();
     for sig in module.types.values() {
         *ptr = match sig {
-            ModuleType::Function(sig) => (req.lookup_shared_signature)(*sig),
+            ModuleType::Function(sig) => req.shared_signatures.lookup(*sig),
             _ => VMSharedSignatureIndex::new(u32::max_value()),
         };
         ptr = ptr.add(1);
@@ -453,7 +493,7 @@ unsafe fn initialize_vmcontext(instance: &Instance, req: InstanceAllocationReque
 
     // Initialize the functions
     for (index, sig) in instance.module.functions.iter() {
-        let type_index = (req.lookup_shared_signature)(*sig);
+        let type_index = req.shared_signatures.lookup(*sig);
 
         let (func_ptr, vmctx) = if let Some(def_index) = instance.module.defined_func_index(index) {
             (

crates/runtime/src/instance/allocator/pooling.rs

@@ -1366,7 +1366,7 @@ mod test {
                                 memories: &[],
                                 globals: &[],
                             },
-                            lookup_shared_signature: &|_| VMSharedSignatureIndex::default(),
+                            shared_signatures: VMSharedSignatureIndex::default().into(),
                             host_state: Box::new(()),
                             interrupts: std::ptr::null(),
                             externref_activations_table: std::ptr::null_mut(),
@@ -1390,7 +1390,7 @@ mod test {
                     memories: &[],
                     globals: &[],
                 },
-                lookup_shared_signature: &|_| VMSharedSignatureIndex::default(),
+                shared_signatures: VMSharedSignatureIndex::default().into(),
                 host_state: Box::new(()),
                 interrupts: std::ptr::null(),
                 externref_activations_table: std::ptr::null_mut(),

crates/runtime/src/instance/allocator/pooling/uffd.rs

@@ -519,7 +519,7 @@ mod test {
                                     memories: &[],
                                     globals: &[],
                                 },
-                                lookup_shared_signature: &|_| VMSharedSignatureIndex::default(),
+                                shared_signatures: VMSharedSignatureIndex::default().into(),
                                 host_state: Box::new(()),
                                 interrupts: ptr::null(),
                                 externref_activations_table: ptr::null_mut(),