Move compilation to worker threads

Author: bjorn3

src/driver/aot.rs

@@ -245,7 +245,7 @@ fn codegen_cgu_content(
     tcx: TyCtxt<'_>,
     module: &mut dyn Module,
     cgu_name: rustc_span::Symbol,
-) -> (Vec<SerializableModule>, String) {
+) -> (Vec<(SerializableModule, Option<Fingerprint>)>, String) {
     let _timer = tcx.prof.generic_activity_with_arg("codegen cgu", cgu_name.as_str());
 
     let cgu = tcx.codegen_unit(cgu_name);
@@ -290,7 +290,8 @@ fn codegen_cgu_content(
                     && tcx.dep_graph.try_mark_green(tcx, &dep_node).is_some()
                 {
                     let data = FileCache.get(&cache_key.to_le_bytes()).unwrap();
-                    codegened_functions.push(SerializableModule::deserialize(&data, isa.clone()));
+                    codegened_functions
+                        .push((SerializableModule::deserialize(&data, isa.clone()), None));
                     continue;
                 };
 
@@ -320,15 +321,12 @@ fn codegen_cgu_content(
                         );
                         ser_module.add_global_asm(&global_asm);
 
-                        let data = ser_module.serialize();
-                        FileCache.insert(&cache_key.to_le_bytes(), data.to_vec());
-
                         ser_module
                     },
                     Some(rustc_middle::dep_graph::hash_result),
                 );
 
-                codegened_functions.push(ser_module);
+                codegened_functions.push((ser_module, Some(cache_key)));
             }
             MonoItem::Static(def_id) => {
                 crate::constant::codegen_static(tcx, module, def_id);
@@ -384,7 +382,7 @@ fn compile_cgu(
     producer: String,
     global_asm_config: GlobalAsmConfig,
     mut module: UnwindModule<ObjectModule>,
-    codegened_functions: Vec<SerializableModule>,
+    codegened_functions: Vec<(SerializableModule, Option<Fingerprint>)>,
     mut global_asm: String,
     cgu_name: String,
 ) -> Result<CompiledModule, String> {
@@ -393,7 +391,12 @@ fn compile_cgu(
             prof.clone(),
         )));
 
-        for codegened_func in codegened_functions {
+        for (codegened_func, cache_key) in codegened_functions {
+            if let Some(cache_key) = cache_key {
+                let data = codegened_func.serialize();
+                FileCache.insert(&cache_key.to_le_bytes(), data.to_vec());
+            }
+
             let asm = codegened_func.apply_to(&mut module);
             global_asm.push_str(&asm);
         }

src/serializable_module.rs

@@ -1,10 +1,14 @@
+use std::cell::RefCell;
 use std::collections::BTreeMap;
+use std::mem;
 use std::sync::{Arc, OnceLock};
 
 use cranelift_codegen::control::ControlPlane;
 use cranelift_codegen::entity::SecondaryMap;
-use cranelift_codegen::ir::{Signature, UserExternalName};
+use cranelift_codegen::ir::function::FunctionParameters;
+use cranelift_codegen::ir::{ExternalName, Signature, UserExternalName};
 use cranelift_codegen::isa::TargetIsa;
+use cranelift_codegen::{Final, FinalizedMachReloc, FinalizedRelocTarget, MachBufferFinalized};
 use cranelift_module::{
     DataId, ModuleDeclarations, ModuleError, ModuleReloc, ModuleRelocTarget, ModuleResult,
 };
@@ -18,14 +22,20 @@ pub(super) struct SerializableModule {
     serialized: OnceLock<Vec<u8>>,
 }
 
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
+#[derive(serde::Serialize, serde::Deserialize)]
 struct SerializableModuleInner {
     declarations: ModuleDeclarations,
-    functions: BTreeMap<FuncId, Function>,
+    functions: BTreeMap<FuncId, RefCell<FunctionMaybeCompiled>>,
     data_objects: BTreeMap<DataId, DataDescription>,
     global_asm: String,
 }
 
+#[derive(serde::Serialize, serde::Deserialize)]
+enum FunctionMaybeCompiled {
+    Ir(Function),
+    Compiled(MachBufferFinalized<Final>, FunctionParameters),
+}
+
 impl StableHash for SerializableModule {
     fn stable_hash<Hcx: StableHashCtxt>(&self, hcx: &mut Hcx, hasher: &mut StableHasher) {
         self.serialize().stable_hash(hcx, hasher);
@@ -47,6 +57,7 @@ impl SerializableModule {
     }
 
     pub(crate) fn serialize(&self) -> &[u8] {
+        self.compile_funcs();
         self.serialized.get_or_init(|| postcard::to_stdvec(&self.inner).unwrap())
     }
 
@@ -63,7 +74,26 @@ impl SerializableModule {
         self.inner.global_asm.push_str(asm);
     }
 
+    fn compile_funcs(&self) {
+        let mut ctx = Context::new();
+        for func in self.inner.functions.values() {
+            let mut func = func.borrow_mut();
+            match &mut *func {
+                FunctionMaybeCompiled::Ir(ir_func) => {
+                    // FIXME lazily do this during serialize/apply_to
+                    ctx.func = mem::replace(ir_func, Function::new());
+                    let res = ctx.compile(&*self.isa, &mut ControlPlane::default()).unwrap();
+
+                    let buffer = res.buffer.clone();
+                    *func = FunctionMaybeCompiled::Compiled(buffer, ctx.func.params);
+                }
+                FunctionMaybeCompiled::Compiled(_, _) => {}
+            }
+        }
+    }
+
     pub(crate) fn apply_to(self, module: &mut dyn Module) -> String {
+        self.compile_funcs();
         let mut function_map: SecondaryMap<FuncId, Option<FuncId>> = SecondaryMap::new();
         let mut data_object_map: SecondaryMap<DataId, Option<DataId>> = SecondaryMap::new();
 
@@ -93,41 +123,62 @@ impl SerializableModule {
                 data_object_map[data_id].unwrap()
             };
 
-        for (func_id, mut func) in self.inner.functions {
+        for (func_id, func) in self.inner.functions {
             let func_id = remap_func_id(module, &self.inner.declarations, func_id);
-            let user_named_funcs = func.params.user_named_funcs().clone();
-            for (ext_name_ref, ext_name) in user_named_funcs {
-                if ext_name.namespace == 0 {
-                    func.params.reset_user_func_name(
-                        ext_name_ref,
-                        UserExternalName::new(
-                            0,
-                            remap_func_id(
-                                module,
-                                &self.inner.declarations,
-                                FuncId::from_u32(ext_name.index),
+
+            let FunctionMaybeCompiled::Compiled(buffer, params) = &*func.borrow() else {
+                unreachable!()
+            };
+
+            let remap_reloc = |reloc: &FinalizedMachReloc| {
+                let name = match reloc.target {
+                    FinalizedRelocTarget::ExternalName(ExternalName::User(reff)) => {
+                        let ext_name = &params.user_named_funcs()[reff];
+                        let ext_name = if ext_name.namespace == 0 {
+                            UserExternalName::new(
+                                0,
+                                remap_func_id(
+                                    module,
+                                    &self.inner.declarations,
+                                    FuncId::from_u32(ext_name.index),
+                                )
+                                .as_u32(),
                             )
-                            .as_u32(),
-                        ),
-                    );
-                } else if ext_name.namespace == 1 {
-                    func.params.reset_user_func_name(
-                        ext_name_ref,
-                        UserExternalName::new(
-                            1,
-                            remap_data_id(
-                                module,
-                                &self.inner.declarations,
-                                DataId::from_u32(ext_name.index),
+                        } else if ext_name.namespace == 1 {
+                            UserExternalName::new(
+                                1,
+                                remap_data_id(
+                                    module,
+                                    &self.inner.declarations,
+                                    DataId::from_u32(ext_name.index),
+                                )
+                                .as_u32(),
                             )
-                            .as_u32(),
-                        ),
-                    );
-                } else {
-                    unreachable!();
-                }
-            }
-            module.define_function(func_id, &mut Context::for_function(func)).unwrap();
+                        } else {
+                            unreachable!();
+                        };
+                        ModuleRelocTarget::user(ext_name.namespace, ext_name.index)
+                    }
+                    FinalizedRelocTarget::ExternalName(ExternalName::TestCase(_)) => {
+                        unimplemented!()
+                    }
+                    FinalizedRelocTarget::ExternalName(ExternalName::LibCall(libcall)) => {
+                        ModuleRelocTarget::LibCall(libcall)
+                    }
+                    FinalizedRelocTarget::ExternalName(ExternalName::KnownSymbol(ks)) => {
+                        ModuleRelocTarget::KnownSymbol(ks)
+                    }
+                    FinalizedRelocTarget::Func(offset) => {
+                        ModuleRelocTarget::FunctionOffset(func_id, offset)
+                    }
+                };
+                ModuleReloc { offset: reloc.offset, kind: reloc.kind, name, addend: reloc.addend }
+            };
+
+            let relocs = buffer.relocs().iter().map(remap_reloc).collect::<Vec<_>>();
+            module
+                .define_function_bytes(func_id, buffer.alignment as u64, buffer.data(), &relocs)
+                .unwrap();
         }
 
         for (data_id, mut data) in self.inner.data_objects {
@@ -214,7 +265,7 @@ impl Module for SerializableModule {
         &mut self,
         func_id: FuncId,
         ctx: &mut Context,
-        ctrl_plane: &mut ControlPlane,
+        _ctrl_plane: &mut ControlPlane,
     ) -> ModuleResult<()> {
         let decl = self.inner.declarations.get_function_decl(func_id);
         if !decl.linkage.is_definable() {
@@ -229,12 +280,9 @@ impl Module for SerializableModule {
             ));
         }
 
-        ctx.verify_if(&*self.isa)?;
-        ctx.optimize(&*self.isa, ctrl_plane)?;
-
-        // FIXME compile to machine code
-
-        self.inner.functions.insert(func_id, ctx.func.clone());
+        self.inner
+            .functions
+            .insert(func_id, RefCell::new(FunctionMaybeCompiled::Ir(ctx.func.clone())));
 
         Ok(())
     }

src/unwind_module.rs

@@ -136,12 +136,12 @@ impl<T: Module> Module for UnwindModule<T> {
 
     fn define_function_bytes(
         &mut self,
-        _func_id: FuncId,
-        _alignment: u64,
-        _bytes: &[u8],
-        _relocs: &[ModuleReloc],
+        func_id: FuncId,
+        alignment: u64,
+        bytes: &[u8],
+        relocs: &[ModuleReloc],
     ) -> ModuleResult<()> {
-        unimplemented!()
+        self.module.define_function_bytes(func_id, alignment, bytes, relocs)
     }
 
     fn define_data(&mut self, data_id: DataId, data: &DataDescription) -> ModuleResult<()> {