rust-lang/rust
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix LooseTypes flag and PrintMod behaviour, add debug helper

Browse Source
This commit is contained in:
Manuel Drehwald
2025-04-12 00:50:41 -04:00
parent e643f59f6d
commit 75f86e6e2e
6 changed files with 68 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
compiler/rustc_codegen_llvm/src/back/lto.rs
View File

@@ -584,12 +584,10 @@ fn thin_lto(
}
}
fn enable_autodiff_settings(ad: &[config::AutoDiff], module: &mut ModuleCodegen<ModuleLlvm>) {
fn enable_autodiff_settings(ad: &[config::AutoDiff]) {
for &val in ad {
// We intentionally don't use a wildcard, to not forget handling anything new.
match val {
config::AutoDiff::PrintModBefore => {
unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
}
config::AutoDiff::PrintPerf => {
llvm::set_print_perf(true);
}
@@ -603,17 +601,23 @@ fn enable_autodiff_settings(ad: &[config::AutoDiff], module: &mut ModuleCodegen<
llvm::set_inline(true);
}
config::AutoDiff::LooseTypes => {
llvm::set_loose_types(false);
llvm::set_loose_types(true);
}
config::AutoDiff::PrintSteps => {
llvm::set_print(true);
}
// We handle this below
// We handle this in the PassWrapper.cpp
config::AutoDiff::PrintPasses => {}
// We handle this in the PassWrapper.cpp
config::AutoDiff::PrintModBefore => {}
// We handle this in the PassWrapper.cpp
config::AutoDiff::PrintModAfter => {}
// We handle this below
// We handle this in the PassWrapper.cpp
config::AutoDiff::PrintModFinal => {}
// This is required and already checked
config::AutoDiff::Enable => {}
// We handle this below
config::AutoDiff::NoPostopt => {}
}
}
// This helps with handling enums for now.
@@ -647,27 +651,27 @@ pub(crate) fn run_pass_manager(
// We then run the llvm_optimize function a second time, to optimize the code which we generated
// in the enzyme differentiation pass.
let enable_ad = config.autodiff.contains(&config::AutoDiff::Enable);
let stage =
if enable_ad { write::AutodiffStage::DuringAD } else { write::AutodiffStage::PostAD };
let stage = if thin {
write::AutodiffStage::PreAD
} else {
if enable_ad { write::AutodiffStage::DuringAD } else { write::AutodiffStage::PostAD }
};
if enable_ad {
enable_autodiff_settings(&config.autodiff, module);
enable_autodiff_settings(&config.autodiff);
}
unsafe {
write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
}
if cfg!(llvm_enzyme) && enable_ad {
// This is the post-autodiff IR, mainly used for testing and educational purposes.
if config.autodiff.contains(&config::AutoDiff::PrintModAfter) {
unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
}
if cfg!(llvm_enzyme) && enable_ad && !thin {
let opt_stage = llvm::OptStage::FatLTO;
let stage = write::AutodiffStage::PostAD;
unsafe {
write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
if !config.autodiff.contains(&config::AutoDiff::NoPostopt) {
unsafe {
write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
}
}
// This is the final IR, so people should be able to inspect the optimized autodiff output,