1
Fork 0

upstream rustc_codegen_ssa/rustc_middle changes for enzyme/autodiff

This commit is contained in:
Manuel Drehwald 2025-01-29 21:31:13 -05:00
parent ebcf860e73
commit 1f30517d40
27 changed files with 482 additions and 38 deletions

View file

@ -189,6 +189,39 @@ pub enum CoverageLevel {
Mcdc,
}
/// The different settings that the `-Z autodiff` flag can have.
#[derive(Clone, Copy, PartialEq, Hash, Debug)]
pub enum AutoDiff {
/// Print TypeAnalysis information
PrintTA,
/// Print ActivityAnalysis Information
PrintAA,
/// Print Performance Warnings from Enzyme
PrintPerf,
/// Combines the three print flags above.
Print,
/// Print the whole module, before running opts.
PrintModBefore,
/// Print the whole module just before we pass it to Enzyme.
/// For Debug purpose, prefer the OPT flag below
PrintModAfterOpts,
/// Print the module after Enzyme differentiated everything.
PrintModAfterEnzyme,
/// Enzyme's loose type debug helper (can cause incorrect gradients)
LooseTypes,
/// More flags
NoModOptAfter,
/// Tell Enzyme to run LLVM Opts on each function it generated. By default off,
/// since we already optimize the whole module after Enzyme is done.
EnableFncOpt,
NoVecUnroll,
RuntimeActivity,
/// Runs Enzyme specific Inlining
Inline,
}
/// Settings for `-Z instrument-xray` flag.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub struct InstrumentXRay {
@ -2902,7 +2935,7 @@ pub(crate) mod dep_tracking {
};
use super::{
BranchProtection, CFGuard, CFProtection, CollapseMacroDebuginfo, CoverageOptions,
AutoDiff, BranchProtection, CFGuard, CFProtection, CollapseMacroDebuginfo, CoverageOptions,
CrateType, DebugInfo, DebugInfoCompression, ErrorOutputType, FmtDebug, FunctionReturn,
InliningThreshold, InstrumentCoverage, InstrumentXRay, LinkerPluginLto, LocationDetail,
LtoCli, MirStripDebugInfo, NextSolverConfig, OomStrategy, OptLevel, OutFileName,
@ -2950,6 +2983,7 @@ pub(crate) mod dep_tracking {
}
impl_dep_tracking_hash_via_hash!(
AutoDiff,
bool,
usize,
NonZero<usize>,

View file

@ -398,6 +398,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: `Print`, `PrintTA`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfterOpts`, `PrintModAfterEnzyme`, `LooseTypes`, `NoModOptAfter`, `EnableFncOpt`, `NoVecUnroll`, `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";
@ -1029,6 +1030,38 @@ pub mod parse {
}
}
pub(crate) fn parse_autodiff(slot: &mut Vec<AutoDiff>, v: Option<&str>) -> bool {
let Some(v) = v else {
*slot = vec![];
return true;
};
let mut v: Vec<&str> = v.split(",").collect();
v.sort_unstable();
for &val in v.iter() {
let variant = match val {
"PrintTA" => AutoDiff::PrintTA,
"PrintAA" => AutoDiff::PrintAA,
"PrintPerf" => AutoDiff::PrintPerf,
"Print" => AutoDiff::Print,
"PrintModBefore" => AutoDiff::PrintModBefore,
"PrintModAfterOpts" => AutoDiff::PrintModAfterOpts,
"PrintModAfterEnzyme" => AutoDiff::PrintModAfterEnzyme,
"LooseTypes" => AutoDiff::LooseTypes,
"NoModOptAfter" => AutoDiff::NoModOptAfter,
"EnableFncOpt" => AutoDiff::EnableFncOpt,
"NoVecUnroll" => AutoDiff::NoVecUnroll,
"Inline" => AutoDiff::Inline,
_ => {
// FIXME(ZuseZ4): print an error saying which value is not recognized
return false;
}
};
slot.push(variant);
}
true
}
pub(crate) fn parse_instrument_coverage(
slot: &mut InstrumentCoverage,
v: Option<&str>,
@ -1736,6 +1769,22 @@ options! {
either `loaded` or `not-loaded`."),
assume_incomplete_release: bool = (false, parse_bool, [TRACKED],
"make cfg(version) treat the current version as incomplete (default: no)"),
autodiff: Vec<crate::config::AutoDiff> = (Vec::new(), parse_autodiff, [TRACKED],
"a list of optional autodiff flags to enable
Optional extra settings:
`=PrintTA`
`=PrintAA`
`=PrintPerf`
`=Print`
`=PrintModBefore`
`=PrintModAfterOpts`
`=PrintModAfterEnzyme`
`=LooseTypes`
`=NoModOptAfter`
`=EnableFncOpt`
`=NoVecUnroll`
`=Inline`
Multiple options can be combined with commas."),
#[rustc_lint_opt_deny_field_access("use `Session::binary_dep_depinfo` instead of this field")]
binary_dep_depinfo: bool = (false, parse_bool, [TRACKED],
"include artifacts (sysroot, crate dependencies) used during compilation in dep-info \