Auto merge of rust-lang#136428 - EnzymeAD:enable-autodiff, r=<try>

test building enzyme in CI

We will later need to make it more granular before merging, but let's start by figuring out how many things we break right now in the first place.

For now I modified `src/ci/docker/host-x86_64/dist-x86_64-linux/Dockerfile`, so let's try that dist job first.

try-job: dist-x86_64-apple
try-job: dist-apple-various

r? `@oli-obk`

Tracking:

- rust-lang#124509

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_builtin_macros/src/autodiff.rs

@@ -242,6 +242,7 @@ mod llvm_enzyme {
             defaultness: ast::Defaultness::Final,
             sig: d_sig,
             generics: Generics::default(),
+            contract: None,
             body: Some(d_body),
         });
         let mut rustc_ad_attr =

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_monomorphize/src/partitioning/autodiff.rs

@@ -66,7 +66,7 @@ pub(crate) fn find_autodiff_source_functions<'tcx>(
     let mut autodiff_items: Vec<AutoDiffItem> = vec![];
     for (item, instance) in autodiff_mono_items {
         let target_id = instance.def_id();
-        let cg_fn_attr = tcx.codegen_fn_attrs(target_id).autodiff_item.clone();
+        let cg_fn_attr = &tcx.codegen_fn_attrs(target_id).autodiff_item;
         let Some(target_attrs) = cg_fn_attr else {
             continue;
         };

compiler/rustc_session/src/config.rs

@@ -192,33 +192,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 \

config.example.toml

@@ -79,10 +79,10 @@
 #tests = false
 
 # Indicates whether the LLVM plugin is enabled or not
-#plugins = false
+#plugins = true
 
 # Whether to build Enzyme as AutoDiff backend.
-#enzyme = false
+#enzyme = true
 
 # Whether to build LLVM with support for it's gpu offload runtime.
 #offload = false

src/bootstrap/src/core/build_steps/compile.rs

@@ -1092,8 +1092,11 @@ pub fn rustc_cargo(
 
     // We want to link against registerEnzyme and in the future we want to use additional
     // functionality from Enzyme core. For that we need to link against Enzyme.
-    // FIXME(ZuseZ4): Get the LLVM version number automatically instead of hardcoding it.
     if builder.config.llvm_enzyme {
+        let arch = builder.build.build;
+        let enzyme_dir = builder.build.out.join(arch).join("enzyme").join("lib");
+        cargo.rustflag("-L").rustflag(enzyme_dir.to_str().expect("Invalid path"));
+        // FIXME(ZuseZ4): Get the LLVM version number automatically instead of hardcoding it.
         cargo.rustflag("-l").rustflag("Enzyme-19");
     }
 

src/ci/docker/host-x86_64/dist-x86_64-linux/Dockerfile

@@ -89,6 +89,8 @@ ENV RUST_CONFIGURE_ARGS \
       --set target.x86_64-unknown-linux-gnu.ar=/rustroot/bin/llvm-ar \
       --set target.x86_64-unknown-linux-gnu.ranlib=/rustroot/bin/llvm-ranlib \
       --set llvm.thin-lto=true \
+      --set llvm.plugins=true \
+      --set llvm.enzyme=true \
       --set llvm.ninja=false \
       --set llvm.libzstd=true \
       --set rust.jemalloc \

src/ci/docker/host-x86_64/x86_64-gnu/Dockerfile

@@ -29,5 +29,7 @@ ENV RUST_CONFIGURE_ARGS \
  --enable-sanitizers \
  --enable-profiler \
  --enable-compiler-docs \
+ --set llvm.plugins=true \
+ --set llvm.enzyme=true \
  --set llvm.libzstd=true
 ENV SCRIPT python3 ../x.py --stage 2 test

src/ci/github-actions/jobs.yml

@@ -369,7 +369,7 @@ auto:
   - name: dist-x86_64-apple
     env:
       SCRIPT: ./x.py dist bootstrap --include-default-paths --host=x86_64-apple-darwin --target=x86_64-apple-darwin
-      RUST_CONFIGURE_ARGS: --enable-full-tools --enable-sanitizers --enable-profiler --set rust.jemalloc --set rust.lto=thin --set rust.codegen-units=1
+      RUST_CONFIGURE_ARGS: --enable-full-tools --enable-sanitizers --enable-profiler --set rust.jemalloc --set rust.lto=thin --set rust.codegen-units=1 --set llvm.plugins=true --set llvm.enzyme=true
       RUSTC_RETRY_LINKER_ON_SEGFAULT: 1
       # Ensure that host tooling is built to support our minimum support macOS version.
       MACOSX_DEPLOYMENT_TARGET: 10.12
@@ -387,7 +387,7 @@ auto:
       SCRIPT: ./x.py dist bootstrap --include-default-paths --host='' --target=aarch64-apple-ios,x86_64-apple-ios,aarch64-apple-ios-sim,aarch64-apple-ios-macabi,x86_64-apple-ios-macabi
       # Mac Catalyst cannot currently compile the sanitizer:
       # https://github.com/rust-lang/rust/issues/129069
-      RUST_CONFIGURE_ARGS: --enable-sanitizers --enable-profiler --set rust.jemalloc --set target.aarch64-apple-ios-macabi.sanitizers=false --set target.x86_64-apple-ios-macabi.sanitizers=false
+      RUST_CONFIGURE_ARGS: --enable-sanitizers --enable-profiler --set rust.jemalloc --set target.aarch64-apple-ios-macabi.sanitizers=false --set target.x86_64-apple-ios-macabi.sanitizers=false --set llvm.plugins=true --set llvm.enzyme=true
       RUSTC_RETRY_LINKER_ON_SEGFAULT: 1
       # Ensure that host tooling is built to support our minimum support macOS version.
       # FIXME(madsmtm): This might be redundant, as we're not building host tooling here (?)

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

tests/codegen/autodiff.rs

@@ -1,4 +1,4 @@
-//@ compile-flags: -C opt-level=3  -Clto=fat
+//@ compile-flags: -Zautodiff=Enable -C opt-level=3  -Clto=fat
 //@ no-prefer-dynamic
 //@ needs-enzyme
 #![feature(autodiff)]