Skip to content

Commit 7d8263c

Browse files
bjorn3bjorn3
committed
Move compilation to worker threads
1 parent 13c0092 commit 7d8263c

File tree

3 files changed

+107
-56
lines changed

3 files changed

+107
-56
lines changed

‎src/driver/aot.rs‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -516,7 +516,7 @@ fn codegen_cgu_content(
516516
tcx: TyCtxt<'_>,
517517
module: &mut dyn Module,
518518
cgu_name: rustc_span::Symbol,
519-
) -> (Vec<SerializableModule>, String) {
519+
) -> (Vec<(SerializableModule, Option<Fingerprint>)>, String) {
520520
let _timer = tcx.prof.generic_activity_with_arg("codegen cgu", cgu_name.as_str());
521521

522522
let cgu = tcx.codegen_unit(cgu_name);
@@ -560,15 +560,16 @@ fn codegen_cgu_content(
560560

561561
if tcx.dep_graph.is_fully_enabled() && tcx.try_mark_green(&dep_node) {
562562
let data = FileCache.get(&cache_key.to_le_bytes()).unwrap();
563-
codegened_functions.push(SerializableModule::deserialize(&data, isa.clone()));
563+
codegened_functions
564+
.push((SerializableModule::deserialize(&data, isa.clone()), None));
564565
continue;
565566
};
566567

567568
let (ser_module, _) = tcx.dep_graph.with_task(
568569
dep_node,
569570
tcx,
570-
(instance, cache_key),
571-
|tcx, (instance, cache_key)| {
571+
instance,
572+
|tcx, instance| {
572573
let mut ser_module =
573574
SerializableModule::new(crate::build_isa(tcx.sess, false));
574575
let codegened_function = crate::base::codegen_fn(
@@ -594,15 +595,12 @@ fn codegen_cgu_content(
594595
);
595596
ser_module.add_global_asm(&global_asm);
596597

597-
let data = ser_module.serialize();
598-
FileCache.insert(&cache_key.to_le_bytes(), data);
599-
600598
ser_module
601599
},
602600
Some(rustc_middle::dep_graph::hash_result),
603601
);
604602

605-
codegened_functions.push(ser_module);
603+
codegened_functions.push((ser_module, Some(cache_key)));
606604
}
607605
MonoItem::Static(def_id) => {
608606
let data_id = crate::constant::codegen_static(tcx, module, def_id);
@@ -648,7 +646,12 @@ fn module_codegen(
648646
profiler.clone(),
649647
)));
650648

651-
for codegened_func in codegened_functions {
649+
for (codegened_func, cache_key) in codegened_functions {
650+
if let Some(cache_key) = cache_key {
651+
let data = codegened_func.serialize();
652+
FileCache.insert(&cache_key.to_le_bytes(), data);
653+
}
654+
652655
let asm = codegened_func.apply_to(&mut module);
653656
global_asm.push_str(&asm);
654657
}

‎src/serializable_module.rs‎

Lines changed: 90 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
1+
use std::cell::RefCell;
12
use std::collections::BTreeMap;
3+
use std::mem;
24
use std::sync::{Arc, OnceLock};
35

46
use cranelift_codegen::control::ControlPlane;
57
use cranelift_codegen::entity::SecondaryMap;
6-
use cranelift_codegen::ir::{Signature, UserExternalName};
8+
use cranelift_codegen::ir::function::FunctionParameters;
9+
use cranelift_codegen::ir::{ExternalName, Signature, UserExternalName};
710
use cranelift_codegen::isa::TargetIsa;
11+
use cranelift_codegen::{Final, FinalizedMachReloc, FinalizedRelocTarget, MachBufferFinalized};
812
use cranelift_module::{
913
DataId, ModuleDeclarations, ModuleError, ModuleReloc, ModuleRelocTarget, ModuleResult,
1014
};
@@ -18,14 +22,20 @@ pub(super) struct SerializableModule {
1822
serialized: OnceLock<Vec<u8>>,
1923
}
2024

21-
#[derive(Debug, serde::Serialize, serde::Deserialize)]
25+
#[derive(serde::Serialize, serde::Deserialize)]
2226
struct SerializableModuleInner {
2327
declarations: ModuleDeclarations,
24-
functions: BTreeMap<FuncId, Function>,
28+
functions: BTreeMap<FuncId, RefCell<FunctionMaybeCompiled>>,
2529
data_objects: BTreeMap<DataId, DataDescription>,
2630
global_asm: String,
2731
}
2832

33+
#[derive(serde::Serialize, serde::Deserialize)]
34+
enum FunctionMaybeCompiled {
35+
Ir(Function),
36+
Compiled(MachBufferFinalized<Final>, FunctionParameters),
37+
}
38+
2939
impl<CTX> HashStable<CTX> for SerializableModule {
3040
fn hash_stable(
3141
&self,
@@ -52,6 +62,7 @@ impl SerializableModule {
5262
}
5363

5464
pub(crate) fn serialize(&self) -> Vec<u8> {
65+
self.compile_funcs();
5566
postcard::to_stdvec(&self.inner).unwrap()
5667
}
5768

@@ -68,7 +79,26 @@ impl SerializableModule {
6879
self.inner.global_asm.push_str(asm);
6980
}
7081

82+
fn compile_funcs(&self) {
83+
let mut ctx = Context::new();
84+
for func in self.inner.functions.values() {
85+
let mut func = func.borrow_mut();
86+
match &mut *func {
87+
FunctionMaybeCompiled::Ir(ir_func) => {
88+
// FIXME lazily do this during serialize/apply_to
89+
ctx.func = mem::replace(ir_func, Function::new());
90+
let res = ctx.compile(&*self.isa, &mut ControlPlane::default()).unwrap();
91+
92+
let buffer = res.buffer.clone();
93+
*func = FunctionMaybeCompiled::Compiled(buffer, ctx.func.params);
94+
}
95+
FunctionMaybeCompiled::Compiled(_, _) => {}
96+
}
97+
}
98+
}
99+
71100
pub(crate) fn apply_to(self, module: &mut dyn Module) -> String {
101+
self.compile_funcs();
72102
let mut function_map: SecondaryMap<FuncId, Option<FuncId>> = SecondaryMap::new();
73103
let mut data_object_map: SecondaryMap<DataId, Option<DataId>> = SecondaryMap::new();
74104

@@ -98,41 +128,62 @@ impl SerializableModule {
98128
data_object_map[data_id].unwrap()
99129
};
100130

101-
for (func_id, mut func) in self.inner.functions {
131+
for (func_id, func) in self.inner.functions {
102132
let func_id = remap_func_id(module, &self.inner.declarations, func_id);
103-
let user_named_funcs = func.params.user_named_funcs().clone();
104-
for (ext_name_ref, ext_name) in user_named_funcs {
105-
if ext_name.namespace == 0 {
106-
func.params.reset_user_func_name(
107-
ext_name_ref,
108-
UserExternalName::new(
109-
0,
110-
remap_func_id(
111-
module,
112-
&self.inner.declarations,
113-
FuncId::from_u32(ext_name.index),
133+
134+
let FunctionMaybeCompiled::Compiled(buffer, params) = &*func.borrow() else {
135+
unreachable!()
136+
};
137+
138+
let remap_reloc = |reloc: &FinalizedMachReloc| {
139+
let name = match reloc.target {
140+
FinalizedRelocTarget::ExternalName(ExternalName::User(reff)) => {
141+
let ext_name = &params.user_named_funcs()[reff];
142+
let ext_name = if ext_name.namespace == 0 {
143+
UserExternalName::new(
144+
0,
145+
remap_func_id(
146+
module,
147+
&self.inner.declarations,
148+
FuncId::from_u32(ext_name.index),
149+
)
150+
.as_u32(),
114151
)
115-
.as_u32(),
116-
),
117-
);
118-
} else if ext_name.namespace == 1 {
119-
func.params.reset_user_func_name(
120-
ext_name_ref,
121-
UserExternalName::new(
122-
1,
123-
remap_data_id(
124-
module,
125-
&self.inner.declarations,
126-
DataId::from_u32(ext_name.index),
152+
} else if ext_name.namespace == 1 {
153+
UserExternalName::new(
154+
1,
155+
remap_data_id(
156+
module,
157+
&self.inner.declarations,
158+
DataId::from_u32(ext_name.index),
159+
)
160+
.as_u32(),
127161
)
128-
.as_u32(),
129-
),
130-
);
131-
} else {
132-
unreachable!();
133-
}
134-
}
135-
module.define_function(func_id, &mut Context::for_function(func)).unwrap();
162+
} else {
163+
unreachable!();
164+
};
165+
ModuleRelocTarget::user(ext_name.namespace, ext_name.index)
166+
}
167+
FinalizedRelocTarget::ExternalName(ExternalName::TestCase(_)) => {
168+
unimplemented!()
169+
}
170+
FinalizedRelocTarget::ExternalName(ExternalName::LibCall(libcall)) => {
171+
ModuleRelocTarget::LibCall(libcall)
172+
}
173+
FinalizedRelocTarget::ExternalName(ExternalName::KnownSymbol(ks)) => {
174+
ModuleRelocTarget::KnownSymbol(ks)
175+
}
176+
FinalizedRelocTarget::Func(offset) => {
177+
ModuleRelocTarget::FunctionOffset(func_id, offset)
178+
}
179+
};
180+
ModuleReloc { offset: reloc.offset, kind: reloc.kind, name, addend: reloc.addend }
181+
};
182+
183+
let relocs = buffer.relocs().iter().map(remap_reloc).collect::<Vec<_>>();
184+
module
185+
.define_function_bytes(func_id, buffer.alignment as u64, buffer.data(), &relocs)
186+
.unwrap();
136187
}
137188

138189
for (data_id, mut data) in self.inner.data_objects {
@@ -219,7 +270,7 @@ impl Module for SerializableModule {
219270
&mut self,
220271
func_id: FuncId,
221272
ctx: &mut Context,
222-
ctrl_plane: &mut ControlPlane,
273+
_ctrl_plane: &mut ControlPlane,
223274
) -> ModuleResult<()> {
224275
let decl = self.inner.declarations.get_function_decl(func_id);
225276
if !decl.linkage.is_definable() {
@@ -234,12 +285,9 @@ impl Module for SerializableModule {
234285
));
235286
}
236287

237-
ctx.verify_if(&*self.isa)?;
238-
ctx.optimize(&*self.isa, ctrl_plane)?;
239-
240-
// FIXME compile to machine code
241-
242-
self.inner.functions.insert(func_id, ctx.func.clone());
288+
self.inner
289+
.functions
290+
.insert(func_id, RefCell::new(FunctionMaybeCompiled::Ir(ctx.func.clone())));
243291

244292
Ok(())
245293
}

‎src/unwind_module.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -136,12 +136,12 @@ impl<T: Module> Module for UnwindModule<T> {
136136

137137
fn define_function_bytes(
138138
&mut self,
139-
_func_id: FuncId,
140-
_alignment: u64,
141-
_bytes: &[u8],
142-
_relocs: &[ModuleReloc],
139+
func_id: FuncId,
140+
alignment: u64,
141+
bytes: &[u8],
142+
relocs: &[ModuleReloc],
143143
) -> ModuleResult<()> {
144-
unimplemented!()
144+
self.module.define_function_bytes(func_id, alignment, bytes, relocs)
145145
}
146146

147147
fn define_data(&mut self, data_id: DataId, data: &DataDescription) -> ModuleResult<()> {

0 commit comments

Comments
 (0)