Analyze whether a module requires a GC heap during its compilation

And avoid forcing the allocation of a GC heap for modules that will not use it.

Split out from bytecodealliance#9350

Author: fitzgen

crates/cranelift/src/compiler.rs

@@ -27,10 +27,10 @@ use std::path;
 use std::sync::{Arc, Mutex};
 use wasmparser::{FuncValidatorAllocations, FunctionBody};
 use wasmtime_environ::{
-    AddressMapSection, BuiltinFunctionIndex, CacheStore, CompileError, DefinedFuncIndex, FlagValue,
-    FunctionBodyData, FunctionLoc, HostCall, ModuleTranslation, ModuleTypesBuilder, PtrSize,
-    RelocationTarget, StackMapSection, StaticModuleIndex, TrapEncodingBuilder, TrapSentinel,
-    TripleExt, Tunables, VMOffsets, WasmFuncType, WasmValType,
+    AddressMapSection, BuiltinFunctionIndex, CacheStore, CompileError, CompiledFunctionBody,
+    DefinedFuncIndex, FlagValue, FunctionBodyData, FunctionLoc, HostCall, ModuleTranslation,
+    ModuleTypesBuilder, PtrSize, RelocationTarget, StackMapSection, StaticModuleIndex,
+    TrapEncodingBuilder, TrapSentinel, TripleExt, Tunables, VMOffsets, WasmFuncType, WasmValType,
 };
 
 #[cfg(feature = "component-model")]
@@ -187,7 +187,7 @@ impl wasmtime_environ::Compiler for Compiler {
         func_index: DefinedFuncIndex,
         input: FunctionBodyData<'_>,
         types: &ModuleTypesBuilder,
-    ) -> Result<Box<dyn Any + Send>, CompileError> {
+    ) -> Result<CompiledFunctionBody, CompileError> {
         let isa = &*self.isa;
         let module = &translation.module;
         let func_index = module.func_index(func_index);
@@ -284,15 +284,18 @@ impl wasmtime_environ::Compiler for Compiler {
         log::debug!("{:?} translated in {:?}", func_index, timing.total());
         log::trace!("{:?} timing info\n{}", func_index, timing);
 
-        Ok(Box::new(func))
+        Ok(CompiledFunctionBody {
+            code: Box::new(func),
+            needs_gc_heap: func_env.needs_gc_heap(),
+        })
     }
 
     fn compile_array_to_wasm_trampoline(
         &self,
         translation: &ModuleTranslation<'_>,
         types: &ModuleTypesBuilder,
         def_func_index: DefinedFuncIndex,
-    ) -> Result<Box<dyn Any + Send>, CompileError> {
+    ) -> Result<CompiledFunctionBody, CompileError> {
         let func_index = translation.module.func_index(def_func_index);
         let sig = translation.module.functions[func_index]
             .signature
@@ -359,16 +362,16 @@ impl wasmtime_environ::Compiler for Compiler {
         builder.ins().return_(&[true_return]);
         builder.finalize();
 
-        Ok(Box::new(compiler.finish(&format!(
-            "array_to_wasm_{}",
-            func_index.as_u32(),
-        ))?))
+        Ok(CompiledFunctionBody {
+            code: Box::new(compiler.finish(&format!("array_to_wasm_{}", func_index.as_u32(),))?),
+            needs_gc_heap: false,
+        })
     }
 
     fn compile_wasm_to_array_trampoline(
         &self,
         wasm_func_ty: &WasmFuncType,
-    ) -> Result<Box<dyn Any + Send>, CompileError> {
+    ) -> Result<CompiledFunctionBody, CompileError> {
         let isa = &*self.isa;
         let pointer_type = isa.pointer_type();
         let wasm_call_sig = wasm_call_signature(isa, wasm_func_ty, &self.tunables);
@@ -432,9 +435,10 @@ impl wasmtime_environ::Compiler for Compiler {
         builder.ins().return_(&results);
         builder.finalize();
 
-        Ok(Box::new(compiler.finish(&format!(
-            "wasm_to_array_trampoline_{wasm_func_ty}"
-        ))?))
+        Ok(CompiledFunctionBody {
+            code: Box::new(compiler.finish(&format!("wasm_to_array_trampoline_{wasm_func_ty}"))?),
+            needs_gc_heap: false,
+        })
     }
 
     fn append_code(
@@ -582,7 +586,7 @@ impl wasmtime_environ::Compiler for Compiler {
     fn compile_wasm_to_builtin(
         &self,
         index: BuiltinFunctionIndex,
-    ) -> Result<Box<dyn Any + Send>, CompileError> {
+    ) -> Result<CompiledFunctionBody, CompileError> {
         let isa = &*self.isa;
         let ptr_size = isa.pointer_bytes();
         let pointer_type = isa.pointer_type();
@@ -651,9 +655,10 @@ impl wasmtime_environ::Compiler for Compiler {
         builder.ins().return_(&results);
         builder.finalize();
 
-        Ok(Box::new(
-            compiler.finish(&format!("wasm_to_builtin_{}", index.name()))?,
-        ))
+        Ok(CompiledFunctionBody {
+            code: Box::new(compiler.finish(&format!("wasm_to_builtin_{}", index.name()))?),
+            needs_gc_heap: false,
+        })
     }
 
     fn compiled_function_relocation_targets<'a>(

crates/cranelift/src/compiler/component.rs

@@ -6,9 +6,8 @@ use cranelift_codegen::ir::condcodes::IntCC;
 use cranelift_codegen::ir::{self, InstBuilder, MemFlags, Value};
 use cranelift_codegen::isa::{CallConv, TargetIsa};
 use cranelift_frontend::FunctionBuilder;
-use std::any::Any;
-use wasmtime_environ::component::*;
 use wasmtime_environ::fact::SYNC_ENTER_FIXED_PARAMS;
+use wasmtime_environ::{component::*, CompiledFunctionBody};
 use wasmtime_environ::{
     HostCall, ModuleInternedTypeIndex, PtrSize, TrapSentinel, Tunables, WasmFuncType, WasmValType,
 };
@@ -1450,7 +1449,7 @@ impl ComponentCompiler for Compiler {
         types: &ComponentTypesBuilder,
         index: TrampolineIndex,
         tunables: &Tunables,
-    ) -> Result<AllCallFunc<Box<dyn Any + Send>>> {
+    ) -> Result<AllCallFunc<CompiledFunctionBody>> {
         let compile = |abi: Abi| -> Result<_> {
             let mut compiler = self.function_compiler();
             let mut c = TrampolineCompiler::new(
@@ -1494,10 +1493,12 @@ impl ComponentCompiler for Compiler {
             c.translate(&component.trampolines[index]);
             c.builder.finalize();
 
-            Ok(Box::new(compiler.finish(&format!(
-                "component_trampoline_{}_{abi:?}",
-                index.as_u32(),
-            ))?))
+            Ok(CompiledFunctionBody {
+                code: Box::new(
+                    compiler.finish(&format!("component_trampoline_{}_{abi:?}", index.as_u32()))?,
+                ),
+                needs_gc_heap: false,
+            })
         };
         Ok(AllCallFunc {
             wasm_call: compile(Abi::Wasm)?,

crates/cranelift/src/func_environ.rs

@@ -14,7 +14,7 @@ use cranelift_codegen::ir::types::*;
 use cranelift_codegen::ir::{self, types};
 use cranelift_codegen::ir::{ArgumentPurpose, Function, InstBuilder, MemFlags};
 use cranelift_codegen::isa::{TargetFrontendConfig, TargetIsa};
-use cranelift_entity::packed_option::ReservedValue;
+use cranelift_entity::packed_option::ReservedValue as _;
 use cranelift_entity::{EntityRef, PrimaryMap, SecondaryMap};
 use cranelift_frontend::FunctionBuilder;
 use cranelift_frontend::Variable;
@@ -98,6 +98,7 @@ pub struct FuncEnvironment<'module_environment> {
     types: &'module_environment ModuleTypesBuilder,
     wasm_func_ty: &'module_environment WasmFuncType,
     sig_ref_to_ty: SecondaryMap<ir::SigRef, Option<&'module_environment WasmFuncType>>,
+    needs_gc_heap: bool,
 
     #[cfg(feature = "gc")]
     ty_to_gc_layout: std::collections::HashMap<
@@ -187,6 +188,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             types,
             wasm_func_ty,
             sig_ref_to_ty: SecondaryMap::default(),
+            needs_gc_heap: false,
 
             #[cfg(feature = "gc")]
             ty_to_gc_layout: std::collections::HashMap::new(),
@@ -1194,6 +1196,11 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             i32::from(self.offsets.ptr.vm_func_ref_type_index()),
         )
     }
+
+    /// Does this function need a GC heap?
+    pub fn needs_gc_heap(&self) -> bool {
+        self.needs_gc_heap
+    }
 }
 
 struct Call<'a, 'func, 'module_env> {

crates/cranelift/src/func_environ/gc.rs

@@ -163,15 +163,19 @@ pub mod builtins {
                     func: &mut ir::Function,
                 ) -> WasmResult<ir::FuncRef> {
                     #[cfg(feature = "gc")]
-                    return Ok(func_env.builtin_functions.$name(func));
+                    {
+                        func_env.needs_gc_heap = true;
+                        return Ok(func_env.builtin_functions.$name(func));
+                    }
 
                     #[cfg(not(feature = "gc"))]
-                    let _ = (func, func_env);
-                    #[cfg(not(feature = "gc"))]
-                    return Err(wasmtime_environ::wasm_unsupported!(
-                        "support for Wasm GC disabled at compile time because the `gc` cargo \
-                         feature was not enabled"
-                    ));
+                    {
+                        let _ = (func, func_env);
+                        return Err(wasmtime_environ::wasm_unsupported!(
+                            "support for Wasm GC disabled at compile time because the `gc` cargo \
+                             feature was not enabled"
+                        ));
+                    }
                 }
             )*
         };

crates/cranelift/src/func_environ/gc/enabled.rs

@@ -21,7 +21,11 @@ mod drc;
 mod null;
 
 /// Get the default GC compiler.
-pub fn gc_compiler(func_env: &FuncEnvironment<'_>) -> WasmResult<Box<dyn GcCompiler>> {
+pub fn gc_compiler(func_env: &mut FuncEnvironment<'_>) -> WasmResult<Box<dyn GcCompiler>> {
+    // If this function requires a GC compiler, that is not too bad of an
+    // over-approximation for it requiring a GC heap.
+    func_env.needs_gc_heap = true;
+
     match func_env.tunables.collector {
         #[cfg(feature = "gc-drc")]
         Some(Collector::DeferredReferenceCounting) => Ok(Box::new(drc::DrcCompiler::default())),

crates/environ/src/compile/mod.rs

@@ -178,6 +178,16 @@ pub enum SettingKind {
     Preset,
 }
 
+/// The result of compiling a single function body.
+pub struct CompiledFunctionBody {
+    /// The code. This is whatever type the `Compiler` implementation wants it
+    /// to be, we just shepherd it around.
+    pub code: Box<dyn Any + Send>,
+    /// Whether the compiled function needs a GC heap to run; that is, whether
+    /// it reads a struct field, allocates, an array, or etc...
+    pub needs_gc_heap: bool,
+}
+
 /// An implementation of a compiler which can compile WebAssembly functions to
 /// machine code and perform other miscellaneous tasks needed by the JIT runtime.
 pub trait Compiler: Send + Sync {
@@ -186,19 +196,13 @@ pub trait Compiler: Send + Sync {
     /// The body of the function is available in `data` and configuration
     /// values are also passed in via `tunables`. Type information in
     /// `translation` is all relative to `types`.
-    ///
-    /// This function returns a tuple:
-    ///
-    /// 1. Metadata about the wasm function itself.
-    /// 2. The function itself, as an `Any` to get downcasted later when passed
-    ///    to `append_code`.
     fn compile_function(
         &self,
         translation: &ModuleTranslation<'_>,
         index: DefinedFuncIndex,
         data: FunctionBodyData<'_>,
         types: &ModuleTypesBuilder,
-    ) -> Result<Box<dyn Any + Send>, CompileError>;
+    ) -> Result<CompiledFunctionBody, CompileError>;
 
     /// Compile a trampoline for an array-call host function caller calling the
     /// `index`th Wasm function.
@@ -210,7 +214,7 @@ pub trait Compiler: Send + Sync {
         translation: &ModuleTranslation<'_>,
         types: &ModuleTypesBuilder,
         index: DefinedFuncIndex,
-    ) -> Result<Box<dyn Any + Send>, CompileError>;
+    ) -> Result<CompiledFunctionBody, CompileError>;
 
     /// Compile a trampoline for a Wasm caller calling a array callee with the
     /// given signature.
@@ -220,7 +224,7 @@ pub trait Compiler: Send + Sync {
     fn compile_wasm_to_array_trampoline(
         &self,
         wasm_func_ty: &WasmFuncType,
-    ) -> Result<Box<dyn Any + Send>, CompileError>;
+    ) -> Result<CompiledFunctionBody, CompileError>;
 
     /// Creates a trampoline that can be used to call Wasmtime's implementation
     /// of the builtin function specified by `index`.
@@ -234,7 +238,7 @@ pub trait Compiler: Send + Sync {
     fn compile_wasm_to_builtin(
         &self,
         index: BuiltinFunctionIndex,
-    ) -> Result<Box<dyn Any + Send>, CompileError>;
+    ) -> Result<CompiledFunctionBody, CompileError>;
 
     /// Returns the list of relocations required for a function from one of the
     /// previous `compile_*` functions above.

crates/environ/src/component/compiler.rs

@@ -1,8 +1,6 @@
 use crate::component::{AllCallFunc, ComponentTranslation, ComponentTypesBuilder, TrampolineIndex};
-use crate::prelude::*;
-use crate::Tunables;
+use crate::{CompiledFunctionBody, Tunables};
 use anyhow::Result;
-use std::any::Any;
 
 /// Compilation support necessary for components.
 pub trait ComponentCompiler: Send + Sync {
@@ -18,5 +16,5 @@ pub trait ComponentCompiler: Send + Sync {
         types: &ComponentTypesBuilder,
         trampoline: TrampolineIndex,
         tunables: &Tunables,
-    ) -> Result<AllCallFunc<Box<dyn Any + Send>>>;
+    ) -> Result<AllCallFunc<CompiledFunctionBody>>;
 }

crates/environ/src/module.rs

@@ -336,6 +336,9 @@ pub struct Module {
     /// Number of imported or aliased tags in the module.
     pub num_imported_tags: usize,
 
+    /// Does this module need a GC heap to run?
+    pub needs_gc_heap: bool,
+
     /// Number of functions that "escape" from this module may need to have a
     /// `VMFuncRef` constructed for them.
     ///
@@ -616,6 +619,7 @@ impl TypeTrace for Module {
             num_imported_globals: _,
             num_imported_tags: _,
             num_escaped_funcs: _,
+            needs_gc_heap: _,
             functions,
             tables,
             memories: _,
@@ -665,6 +669,7 @@ impl TypeTrace for Module {
             num_imported_globals: _,
             num_imported_tags: _,
             num_escaped_funcs: _,
+            needs_gc_heap: _,
             functions,
             tables,
             memories: _,