Skip to content

Commit be5f8f2

Browse files
fmeasefmease
authored
Rollup merge of #143388 - bjorn3:lto_refactors, r=compiler-errors
Various refactors to the LTO handling code In particular reducing the sharing of code paths between fat and thin-LTO and making the fat LTO implementation more self-contained. This also moves some autodiff handling out of cg_ssa into cg_llvm given that Enzyme only works with LLVM anyway and an implementation for another backend may do things entirely differently. This will also make it a bit easier to split LTO handling out of the coordinator thread main loop into a separate loop, which should reduce the complexity of the coordinator thread.
2 parents 9ac88ea + 21026ca commit be5f8f2

File tree

13 files changed

+155
-241
lines changed

13 files changed

+155
-241
lines changed

‎compiler/rustc_codegen_gcc/src/back/lto.rs‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ use std::sync::Arc;
2424

2525
use gccjit::{Context, OutputKind};
2626
use object::read::archive::ArchiveFile;
27-
use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule, ThinShared};
27+
use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule, ThinShared};
2828
use rustc_codegen_ssa::back::symbol_export;
2929
use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
3030
use rustc_codegen_ssa::traits::*;
@@ -176,7 +176,7 @@ pub(crate) fn run_fat(
176176
cgcx: &CodegenContext<GccCodegenBackend>,
177177
modules: Vec<FatLtoInput<GccCodegenBackend>>,
178178
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
179-
) -> Result<LtoModuleCodegen<GccCodegenBackend>, FatalError> {
179+
) -> Result<ModuleCodegen<GccContext>, FatalError> {
180180
let dcx = cgcx.create_dcx();
181181
let dcx = dcx.handle();
182182
let lto_data = prepare_lto(cgcx, dcx)?;
@@ -201,7 +201,7 @@ fn fat_lto(
201201
mut serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
202202
tmp_path: TempDir,
203203
//symbols_below_threshold: &[String],
204-
) -> Result<LtoModuleCodegen<GccCodegenBackend>, FatalError> {
204+
) -> Result<ModuleCodegen<GccContext>, FatalError> {
205205
let _timer = cgcx.prof.generic_activity("GCC_fat_lto_build_monolithic_module");
206206
info!("going for a fat lto");
207207

@@ -334,7 +334,7 @@ fn fat_lto(
334334
// of now.
335335
module.module_llvm.temp_dir = Some(tmp_path);
336336

337-
Ok(LtoModuleCodegen::Fat(module))
337+
Ok(module)
338338
}
339339

340340
pub struct ModuleBuffer(PathBuf);
@@ -358,7 +358,7 @@ pub(crate) fn run_thin(
358358
cgcx: &CodegenContext<GccCodegenBackend>,
359359
modules: Vec<(String, ThinBuffer)>,
360360
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
361-
) -> Result<(Vec<LtoModuleCodegen<GccCodegenBackend>>, Vec<WorkProduct>), FatalError> {
361+
) -> Result<(Vec<ThinModule<GccCodegenBackend>>, Vec<WorkProduct>), FatalError> {
362362
let dcx = cgcx.create_dcx();
363363
let dcx = dcx.handle();
364364
let lto_data = prepare_lto(cgcx, dcx)?;
@@ -427,7 +427,7 @@ fn thin_lto(
427427
tmp_path: TempDir,
428428
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
429429
//_symbols_below_threshold: &[String],
430-
) -> Result<(Vec<LtoModuleCodegen<GccCodegenBackend>>, Vec<WorkProduct>), FatalError> {
430+
) -> Result<(Vec<ThinModule<GccCodegenBackend>>, Vec<WorkProduct>), FatalError> {
431431
let _timer = cgcx.prof.generic_activity("LLVM_thin_lto_global_analysis");
432432
info!("going for that thin, thin LTO");
433433

@@ -573,8 +573,7 @@ fn thin_lto(
573573
}*/
574574

575575
info!(" - {}: re-compiled", module_name);
576-
opt_jobs
577-
.push(LtoModuleCodegen::Thin(ThinModule { shared: shared.clone(), idx: module_index }));
576+
opt_jobs.push(ThinModule { shared: shared.clone(), idx: module_index });
578577
}
579578

580579
// Save the current ThinLTO import information for the next compilation

‎compiler/rustc_codegen_gcc/src/back/write.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,12 @@ use crate::{GccCodegenBackend, GccContext};
1616

1717
pub(crate) fn codegen(
1818
cgcx: &CodegenContext<GccCodegenBackend>,
19-
dcx: DiagCtxtHandle<'_>,
2019
module: ModuleCodegen<GccContext>,
2120
config: &ModuleConfig,
2221
) -> Result<CompiledModule, FatalError> {
22+
let dcx = cgcx.create_dcx();
23+
let dcx = dcx.handle();
24+
2325
let _timer = cgcx.prof.generic_activity_with_arg("GCC_module_codegen", &*module.name);
2426
{
2527
let context = &module.module_llvm.context;

‎compiler/rustc_codegen_gcc/src/lib.rs‎

Lines changed: 10 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ use gccjit::{CType, Context, OptimizationLevel};
9393
use gccjit::{TargetInfo, Version};
9494
use rustc_ast::expand::allocator::AllocatorKind;
9595
use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
96-
use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};
96+
use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule};
9797
use rustc_codegen_ssa::back::write::{
9898
CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryFn,
9999
};
@@ -353,19 +353,24 @@ impl WriteBackendMethods for GccCodegenBackend {
353353
type ThinData = ThinData;
354354
type ThinBuffer = ThinBuffer;
355355

356-
fn run_fat_lto(
356+
fn run_and_optimize_fat_lto(
357357
cgcx: &CodegenContext<Self>,
358358
modules: Vec<FatLtoInput<Self>>,
359359
cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
360-
) -> Result<LtoModuleCodegen<Self>, FatalError> {
360+
diff_fncs: Vec<AutoDiffItem>,
361+
) -> Result<ModuleCodegen<Self::Module>, FatalError> {
362+
if !diff_fncs.is_empty() {
363+
unimplemented!();
364+
}
365+
361366
back::lto::run_fat(cgcx, modules, cached_modules)
362367
}
363368

364369
fn run_thin_lto(
365370
cgcx: &CodegenContext<Self>,
366371
modules: Vec<(String, Self::ThinBuffer)>,
367372
cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
368-
) -> Result<(Vec<LtoModuleCodegen<Self>>, Vec<WorkProduct>), FatalError> {
373+
) -> Result<(Vec<ThinModule<Self>>, Vec<WorkProduct>), FatalError> {
369374
back::lto::run_thin(cgcx, modules, cached_modules)
370375
}
371376

@@ -387,14 +392,6 @@ impl WriteBackendMethods for GccCodegenBackend {
387392
Ok(())
388393
}
389394

390-
fn optimize_fat(
391-
_cgcx: &CodegenContext<Self>,
392-
_module: &mut ModuleCodegen<Self::Module>,
393-
) -> Result<(), FatalError> {
394-
// TODO(antoyo)
395-
Ok(())
396-
}
397-
398395
fn optimize_thin(
399396
cgcx: &CodegenContext<Self>,
400397
thin: ThinModule<Self>,
@@ -404,11 +401,10 @@ impl WriteBackendMethods for GccCodegenBackend {
404401

405402
fn codegen(
406403
cgcx: &CodegenContext<Self>,
407-
dcx: DiagCtxtHandle<'_>,
408404
module: ModuleCodegen<Self::Module>,
409405
config: &ModuleConfig,
410406
) -> Result<CompiledModule, FatalError> {
411-
back::write::codegen(cgcx, dcx, module, config)
407+
back::write::codegen(cgcx, module, config)
412408
}
413409

414410
fn prepare_thin(
@@ -429,15 +425,6 @@ impl WriteBackendMethods for GccCodegenBackend {
429425
) -> Result<ModuleCodegen<Self::Module>, FatalError> {
430426
back::write::link(cgcx, dcx, modules)
431427
}
432-
433-
fn autodiff(
434-
_cgcx: &CodegenContext<Self>,
435-
_module: &ModuleCodegen<Self::Module>,
436-
_diff_functions: Vec<AutoDiffItem>,
437-
_config: &ModuleConfig,
438-
) -> Result<(), FatalError> {
439-
unimplemented!()
440-
}
441428
}
442429

443430
/// This is the entrypoint for a hot plugged rustc_codegen_gccjit

‎compiler/rustc_codegen_llvm/messages.ftl‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
codegen_llvm_autodiff_without_enable = using the autodiff feature requires -Z autodiff=Enable
2-
codegen_llvm_autodiff_without_lto = using the autodiff feature requires using fat-lto
32
43
codegen_llvm_copy_bitcode = failed to copy bitcode to object file: {$err}
54

‎compiler/rustc_codegen_llvm/src/back/lto.rs‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use std::sync::Arc;
77
use std::{io, iter, slice};
88

99
use object::read::archive::ArchiveFile;
10-
use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule, ThinShared};
10+
use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule, ThinShared};
1111
use rustc_codegen_ssa::back::symbol_export;
1212
use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
1313
use rustc_codegen_ssa::traits::*;
@@ -201,7 +201,7 @@ pub(crate) fn run_fat(
201201
cgcx: &CodegenContext<LlvmCodegenBackend>,
202202
modules: Vec<FatLtoInput<LlvmCodegenBackend>>,
203203
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
204-
) -> Result<LtoModuleCodegen<LlvmCodegenBackend>, FatalError> {
204+
) -> Result<ModuleCodegen<ModuleLlvm>, FatalError> {
205205
let dcx = cgcx.create_dcx();
206206
let dcx = dcx.handle();
207207
let (symbols_below_threshold, upstream_modules) = prepare_lto(cgcx, dcx)?;
@@ -217,7 +217,7 @@ pub(crate) fn run_thin(
217217
cgcx: &CodegenContext<LlvmCodegenBackend>,
218218
modules: Vec<(String, ThinBuffer)>,
219219
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
220-
) -> Result<(Vec<LtoModuleCodegen<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
220+
) -> Result<(Vec<ThinModule<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
221221
let dcx = cgcx.create_dcx();
222222
let dcx = dcx.handle();
223223
let (symbols_below_threshold, upstream_modules) = prepare_lto(cgcx, dcx)?;
@@ -248,7 +248,7 @@ fn fat_lto(
248248
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
249249
mut serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
250250
symbols_below_threshold: &[*const libc::c_char],
251-
) -> Result<LtoModuleCodegen<LlvmCodegenBackend>, FatalError> {
251+
) -> Result<ModuleCodegen<ModuleLlvm>, FatalError> {
252252
let _timer = cgcx.prof.generic_activity("LLVM_fat_lto_build_monolithic_module");
253253
info!("going for a fat lto");
254254

@@ -366,7 +366,7 @@ fn fat_lto(
366366
save_temp_bitcode(cgcx, &module, "lto.after-restriction");
367367
}
368368

369-
Ok(LtoModuleCodegen::Fat(module))
369+
Ok(module)
370370
}
371371

372372
pub(crate) struct Linker<'a>(&'a mut llvm::Linker<'a>);
@@ -436,7 +436,7 @@ fn thin_lto(
436436
serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
437437
cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
438438
symbols_below_threshold: &[*const libc::c_char],
439-
) -> Result<(Vec<LtoModuleCodegen<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
439+
) -> Result<(Vec<ThinModule<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
440440
let _timer = cgcx.prof.generic_activity("LLVM_thin_lto_global_analysis");
441441
unsafe {
442442
info!("going for that thin, thin LTO");
@@ -568,10 +568,7 @@ fn thin_lto(
568568
}
569569

570570
info!(" - {}: re-compiled", module_name);
571-
opt_jobs.push(LtoModuleCodegen::Thin(ThinModule {
572-
shared: Arc::clone(&shared),
573-
idx: module_index,
574-
}));
571+
opt_jobs.push(ThinModule { shared: Arc::clone(&shared), idx: module_index });
575572
}
576573

577574
// Save the current ThinLTO import information for the next compilation

‎compiler/rustc_codegen_llvm/src/back/write.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -817,10 +817,12 @@ pub(crate) fn link(
817817

818818
pub(crate) fn codegen(
819819
cgcx: &CodegenContext<LlvmCodegenBackend>,
820-
dcx: DiagCtxtHandle<'_>,
821820
module: ModuleCodegen<ModuleLlvm>,
822821
config: &ModuleConfig,
823822
) -> Result<CompiledModule, FatalError> {
823+
let dcx = cgcx.create_dcx();
824+
let dcx = dcx.handle();
825+
824826
let _timer = cgcx.prof.generic_activity_with_arg("LLVM_module_codegen", &*module.name);
825827
{
826828
let llmod = module.module_llvm.llmod();

‎compiler/rustc_codegen_llvm/src/builder/autodiff.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ use std::ptr;
22

33
use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivity, DiffMode};
44
use rustc_codegen_ssa::ModuleCodegen;
5-
use rustc_codegen_ssa::back::write::ModuleConfig;
65
use rustc_codegen_ssa::common::TypeKind;
76
use rustc_codegen_ssa::traits::BaseTypeCodegenMethods;
87
use rustc_errors::FatalError;
@@ -461,7 +460,6 @@ pub(crate) fn differentiate<'ll>(
461460
module: &'ll ModuleCodegen<ModuleLlvm>,
462461
cgcx: &CodegenContext<LlvmCodegenBackend>,
463462
diff_items: Vec<AutoDiffItem>,
464-
_config: &ModuleConfig,
465463
) -> Result<(), FatalError> {
466464
for item in &diff_items {
467465
trace!("{}", item);

‎compiler/rustc_codegen_llvm/src/errors.rs‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,6 @@ impl<G: EmissionGuarantee> Diagnostic<'_, G> for ParseTargetMachineConfig<'_> {
3737
}
3838
}
3939

40-
#[derive(Diagnostic)]
41-
#[diag(codegen_llvm_autodiff_without_lto)]
42-
pub(crate) struct AutoDiffWithoutLTO;
43-
4440
#[derive(Diagnostic)]
4541
#[diag(codegen_llvm_autodiff_without_enable)]
4642
pub(crate) struct AutoDiffWithoutEnable;

‎compiler/rustc_codegen_llvm/src/lib.rs‎

Lines changed: 19 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,11 @@ use std::mem::ManuallyDrop;
2626
use back::owned_target_machine::OwnedTargetMachine;
2727
use back::write::{create_informational_target_machine, create_target_machine};
2828
use context::SimpleCx;
29-
use errors::{AutoDiffWithoutLTO, ParseTargetMachineConfig};
29+
use errors::ParseTargetMachineConfig;
3030
use llvm_util::target_config;
3131
use rustc_ast::expand::allocator::AllocatorKind;
3232
use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
33-
use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};
33+
use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule};
3434
use rustc_codegen_ssa::back::write::{
3535
CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryConfig, TargetMachineFactoryFn,
3636
};
@@ -43,7 +43,7 @@ use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
4343
use rustc_middle::ty::TyCtxt;
4444
use rustc_middle::util::Providers;
4545
use rustc_session::Session;
46-
use rustc_session::config::{Lto, OptLevel, OutputFilenames, PrintKind, PrintRequest};
46+
use rustc_session::config::{OptLevel, OutputFilenames, PrintKind, PrintRequest};
4747
use rustc_span::Symbol;
4848

4949
mod back {
@@ -174,18 +174,29 @@ impl WriteBackendMethods for LlvmCodegenBackend {
174174
) -> Result<ModuleCodegen<Self::Module>, FatalError> {
175175
back::write::link(cgcx, dcx, modules)
176176
}
177-
fn run_fat_lto(
177+
fn run_and_optimize_fat_lto(
178178
cgcx: &CodegenContext<Self>,
179179
modules: Vec<FatLtoInput<Self>>,
180180
cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
181-
) -> Result<LtoModuleCodegen<Self>, FatalError> {
182-
back::lto::run_fat(cgcx, modules, cached_modules)
181+
diff_fncs: Vec<AutoDiffItem>,
182+
) -> Result<ModuleCodegen<Self::Module>, FatalError> {
183+
let mut module = back::lto::run_fat(cgcx, modules, cached_modules)?;
184+
185+
if !diff_fncs.is_empty() {
186+
builder::autodiff::differentiate(&module, cgcx, diff_fncs)?;
187+
}
188+
189+
let dcx = cgcx.create_dcx();
190+
let dcx = dcx.handle();
191+
back::lto::run_pass_manager(cgcx, dcx, &mut module, false)?;
192+
193+
Ok(module)
183194
}
184195
fn run_thin_lto(
185196
cgcx: &CodegenContext<Self>,
186197
modules: Vec<(String, Self::ThinBuffer)>,
187198
cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
188-
) -> Result<(Vec<LtoModuleCodegen<Self>>, Vec<WorkProduct>), FatalError> {
199+
) -> Result<(Vec<ThinModule<Self>>, Vec<WorkProduct>), FatalError> {
189200
back::lto::run_thin(cgcx, modules, cached_modules)
190201
}
191202
fn optimize(
@@ -196,14 +207,6 @@ impl WriteBackendMethods for LlvmCodegenBackend {
196207
) -> Result<(), FatalError> {
197208
back::write::optimize(cgcx, dcx, module, config)
198209
}
199-
fn optimize_fat(
200-
cgcx: &CodegenContext<Self>,
201-
module: &mut ModuleCodegen<Self::Module>,
202-
) -> Result<(), FatalError> {
203-
let dcx = cgcx.create_dcx();
204-
let dcx = dcx.handle();
205-
back::lto::run_pass_manager(cgcx, dcx, module, false)
206-
}
207210
fn optimize_thin(
208211
cgcx: &CodegenContext<Self>,
209212
thin: ThinModule<Self>,
@@ -212,11 +215,10 @@ impl WriteBackendMethods for LlvmCodegenBackend {
212215
}
213216
fn codegen(
214217
cgcx: &CodegenContext<Self>,
215-
dcx: DiagCtxtHandle<'_>,
216218
module: ModuleCodegen<Self::Module>,
217219
config: &ModuleConfig,
218220
) -> Result<CompiledModule, FatalError> {
219-
back::write::codegen(cgcx, dcx, module, config)
221+
back::write::codegen(cgcx, module, config)
220222
}
221223
fn prepare_thin(
222224
module: ModuleCodegen<Self::Module>,
@@ -227,19 +229,6 @@ impl WriteBackendMethods for LlvmCodegenBackend {
227229
fn serialize_module(module: ModuleCodegen<Self::Module>) -> (String, Self::ModuleBuffer) {
228230
(module.name, back::lto::ModuleBuffer::new(module.module_llvm.llmod()))
229231
}
230-
/// Generate autodiff rules
231-
fn autodiff(
232-
cgcx: &CodegenContext<Self>,
233-
module: &ModuleCodegen<Self::Module>,
234-
diff_fncs: Vec<AutoDiffItem>,
235-
config: &ModuleConfig,
236-
) -> Result<(), FatalError> {
237-
if cgcx.lto != Lto::Fat {
238-
let dcx = cgcx.create_dcx();
239-
return Err(dcx.handle().emit_almost_fatal(AutoDiffWithoutLTO));
240-
}
241-
builder::autodiff::differentiate(module, cgcx, diff_fncs, config)
242-
}
243232
}
244233

245234
impl LlvmCodegenBackend {

0 commit comments

Comments
 (0)