Refactor Cache logic to include debug information (bytecodealliance#2065

)

* move caching to the CompilationArtifacts

* mv cache_config from Compiler to CompiledModule

* hash isa flags

* no cache for wasm2obj

* mv caching to wasmtime crate

* account each Compiler field when hash

crates/environ/src/cache.rs

@@ -14,6 +14,7 @@ mod worker;
 pub use config::{create_new_config, CacheConfig};
 use worker::Worker;
 
+/// Module level cache entry.
 pub struct ModuleCacheEntry<'config>(Option<ModuleCacheEntryInner<'config>>);
 
 struct ModuleCacheEntryInner<'config> {
@@ -24,6 +25,7 @@ struct ModuleCacheEntryInner<'config> {
 struct Sha256Hasher(Sha256);
 
 impl<'config> ModuleCacheEntry<'config> {
+    /// Create the cache entry.
     pub fn new<'data>(compiler_name: &str, cache_config: &'config CacheConfig) -> Self {
         if cache_config.enabled() {
             Self(Some(ModuleCacheEntryInner::new(
@@ -40,6 +42,7 @@ impl<'config> ModuleCacheEntry<'config> {
         Self(Some(inner))
     }
 
+    /// Gets cached data if state matches, otherwise calls the `compute`.
     pub fn get_data<T, U, E>(&self, state: T, compute: fn(T) -> Result<U, E>) -> Result<U, E>
     where
         T: Hash,

crates/environ/src/compilation.rs

@@ -2,7 +2,6 @@
 //! module.
 
 use crate::address_map::{ModuleAddressMap, ValueLabelsRanges};
-use crate::CacheConfig;
 use crate::ModuleTranslation;
 use cranelift_codegen::{binemit, ir, isa, isa::unwind::UnwindInfo};
 use cranelift_entity::PrimaryMap;
@@ -201,6 +200,5 @@ pub trait Compiler {
     fn compile_module(
         translation: &ModuleTranslation,
         isa: &dyn isa::TargetIsa,
-        cache_config: &CacheConfig,
     ) -> Result<CompileResult, CompileError>;
 }

crates/environ/src/cranelift.rs

@@ -86,14 +86,13 @@
 // assume no valid stack pointer will ever be `usize::max_value() - 32k`.
 
 use crate::address_map::{FunctionAddressMap, InstructionAddressMap};
-use crate::cache::ModuleCacheEntry;
 use crate::compilation::{
     Compilation, CompileError, CompiledFunction, Relocation, RelocationTarget, StackMapInformation,
     TrapInformation,
 };
 use crate::compilation::{CompileResult, Compiler};
 use crate::func_environ::{get_func_name, FuncEnvironment};
-use crate::{CacheConfig, FunctionBodyData, ModuleLocal, ModuleTranslation, Tunables};
+use crate::{FunctionBodyData, ModuleLocal, ModuleTranslation, Tunables};
 use cranelift_codegen::ir::{self, ExternalName};
 use cranelift_codegen::machinst::buffer::MachSrcLoc;
 use cranelift_codegen::print_errors::pretty_error;
@@ -103,7 +102,6 @@ use cranelift_wasm::{DefinedFuncIndex, FuncIndex, FuncTranslator, ModuleTranslat
 #[cfg(feature = "parallel-compilation")]
 use rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
 use std::convert::TryFrom;
-use std::hash::{Hash, Hasher};
 
 /// Implementation of a relocation sink that just saves all the information for later
 pub struct RelocSink {
@@ -290,49 +288,45 @@ impl Compiler for Cranelift {
     fn compile_module(
         translation: &ModuleTranslation,
         isa: &dyn isa::TargetIsa,
-        cache_config: &CacheConfig,
     ) -> Result<CompileResult, CompileError> {
-        let cache_entry = ModuleCacheEntry::new("cranelift", cache_config);
-
-        let result = cache_entry.get_data(
-            CompileEnv {
-                local: &translation.module.local,
-                module_translation: HashedModuleTranslationState(
-                    translation.module_translation.as_ref().unwrap(),
-                ),
-                function_body_inputs: &translation.function_body_inputs,
-                isa: Isa(isa),
-                tunables: &translation.tunables,
-            },
-            compile,
-        )?;
-        Ok(result)
+        compile(
+            isa,
+            &translation.module.local,
+            translation.module_translation.as_ref().unwrap(),
+            &translation.function_body_inputs,
+            &translation.tunables,
+        )
     }
 }
 
-fn compile(env: CompileEnv<'_>) -> Result<CompileResult, CompileError> {
-    let Isa(isa) = env.isa;
-    let mut functions = PrimaryMap::with_capacity(env.function_body_inputs.len());
-    let mut relocations = PrimaryMap::with_capacity(env.function_body_inputs.len());
-    let mut address_transforms = PrimaryMap::with_capacity(env.function_body_inputs.len());
-    let mut value_ranges = PrimaryMap::with_capacity(env.function_body_inputs.len());
-    let mut stack_slots = PrimaryMap::with_capacity(env.function_body_inputs.len());
-    let mut traps = PrimaryMap::with_capacity(env.function_body_inputs.len());
-    let mut stack_maps = PrimaryMap::with_capacity(env.function_body_inputs.len());
+fn compile(
+    isa: &dyn isa::TargetIsa,
+    local: &ModuleLocal,
+    module_translation: &ModuleTranslationState,
+    function_body_inputs: &PrimaryMap<DefinedFuncIndex, FunctionBodyData<'_>>,
+    tunables: &Tunables,
+) -> Result<CompileResult, CompileError> {
+    let mut functions = PrimaryMap::with_capacity(function_body_inputs.len());
+    let mut relocations = PrimaryMap::with_capacity(function_body_inputs.len());
+    let mut address_transforms = PrimaryMap::with_capacity(function_body_inputs.len());
+    let mut value_ranges = PrimaryMap::with_capacity(function_body_inputs.len());
+    let mut stack_slots = PrimaryMap::with_capacity(function_body_inputs.len());
+    let mut traps = PrimaryMap::with_capacity(function_body_inputs.len());
+    let mut stack_maps = PrimaryMap::with_capacity(function_body_inputs.len());
 
     type FunctionBodyInput<'a> = (DefinedFuncIndex, &'a FunctionBodyData<'a>);
 
     let compile_function = |func_translator: &mut FuncTranslator,
                             (i, input): &FunctionBodyInput| {
-        let func_index = env.local.func_index(*i);
+        let func_index = local.func_index(*i);
         let mut context = Context::new();
         context.func.name = get_func_name(func_index);
-        context.func.signature = env.local.native_func_signature(func_index).clone();
-        if env.tunables.debug_info {
+        context.func.signature = local.native_func_signature(func_index).clone();
+        if tunables.debug_info {
             context.func.collect_debug_info();
         }
 
-        let mut func_env = FuncEnvironment::new(isa.frontend_config(), env.local, env.tunables);
+        let mut func_env = FuncEnvironment::new(isa.frontend_config(), local, tunables);
 
         // We use these as constant offsets below in
         // `stack_limit_from_arguments`, so assert their values here. This
@@ -368,7 +362,7 @@ fn compile(env: CompileEnv<'_>) -> Result<CompileResult, CompileError> {
         });
         context.func.stack_limit = Some(stack_limit);
         func_translator.translate(
-            env.module_translation.0,
+            module_translation,
             input.data,
             input.module_offset,
             &mut context.func,
@@ -397,7 +391,7 @@ fn compile(env: CompileEnv<'_>) -> Result<CompileResult, CompileError> {
 
         let address_transform = get_function_address_map(&context, input, code_buf.len(), isa);
 
-        let ranges = if env.tunables.debug_info {
+        let ranges = if tunables.debug_info {
             let ranges = context.build_value_labels_ranges(isa).map_err(|error| {
                 CompileError::Codegen(pretty_error(&context.func, Some(isa), error))
             })?;
@@ -419,7 +413,7 @@ fn compile(env: CompileEnv<'_>) -> Result<CompileResult, CompileError> {
         ))
     };
 
-    let inputs: Vec<FunctionBodyInput> = env.function_body_inputs.into_iter().collect();
+    let inputs: Vec<FunctionBodyInput> = function_body_inputs.into_iter().collect();
 
     let results: Result<Vec<_>, CompileError> = {
         cfg_if::cfg_if! {
@@ -476,49 +470,3 @@ fn compile(env: CompileEnv<'_>) -> Result<CompileResult, CompileError> {
         stack_maps,
     ))
 }
-
-#[derive(Hash)]
-struct CompileEnv<'a> {
-    local: &'a ModuleLocal,
-    module_translation: HashedModuleTranslationState<'a>,
-    function_body_inputs: &'a PrimaryMap<DefinedFuncIndex, FunctionBodyData<'a>>,
-    isa: Isa<'a, 'a>,
-    tunables: &'a Tunables,
-}
-
-/// This is a wrapper struct to hash the specific bits of `TargetIsa` that
-/// affect the output we care about. The trait itself can't implement `Hash`
-/// (it's not object safe) so we have to implement our own hashing.
-struct Isa<'a, 'b>(&'a (dyn isa::TargetIsa + 'b));
-
-impl Hash for Isa<'_, '_> {
-    fn hash<H: Hasher>(&self, hasher: &mut H) {
-        self.0.triple().hash(hasher);
-        self.0.frontend_config().hash(hasher);
-
-        // TODO: if this `to_string()` is too expensive then we should upstream
-        // a native hashing ability of flags into cranelift itself, but
-        // compilation and/or cache loading is relatively expensive so seems
-        // unlikely.
-        self.0.flags().to_string().hash(hasher);
-
-        // TODO: ... and should we hash anything else? There's a lot of stuff in
-        // `TargetIsa`, like registers/encodings/etc. Should we be hashing that
-        // too? It seems like wasmtime doesn't configure it too too much, but
-        // this may become an issue at some point.
-    }
-}
-
-/// A wrapper struct around cranelift's `ModuleTranslationState` to implement
-/// `Hash` since it's not `Hash` upstream yet.
-///
-/// TODO: we should upstream a `Hash` implementation, it would be very small! At
-/// this moment though based on the definition it should be fine to not hash it
-/// since we'll re-hash the signatures later.
-struct HashedModuleTranslationState<'a>(&'a ModuleTranslationState);
-
-impl Hash for HashedModuleTranslationState<'_> {
-    fn hash<H: Hasher>(&self, _hasher: &mut H) {
-        // nothing to hash right now
-    }
-}

crates/environ/src/lib.rs

@@ -44,7 +44,7 @@ pub use crate::address_map::{
     ModuleVmctxInfo, ValueLabelsRanges,
 };
 pub use crate::cache::create_new_config as cache_create_new_config;
-pub use crate::cache::CacheConfig;
+pub use crate::cache::{CacheConfig, ModuleCacheEntry};
 pub use crate::compilation::{
     Compilation, CompileError, CompiledFunction, Compiler, Relocation, RelocationTarget,
     Relocations, StackMapInformation, StackMaps, TrapInformation, Traps,

crates/environ/src/lightbeam.rs

@@ -20,7 +20,6 @@ impl Compiler for Lightbeam {
     fn compile_module(
         translation: &ModuleTranslation,
         isa: &dyn isa::TargetIsa,
-        _cache_config: &CacheConfig,
     ) -> Result<CompileResult, CompileError> {
         if translation.tunables.debug_info {
             return Err(CompileError::DebugInfoNotSupported);

crates/jit/src/compiler.rs

@@ -4,17 +4,18 @@ use crate::instantiate::SetupError;
 use crate::object::{build_object, ObjectUnwindInfo};
 use cranelift_codegen::ir;
 use object::write::Object;
+use std::hash::{Hash, Hasher};
 use wasmtime_debug::{emit_dwarf, DebugInfoData, DwarfSection};
 use wasmtime_environ::entity::{EntityRef, PrimaryMap};
 use wasmtime_environ::isa::{unwind::UnwindInfo, TargetFrontendConfig, TargetIsa};
 use wasmtime_environ::wasm::{DefinedFuncIndex, DefinedMemoryIndex, MemoryIndex};
 use wasmtime_environ::{
-    CacheConfig, Compiler as _C, Module, ModuleAddressMap, ModuleMemoryOffset, ModuleTranslation,
+    Compiler as _C, Module, ModuleAddressMap, ModuleMemoryOffset, ModuleTranslation,
     ModuleVmctxInfo, StackMaps, Traps, Tunables, VMOffsets, ValueLabelsRanges,
 };
 
 /// Select which kind of compilation to use.
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, Hash)]
 pub enum CompilationStrategy {
     /// Let Wasmtime pick the strategy.
     Auto,
@@ -38,22 +39,15 @@ pub enum CompilationStrategy {
 pub struct Compiler {
     isa: Box<dyn TargetIsa>,
     strategy: CompilationStrategy,
-    cache_config: CacheConfig,
     tunables: Tunables,
 }
 
 impl Compiler {
     /// Construct a new `Compiler`.
-    pub fn new(
-        isa: Box<dyn TargetIsa>,
-        strategy: CompilationStrategy,
-        cache_config: CacheConfig,
-        tunables: Tunables,
-    ) -> Self {
+    pub fn new(isa: Box<dyn TargetIsa>, strategy: CompilationStrategy, tunables: Tunables) -> Self {
         Self {
             isa,
             strategy,
-            cache_config,
             tunables,
         }
     }
@@ -126,6 +120,11 @@ impl Compiler {
         &self.tunables
     }
 
+    /// Return the compilation strategy.
+    pub fn strategy(&self) -> CompilationStrategy {
+        self.strategy
+    }
+
     /// Compile the given function bodies.
     pub(crate) fn compile<'data>(
         &self,
@@ -144,19 +143,11 @@ impl Compiler {
             // For now, interpret `Auto` as `Cranelift` since that's the most stable
             // implementation.
             CompilationStrategy::Auto | CompilationStrategy::Cranelift => {
-                wasmtime_environ::cranelift::Cranelift::compile_module(
-                    translation,
-                    &*self.isa,
-                    &self.cache_config,
-                )
+                wasmtime_environ::cranelift::Cranelift::compile_module(translation, &*self.isa)
             }
             #[cfg(feature = "lightbeam")]
             CompilationStrategy::Lightbeam => {
-                wasmtime_environ::lightbeam::Lightbeam::compile_module(
-                    translation,
-                    &*self.isa,
-                    &self.cache_config,
-                )
+                wasmtime_environ::lightbeam::Lightbeam::compile_module(translation, &*self.isa)
             }
         }
         .map_err(SetupError::Compile)?;
@@ -193,3 +184,30 @@ impl Compiler {
         })
     }
 }
+
+impl Hash for Compiler {
+    fn hash<H: Hasher>(&self, hasher: &mut H) {
+        let Compiler {
+            strategy,
+            isa,
+            tunables,
+        } = self;
+
+        // Hash compiler's flags: compilation strategy, isa, frontend config,
+        // misc tunables.
+        strategy.hash(hasher);
+        isa.triple().hash(hasher);
+        // TODO: if this `to_string()` is too expensive then we should upstream
+        // a native hashing ability of flags into cranelift itself, but
+        // compilation and/or cache loading is relatively expensive so seems
+        // unlikely.
+        isa.flags().to_string().hash(hasher);
+        isa.frontend_config().hash(hasher);
+        tunables.hash(hasher);
+
+        // TODO: ... and should we hash anything else? There's a lot of stuff in
+        // `TargetIsa`, like registers/encodings/etc. Should we be hashing that
+        // too? It seems like wasmtime doesn't configure it too too much, but
+        // this may become an issue at some point.
+    }
+}

crates/wasmtime/src/module.rs

@@ -4,6 +4,7 @@ use crate::types::{EntityType, ExportType, ExternType, ImportType};
 use anyhow::{bail, Context, Result};
 use std::path::Path;
 use std::sync::{Arc, Mutex};
+use wasmtime_environ::ModuleCacheEntry;
 use wasmtime_jit::{CompilationArtifacts, CompiledModule};
 
 /// A compiled WebAssembly module, ready to be instantiated.
@@ -300,7 +301,17 @@ impl Module {
     }
 
     unsafe fn compile(engine: &Engine, binary: &[u8]) -> Result<Self> {
-        let compiled = CompiledModule::new(engine.compiler(), binary, &*engine.config().profiler)?;
+        let cache_entry = ModuleCacheEntry::new("wasmtime", engine.cache_config());
+        let artifacts = cache_entry
+            .get_data((engine.compiler(), binary), |(compiler, binary)| {
+                CompilationArtifacts::build(compiler, binary)
+            })?;
+
+        let compiled = CompiledModule::from_artifacts(
+            artifacts,
+            engine.compiler().isa(),
+            &*engine.config().profiler,
+        )?;
 
         Ok(Module {
             engine: engine.clone(),

crates/wasmtime/src/runtime.rs

@@ -628,12 +628,7 @@ impl Config {
 
     fn build_compiler(&self) -> Compiler {
         let isa = self.target_isa();
-        Compiler::new(
-            isa,
-            self.strategy,
-            self.cache_config.clone(),
-            self.tunables.clone(),
-        )
+        Compiler::new(isa, self.strategy, self.tunables.clone())
     }
 
     /// Hashes/fingerprints compiler setting to ensure that compatible
@@ -798,6 +793,10 @@ impl Engine {
         &self.inner.compiler
     }
 
+    pub(crate) fn cache_config(&self) -> &CacheConfig {
+        &self.config().cache_config
+    }
+
     /// Returns whether the engine `a` and `b` refer to the same configuration.
     pub fn same(a: &Engine, b: &Engine) -> bool {
         Arc::ptr_eq(&a.inner, &b.inner)