Rollup merge of #146598 - bjorn3:feature_llvm_enzyme, r=davidtwco

Make llvm_enzyme a regular cargo feature

This makes it clearer that it is set by the build system rather than by the rustc that compiles the current rustc. It also avoids bootstrap needing to pass `--check-cfg llvm_enzyme` to rustc.

Author: jdonszelmann

compiler/rustc/Cargo.toml

@@ -30,6 +30,7 @@ features = ['unprefixed_malloc_on_supported_platforms']
 check_only = ['rustc_driver_impl/check_only']
 jemalloc = ['dep:tikv-jemalloc-sys']
 llvm = ['rustc_driver_impl/llvm']
+llvm_enzyme = ['rustc_driver_impl/llvm_enzyme']
 max_level_info = ['rustc_driver_impl/max_level_info']
 rustc_randomized_layouts = ['rustc_driver_impl/rustc_randomized_layouts']
 # tidy-alphabetical-end

compiler/rustc_builtin_macros/Cargo.toml

@@ -33,3 +33,8 @@ smallvec = { version = "1.8.1", features = ["union", "may_dangle"] }
 thin-vec = "0.2.12"
 tracing = "0.1"
 # tidy-alphabetical-end
+
+[features]
+# tidy-alphabetical-start
+llvm_enzyme = []
+# tidy-alphabetical-end

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -209,7 +209,8 @@ mod llvm_enzyme {
         mut item: Annotatable,
         mode: DiffMode,
     ) -> Vec<Annotatable> {
-        if cfg!(not(llvm_enzyme)) {
+        // FIXME(bjorn3) maybe have the backend directly tell if autodiff is supported?
+        if cfg!(not(feature = "llvm_enzyme")) {
             ecx.sess.dcx().emit_err(errors::AutoDiffSupportNotBuild { span: meta_item.span });
             return vec![item];
         }

compiler/rustc_codegen_llvm/Cargo.toml

@@ -46,5 +46,6 @@ tracing = "0.1"
 [features]
 # tidy-alphabetical-start
 check_only = ["rustc_llvm/check_only"]
+llvm_enzyme = []
 # tidy-alphabetical-end
 

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -617,7 +617,7 @@ pub(crate) fn run_pass_manager(
         crate::builder::gpu_offload::handle_gpu_code(cgcx, &cx);
     }
 
-    if cfg!(llvm_enzyme) && enable_ad && !thin {
+    if cfg!(feature = "llvm_enzyme") && enable_ad && !thin {
         let opt_stage = llvm::OptStage::FatLTO;
         let stage = write::AutodiffStage::PostAD;
         if !config.autodiff.contains(&config::AutoDiff::NoPostopt) {

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -574,7 +574,8 @@ 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 consider_ad =
+        cfg!(feature = "llvm_enzyme") && config.autodiff.contains(&config::AutoDiff::Enable);
     let run_enzyme = autodiff_stage == AutodiffStage::DuringAD;
     let print_before_enzyme = config.autodiff.contains(&config::AutoDiff::PrintModBefore);
     let print_after_enzyme = config.autodiff.contains(&config::AutoDiff::PrintModAfter);
@@ -740,7 +741,8 @@ pub(crate) 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).
-        let consider_ad = cfg!(llvm_enzyme) && config.autodiff.contains(&config::AutoDiff::Enable);
+        let consider_ad =
+            cfg!(feature = "llvm_enzyme") && config.autodiff.contains(&config::AutoDiff::Enable);
         let autodiff_stage = if consider_ad { AutodiffStage::PreAD } else { AutodiffStage::PostAD };
         // The embedded bitcode is used to run LTO/ThinLTO.
         // The bitcode obtained during the `codegen` phase is no longer suitable for performing LTO.

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -59,10 +59,10 @@ pub(crate) enum LLVMRustVerifierFailureAction {
     LLVMReturnStatusAction = 2,
 }
 
-#[cfg(llvm_enzyme)]
+#[cfg(feature = "llvm_enzyme")]
 pub(crate) use self::Enzyme_AD::*;
 
-#[cfg(llvm_enzyme)]
+#[cfg(feature = "llvm_enzyme")]
 pub(crate) mod Enzyme_AD {
     use std::ffi::{CString, c_char};
 
@@ -134,10 +134,10 @@ pub(crate) mod Enzyme_AD {
     }
 }
 
-#[cfg(not(llvm_enzyme))]
+#[cfg(not(feature = "llvm_enzyme"))]
 pub(crate) use self::Fallback_AD::*;
 
-#[cfg(not(llvm_enzyme))]
+#[cfg(not(feature = "llvm_enzyme"))]
 pub(crate) mod Fallback_AD {
     #![allow(unused_variables)]
 

compiler/rustc_driver_impl/Cargo.toml

@@ -74,6 +74,7 @@ ctrlc = "3.4.4"
 # tidy-alphabetical-start
 check_only = ['rustc_interface/check_only']
 llvm = ['rustc_interface/llvm']
+llvm_enzyme = ['rustc_interface/llvm_enzyme']
 max_level_info = ['rustc_log/max_level_info']
 rustc_randomized_layouts = [
     'rustc_index/rustc_randomized_layouts',

compiler/rustc_interface/Cargo.toml

@@ -58,4 +58,5 @@ rustc_abi = { path = "../rustc_abi" }
 # tidy-alphabetical-start
 check_only = ['rustc_codegen_llvm?/check_only']
 llvm = ['dep:rustc_codegen_llvm']
+llvm_enzyme = ['rustc_builtin_macros/llvm_enzyme', 'rustc_codegen_llvm/llvm_enzyme']
 # tidy-alphabetical-end

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

@@ -1357,10 +1357,6 @@ pub fn rustc_cargo_env(builder: &Builder<'_>, cargo: &mut Cargo, target: TargetS
         cargo.env("RUSTC_VERIFY_LLVM_IR", "1");
     }
 
-    if builder.config.llvm_enzyme {
-        cargo.rustflag("--cfg=llvm_enzyme");
-    }
-
     // These conditionals represent a tension between three forces:
     // - For non-check builds, we need to define some LLVM-related environment
     //   variables, requiring LLVM to have been built.