wasmtime: Option to disable parallel compilation (#3169)

* Introduce parallel-compilation configuration switch

* Plumb parallel_compilation config to compilation

* Adjust obj.rs

* Address review

* Fix compilation fail in `cache` crate

* Fix obj.rs

Also remove the now unneeded feature in /Cargo.toml

* fmt

crates/cache/src/lib.rs

@@ -43,6 +43,8 @@ impl<'config> ModuleCacheEntry<'config> {
     }
 
     /// Gets cached data if state matches, otherwise calls the `compute`.
+    // NOTE: This takes a function pointer instead of a closure so that it doesn't accidentally
+    // close over something not accounted in the cache.
     pub fn get_data<T, U, E>(&self, state: T, compute: fn(T) -> Result<U, E>) -> Result<U, E>
     where
         T: Hash,

crates/jit/src/compiler.rs

@@ -46,6 +46,7 @@ pub struct Compiler {
     strategy: CompilationStrategy,
     tunables: Tunables,
     features: WasmFeatures,
+    parallel_compilation: bool,
 }
 
 impl Compiler {
@@ -55,6 +56,7 @@ impl Compiler {
         strategy: CompilationStrategy,
         tunables: Tunables,
         features: WasmFeatures,
+        parallel_compilation: bool,
     ) -> Self {
         Self {
             isa,
@@ -68,6 +70,7 @@ impl Compiler {
             },
             tunables,
             features,
+            parallel_compilation,
         }
     }
 }
@@ -137,8 +140,9 @@ impl Compiler {
     ) -> Result<Compilation, SetupError> {
         let functions = mem::take(&mut translation.function_body_inputs);
         let functions = functions.into_iter().collect::<Vec<_>>();
-        let funcs = maybe_parallel!(functions.(into_iter | into_par_iter))
-            .map(|(index, func)| {
+
+        let funcs = self
+            .run_maybe_parallel(functions, |(index, func)| {
                 self.compiler.compile_function(
                     translation,
                     index,
@@ -147,8 +151,7 @@ impl Compiler {
                     &self.tunables,
                     types,
                 )
-            })
-            .collect::<Result<Vec<_>, _>>()?
+            })?
             .into_iter()
             .collect::<CompiledFunctions>();
 
@@ -172,6 +175,33 @@ impl Compiler {
             funcs,
         })
     }
+
+    /// Run the given closure in parallel if the compiler is configured to do so.
+    pub(crate) fn run_maybe_parallel<
+        A: Send,
+        B: Send,
+        E: Send,
+        F: Fn(A) -> Result<B, E> + Send + Sync,
+    >(
+        &self,
+        input: Vec<A>,
+        f: F,
+    ) -> Result<Vec<B>, E> {
+        if self.parallel_compilation {
+            #[cfg(feature = "parallel-compilation")]
+            return input
+                .into_par_iter()
+                .map(|a| f(a))
+                .collect::<Result<Vec<B>, E>>();
+        }
+
+        // In case the parallel-compilation feature is disabled or the parallel_compilation config
+        // was turned off dynamically fallback to the non-parallel version.
+        input
+            .into_iter()
+            .map(|a| f(a))
+            .collect::<Result<Vec<B>, E>>()
+    }
 }
 
 impl Hash for Compiler {
@@ -182,6 +212,7 @@ impl Hash for Compiler {
             isa,
             tunables,
             features,
+            parallel_compilation: _,
         } = self;
 
         // Hash compiler's flags: compilation strategy, isa, frontend config,

crates/jit/src/instantiate.rs

@@ -8,8 +8,6 @@ use crate::compiler::{Compilation, Compiler};
 use crate::link::link_module;
 use crate::object::ObjectUnwindInfo;
 use object::File as ObjectFile;
-#[cfg(feature = "parallel-compilation")]
-use rayon::prelude::*;
 use serde::{Deserialize, Serialize};
 use std::ops::Range;
 use std::sync::Arc;
@@ -112,8 +110,9 @@ impl CompilationArtifacts {
         .translate(data)
         .map_err(|error| SetupError::Compile(CompileError::Wasm(error)))?;
 
-        let list = maybe_parallel!(translations.(into_iter | into_par_iter))
-            .map(|mut translation| {
+        let list = compiler.run_maybe_parallel::<_, _, SetupError, _>(
+            translations,
+            |mut translation| {
                 let Compilation {
                     obj,
                     unwind_info,
@@ -159,8 +158,9 @@ impl CompilationArtifacts {
                     },
                     has_unparsed_debuginfo,
                 })
-            })
-            .collect::<Result<Vec<_>, SetupError>>()?;
+            },
+        )?;
+
         Ok((
             main_module,
             list,
@@ -226,10 +226,11 @@ impl CompiledModule {
         artifacts: Vec<CompilationArtifacts>,
         isa: &dyn TargetIsa,
         profiler: &dyn ProfilingAgent,
+        compiler: &Compiler,
     ) -> Result<Vec<Arc<Self>>, SetupError> {
-        maybe_parallel!(artifacts.(into_iter | into_par_iter))
-            .map(|a| CompiledModule::from_artifacts(a, isa, profiler))
-            .collect()
+        compiler.run_maybe_parallel(artifacts, |a| {
+            CompiledModule::from_artifacts(a, isa, profiler)
+        })
     }
 
     /// Creates `CompiledModule` directly from `CompilationArtifacts`.

crates/jit/src/lib.rs

@@ -20,20 +20,6 @@
     )
 )]
 
-#[cfg(feature = "parallel-compilation")]
-macro_rules! maybe_parallel {
-    ($e:ident.($serial:ident | $parallel:ident)) => {
-        $e.$parallel()
-    };
-}
-
-#[cfg(not(feature = "parallel-compilation"))]
-macro_rules! maybe_parallel {
-    ($e:ident.($serial:ident | $parallel:ident)) => {
-        $e.$serial()
-    };
-}
-
 mod code_memory;
 mod compiler;
 mod instantiate;

crates/wasmtime/src/config.rs

@@ -356,6 +356,7 @@ pub struct Config {
     pub(crate) async_stack_size: usize,
     pub(crate) async_support: bool,
     pub(crate) deserialize_check_wasmtime_version: bool,
+    pub(crate) parallel_compilation: bool,
 }
 
 impl Config {
@@ -392,6 +393,7 @@ impl Config {
             async_stack_size: 2 << 20,
             async_support: false,
             deserialize_check_wasmtime_version: true,
+            parallel_compilation: true,
         };
         ret.cranelift_debug_verifier(false);
         ret.cranelift_opt_level(OptLevel::Speed);
@@ -1210,6 +1212,18 @@ impl Config {
         self
     }
 
+    /// Configure wether wasmtime should compile a module using multiple threads.
+    ///
+    /// Disabling this will result in a single thread being used to compile the wasm bytecode.
+    ///
+    /// By default parallel compilation is enabled.
+    #[cfg(feature = "parallel-compilation")]
+    #[cfg_attr(nightlydoc, doc(cfg(feature = "parallel-compilation")))]
+    pub fn parallel_compilation(&mut self, parallel: bool) -> &mut Self {
+        self.parallel_compilation = parallel;
+        self
+    }
+
     pub(crate) fn target_isa(&self) -> Box<dyn TargetIsa> {
         self.isa_flags
             .clone()
@@ -1226,7 +1240,13 @@ impl Config {
         let isa = self.target_isa();
         let mut tunables = self.tunables.clone();
         allocator.adjust_tunables(&mut tunables);
-        Compiler::new(isa, self.strategy, tunables, self.features)
+        Compiler::new(
+            isa,
+            self.strategy,
+            tunables,
+            self.features,
+            self.parallel_compilation,
+        )
     }
 
     pub(crate) fn build_allocator(&self) -> Result<Box<dyn InstanceAllocator>> {

crates/wasmtime/src/module.rs

@@ -300,18 +300,27 @@ impl Module {
                     engine.cache_config(),
                 )
                 .get_data((engine.compiler(), binary), |(compiler, binary)| {
-                    CompilationArtifacts::build(compiler, binary, USE_PAGED_MEM_INIT)
+                    CompilationArtifacts::build(
+                        compiler,
+                        binary,
+                        USE_PAGED_MEM_INIT,
+                    )
                 })?;
             } else {
                 let (main_module, artifacts, types) =
-                    CompilationArtifacts::build(engine.compiler(), binary, USE_PAGED_MEM_INIT)?;
+                    CompilationArtifacts::build(
+                        engine.compiler(),
+                        binary,
+                        USE_PAGED_MEM_INIT,
+                    )?;
             }
         };
 
         let modules = CompiledModule::from_artifacts_list(
             artifacts,
             engine.compiler().isa(),
             &*engine.config().profiler,
+            engine.compiler(),
         )?;
 
         Self::from_parts(engine, modules, main_module, Arc::new(types), &[])

crates/wasmtime/src/module/serialization.rs

@@ -291,6 +291,7 @@ impl<'a> SerializedModule<'a> {
                 .collect(),
             engine.compiler().isa(),
             &*engine.config().profiler,
+            engine.compiler(),
         )?;
 
         assert!(!modules.is_empty());

src/obj.rs

@@ -62,6 +62,7 @@ pub fn compile_to_obj(
         },
         tunables.clone(),
         features.clone(),
+        true, // enable parallel compilation
     );
 
     let environ = ModuleEnvironment::new(compiler.isa().frontend_config(), &tunables, &features);