bjoernager/rust
bjoernager/rust
1
Fork 0
Code Issues Pull requests Activity

add new flag to print the module post-AD, before opts

Browse source
This commit is contained in:
Manuel Drehwald 2025-04-04 14:25:23 -04:00
parent 79e17bc71e
commit 89d8948835
3 changed files with 17 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

12
compiler/rustc_codegen_llvm/src/back/lto.rs
View file

@ -610,6 +610,8 @@ fn enable_autodiff_settings(ad: &[config::AutoDiff], module: &mut ModuleCodegen<
}
// We handle this below
config::AutoDiff::PrintModAfter => {}
// We handle this below
config::AutoDiff::PrintModFinal => {}
// This is required and already checked
config::AutoDiff::Enable => {}
}
@ -657,14 +659,20 @@ pub(crate) fn run_pass_manager(
}
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()) };
}
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)?;
}
// This is the final IR, so people should be able to inspect the optimized autodiff output.
if config.autodiff.contains(&config::AutoDiff::PrintModAfter) {
// This is the final IR, so people should be able to inspect the optimized autodiff output,
// for manual inspection.
if config.autodiff.contains(&config::AutoDiff::PrintModFinal) {
unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
}
}

6
compiler/rustc_session/src/config.rs
View file

@ -235,10 +235,12 @@ pub enum AutoDiff {
PrintPerf,
/// Print intermediate IR generation steps
PrintSteps,
/// Print the whole module, before running opts.
/// Print the module, before running autodiff.
PrintModBefore,
/// Print the module after Enzyme differentiated everything.
/// Print the module after running autodiff.
PrintModAfter,
/// Print the module after running autodiff and optimizations.
PrintModFinal,
/// Enzyme's loose type debug helper (can cause incorrect gradients!!)
/// Usable in cases where Enzyme errors with `can not deduce type of X`.

4
compiler/rustc_session/src/options.rs
View file

@ -707,7 +707,7 @@ mod desc {
pub(crate) const parse_list: &str = "a space-separated list of strings";
pub(crate) const parse_list_with_polarity: &str =
"a comma-separated list of strings, with elements beginning with + or -";
pub(crate) const parse_autodiff: &str = "a comma separated list of settings: `Enable`, `PrintSteps`, `PrintTA`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfter`, `LooseTypes`, `Inline`";
pub(crate) const parse_autodiff: &str = "a comma separated list of settings: `Enable`, `PrintSteps`, `PrintTA`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfter`, `PrintModFinal`, `LooseTypes`, `Inline`";
pub(crate) const parse_comma_list: &str = "a comma-separated list of strings";
pub(crate) const parse_opt_comma_list: &str = parse_comma_list;
pub(crate) const parse_number: &str = "a number";
@ -1355,6 +1355,7 @@ pub mod parse {
"PrintSteps" => AutoDiff::PrintSteps,
"PrintModBefore" => AutoDiff::PrintModBefore,
"PrintModAfter" => AutoDiff::PrintModAfter,
"PrintModFinal" => AutoDiff::PrintModFinal,
"LooseTypes" => AutoDiff::LooseTypes,
"Inline" => AutoDiff::Inline,
_ => {
@ -2088,6 +2089,7 @@ options! {
`=PrintSteps`
`=PrintModBefore`
`=PrintModAfter`
`=PrintModFinal`
`=LooseTypes`
`=Inline`
Multiple options can be combined with commas."),