remove perf impact of autodiff when not in use

Author: ZuseZ4

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -17,7 +17,6 @@ use crate::{Ty, TyKind};
 /// functions. The proper solution is to recognize and resolve this DAG of autodiff invocations,
 /// as it's already done in the C++ and Julia frontend of Enzyme.
 ///
-/// (FIXME) remove *First variants.
 /// Documentation for using [reverse](https://enzyme.mit.edu/rust/rev.html) and
 /// [forward](https://enzyme.mit.edu/rust/fwd.html) mode is available online.
 #[derive(Clone, Copy, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -607,19 +607,20 @@ pub(crate) fn run_pass_manager(
     // If this rustc version was build with enzyme/autodiff enabled, and if users applied the
     // `#[autodiff]` macro at least once, then we will later call llvm_optimize a second time.
     debug!("running llvm pm opt pipeline");
+
+    // The PostAD behavior is the same that we would have if no autodiff was used.
+    // It will run the default optimization pipeline. If AD is enabled we select
+    // the DuringAD stage, which will disable vectorization and loop unrolling, and
+    // schedule two autodiff optimization + differentiation passes.
+    // 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 };
     unsafe {
-        write::llvm_optimize(
-            cgcx,
-            dcx,
-            module,
-            config,
-            opt_level,
-            opt_stage,
-            write::AutodiffStage::DuringAD,
-        )?;
+        write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, stage)?;
     }
-    // FIXME(ZuseZ4): Make this more granular
-    if cfg!(llvm_enzyme) && !thin {
+    if cfg!(llvm_enzyme) && !thin && enable_ad {
         unsafe {
             write::llvm_optimize(
                 cgcx,

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -564,19 +564,16 @@ pub(crate) unsafe fn llvm_optimize(
     // FIXME(ZuseZ4): In a future update we could figure out how to only optimize individual functions getting
     // differentiated.
 
+    let consider_ad = cfg!(llvm_enzyme) && config.autodiff.contains(&config::AutoDiff::Enable);
+    let run_enzyme = autodiff_stage == AutodiffStage::DuringAD;
     let unroll_loops;
     let vectorize_slp;
     let vectorize_loop;
-    let run_enzyme = cfg!(llvm_enzyme) && autodiff_stage == AutodiffStage::DuringAD;
 
     // When we build rustc with enzyme/autodiff support, we want to postpone size-increasing
     // optimizations until after differentiation. Our pipeline is thus: (opt + enzyme), (full opt).
     // We therefore have two calls to llvm_optimize, if autodiff is used.
-    //
-    // FIXME(ZuseZ4): Before shipping on nightly,
-    // we should make this more granular, or at least check that the user has at least one autodiff
-    // call in their code, to justify altering the compilation pipeline.
-    if cfg!(llvm_enzyme) && autodiff_stage != AutodiffStage::PostAD {
+    if consider_ad && autodiff_stage != AutodiffStage::PostAD {
         unroll_loops = false;
         vectorize_slp = false;
         vectorize_loop = false;
@@ -706,10 +703,8 @@ pub(crate) unsafe fn optimize(
 
         // If we know that we will later run AD, then we disable vectorization and loop unrolling.
         // Otherwise we pretend AD is already done and run the normal opt pipeline (=PostAD).
-        // FIXME(ZuseZ4): Make this more granular, only set PreAD if we actually have autodiff
-        // usages, not just if we build rustc with autodiff support.
-        let autodiff_stage =
-            if cfg!(llvm_enzyme) { AutodiffStage::PreAD } else { AutodiffStage::PostAD };
+        let consider_ad = cfg!(llvm_enzyme) && config.autodiff.contains(&config::AutoDiff::Enable);
+        let autodiff_stage = if consider_ad { AutodiffStage::PreAD } else { AutodiffStage::PostAD };
         return unsafe {
             llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, autodiff_stage)
         };

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -46,9 +46,6 @@ fn generate_enzyme_call<'ll>(
     let output = attrs.ret_activity;
 
     // We have to pick the name depending on whether we want forward or reverse mode autodiff.
-    // FIXME(ZuseZ4): The new pass based approach should not need the {Forward/Reverse}First method anymore, since
-    // it will handle higher-order derivatives correctly automatically (in theory). Currently
-    // higher-order derivatives fail, so we should debug that before adjusting this code.
     let mut ad_name: String = match attrs.mode {
         DiffMode::Forward => "__enzyme_fwddiff",
         DiffMode::Reverse => "__enzyme_autodiff",

compiler/rustc_interface/src/tests.rs

@@ -759,7 +759,7 @@ fn test_unstable_options_tracking_hash() {
     tracked!(allow_features, Some(vec![String::from("lang_items")]));
     tracked!(always_encode_mir, true);
     tracked!(assume_incomplete_release, true);
-    tracked!(autodiff, vec![AutoDiff::Print]);
+    tracked!(autodiff, vec![AutoDiff::Enable]);
     tracked!(binary_dep_depinfo, true);
     tracked!(box_noalias, false);
     tracked!(

compiler/rustc_session/src/config.rs

@@ -198,33 +198,28 @@ pub enum CoverageLevel {
 /// The different settings that the `-Z autodiff` flag can have.
 #[derive(Clone, Copy, PartialEq, Hash, Debug)]
 pub enum AutoDiff {
+    /// Enable the autodiff opt pipeline
+    Enable,
+
     /// Print TypeAnalysis information
     PrintTA,
     /// Print ActivityAnalysis Information
     PrintAA,
     /// Print Performance Warnings from Enzyme
     PrintPerf,
-    /// Combines the three print flags above.
-    Print,
+    /// Print intermediate IR generation steps
+    PrintSteps,
     /// 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,
+    PrintModAfter,
 
-    /// Enzyme's loose type debug helper (can cause incorrect gradients)
+    /// Enzyme's loose type debug helper (can cause incorrect gradients!!)
+    /// Usable in cases where Enzyme errors with `can not deduce type of X`.
     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,
+    /// See Enzyme core documentation. FIXME(ZuseZ4): Clarify usages
     RuntimeActivity,
-    /// Runs Enzyme specific Inlining
+    /// Runs Enzyme's aggressive inlining
     Inline,
 }
 

compiler/rustc_session/src/options.rs

@@ -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: `Print`, `PrintTA`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfterOpts`, `PrintModAfterEnzyme`, `LooseTypes`, `NoModOptAfter`, `EnableFncOpt`, `NoVecUnroll`, `Inline`";
+    pub(crate) const parse_autodiff: &str = "a comma separated list of settings: `Enable`, `PrintSteps`, `PrintTA`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfter`, `LooseTypes`, `RuntimeActivity`, `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";
@@ -1348,17 +1348,15 @@ pub mod parse {
         v.sort_unstable();
         for &val in v.iter() {
             let variant = match val {
+                "Enable" => AutoDiff::Enable,
                 "PrintTA" => AutoDiff::PrintTA,
                 "PrintAA" => AutoDiff::PrintAA,
                 "PrintPerf" => AutoDiff::PrintPerf,
-                "Print" => AutoDiff::Print,
+                "PrintSteps" => AutoDiff::PrintSteps,
                 "PrintModBefore" => AutoDiff::PrintModBefore,
-                "PrintModAfterOpts" => AutoDiff::PrintModAfterOpts,
-                "PrintModAfterEnzyme" => AutoDiff::PrintModAfterEnzyme,
+                "PrintModAfter" => AutoDiff::PrintModAfter,
                 "LooseTypes" => AutoDiff::LooseTypes,
-                "NoModOptAfter" => AutoDiff::NoModOptAfter,
-                "EnableFncOpt" => AutoDiff::EnableFncOpt,
-                "NoVecUnroll" => AutoDiff::NoVecUnroll,
+                "RuntimeActivity" => AutoDiff::RuntimeActivity,
                 "Inline" => AutoDiff::Inline,
                 _ => {
                     // FIXME(ZuseZ4): print an error saying which value is not recognized
@@ -2081,21 +2079,20 @@ options! {
     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."),
+        "a list of autodiff flags to enable
+        Mandatory setting:
+        `=Enable`
+        Optional extra settings:
+        `=PrintTA`
+        `=PrintAA`
+        `=PrintPerf`
+        `=PrintSteps`
+        `=PrintModBefore`
+        `=PrintModAfter`
+        `=LooseTypes`
+        '=RuntimeActivity`
+        `=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 \

src/doc/unstable-book/src/compiler-flags/autodiff.md

@@ -8,16 +8,13 @@ This feature allows you to differentiate functions using automatic differentiati
 Set the `-Zautodiff=<options>` compiler flag to adjust the behaviour of the autodiff feature.
 Multiple options can be separated with a comma. Valid options are:
 
+`Enable` - Required flag to enable autodiff
 `PrintTA` - print Type Analysis Information
 `PrintAA` - print Activity Analysis Information
 `PrintPerf` - print Performance Warnings from Enzyme
-`Print` - prints all intermediate transformations
+`PrintSteps` - prints all intermediate transformations
 `PrintModBefore` - print the whole module, before running opts
-`PrintModAfterOpts` - print the whole module just before we pass it to Enzyme
-`PrintModAfterEnzyme` - print the module after Enzyme differentiated everything
+`PrintModAfter` - print the module after Enzyme differentiated everything
 `LooseTypes` - Enzyme's loose type debug helper (can cause incorrect gradients)
 `Inline` - runs Enzyme specific Inlining
-`NoModOptAfter` - do not optimize the module after Enzyme is done
-`EnableFncOpt` - tell Enzyme to run LLVM Opts on each function it generated
-`NoVecUnroll` - do not unroll vectorized loops
 `RuntimeActivity` - allow specifying activity at runtime