cleanup and fix Non-enzyme build

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -19,7 +19,7 @@ use rustc_middle::bug;
 use rustc_middle::dep_graph::WorkProduct;
 use rustc_middle::middle::exported_symbols::{SymbolExportInfo, SymbolExportLevel};
 use rustc_session::config::{self, CrateType, Lto};
-use tracing::{debug, info, trace};
+use tracing::{debug, info};
 
 use crate::back::write::{
     self, CodegenDiagnosticsStage, DiagnosticHandlers, bitcode_section_name, save_temp_bitcode,
@@ -606,7 +606,7 @@ 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.
-    trace!("running llvm pm opt pipeline");
+    debug!("running llvm pm opt pipeline");
     unsafe {
         write::llvm_optimize(
             cgcx,
@@ -632,7 +632,7 @@ pub(crate) fn run_pass_manager(
             )?;
         }
     }
-    trace!("lto done");
+    debug!("lto done");
     Ok(())
 }
 

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -4,10 +4,9 @@ use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivit
 use rustc_codegen_ssa::ModuleCodegen;
 use rustc_codegen_ssa::back::write::ModuleConfig;
 use rustc_errors::FatalError;
-use rustc_session::config::Lto;
 use tracing::{debug, trace};
 
-use crate::back::write::{llvm_err, llvm_optimize};
+use crate::back::write::llvm_err;
 use crate::builder::SBuilder;
 use crate::context::SimpleCx;
 use crate::declare::declare_simple_fn;
@@ -153,7 +152,7 @@ fn generate_enzyme_call<'ll>(
             _ => {}
         }
 
-        trace!("matching autodiff arguments");
+        debug!("matching autodiff arguments");
         // We now handle the issue that Rust level arguments not always match the llvm-ir level
         // arguments. A slice, `&[f32]`, for example, is represented as a pointer and a length on
         // llvm-ir level. The number of activities matches the number of Rust level arguments, so we
@@ -320,29 +319,6 @@ pub(crate) fn differentiate<'ll>(
 
     // FIXME(ZuseZ4): support SanitizeHWAddress and prevent illegal/unsupported opts
 
-    //if let Some(opt_level) = config.opt_level {
-    //    let opt_stage = match cgcx.lto {
-    //        Lto::Fat => llvm::OptStage::PreLinkFatLTO,
-    //        Lto::Thin | Lto::ThinLocal => llvm::OptStage::PreLinkThinLTO,
-    //        _ if cgcx.opts.cg.linker_plugin_lto.enabled() => llvm::OptStage::PreLinkThinLTO,
-    //        _ => llvm::OptStage::PreLinkNoLTO,
-    //    };
-    //    // This is our second opt call, so now we run all opts,
-    //    // to make sure we get the best performance.
-    //    let skip_size_increasing_opts = false;
-    //    trace!("running Module Optimization after differentiation");
-    //    unsafe {
-    //        llvm_optimize(
-    //            cgcx,
-    //            diag_handler.handle(),
-    //            module,
-    //            config,
-    //            opt_level,
-    //            opt_stage,
-    //            skip_size_increasing_opts,
-    //        )?
-    //    };
-    //}
     trace!("done with differentiate()");
 
     Ok(())

compiler/rustc_llvm/build.rs

@@ -193,25 +193,10 @@ fn main() {
         cfg.define(&flag, None);
     }
 
-    // This doesn't work
-    if tracked_env_var_os("llvm_enzyme").is_some() {
-        // If we're just running `check`, there's no need for LLVM to be built.
-        loop {
-            println!("ENZYME!");
-        }
-    }
-
-    // This doesn't work either
-    if cfg!(llvm_enzyme) {
-        // Enzyme is a fork of LLVM, so we need to use a different build script.
-        loop {
-            println!("ENZYME!");
-        }
+    if tracked_env_var_os("LLVM_ENZYME").is_some() {
+        cfg.define("ENZYME", None);
     }
 
-    // FIXME: Only enable it, if the user build rustc with autodiff support
-    cfg.define("ENZYME", None);
-
     if tracked_env_var_os("LLVM_RUSTLLVM").is_some() {
         cfg.define("LLVM_RUSTLLVM", None);
     }

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

@@ -1234,6 +1234,9 @@ fn rustc_llvm_env(builder: &Builder<'_>, cargo: &mut Cargo, target: TargetSelect
     if builder.is_rust_llvm(target) {
         cargo.env("LLVM_RUSTLLVM", "1");
     }
+    if builder.config.llvm_enzyme {
+        cargo.env("LLVM_ENZYME", "1");
+    }
     let llvm::LlvmResult { llvm_config, .. } = builder.ensure(llvm::Llvm { target });
     cargo.env("LLVM_CONFIG", &llvm_config);