various updates

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};
+use tracing::{debug, info, trace};
 
 use crate::back::write::{
     self, CodegenDiagnosticsStage, DiagnosticHandlers, bitcode_section_name, save_temp_bitcode,
@@ -606,12 +606,33 @@ 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.
-    let first_run = true;
-    debug!("running llvm pm opt pipeline");
+    trace!("running llvm pm opt pipeline");
     unsafe {
-        write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, first_run)?;
+        write::llvm_optimize(
+            cgcx,
+            dcx,
+            module,
+            config,
+            opt_level,
+            opt_stage,
+            write::AutodiffStage::DuringAD,
+        )?;
     }
-    debug!("lto done");
+    // FIXME(ZuseZ4): Make this more granular
+    if cfg!(llvm_enzyme) && !thin {
+        unsafe {
+            write::llvm_optimize(
+                cgcx,
+                dcx,
+                module,
+                config,
+                opt_level,
+                llvm::OptStage::FatLTO,
+                write::AutodiffStage::PostAD,
+            )?;
+        }
+    }
+    trace!("lto done");
     Ok(())
 }
 

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -530,14 +530,24 @@ fn get_instr_profile_output_path(config: &ModuleConfig) -> Option<CString> {
     config.instrument_coverage.then(|| c"default_%m_%p.profraw".to_owned())
 }
 
+// PreAD will run llvm opts but disable size increasing opts (vectorization, loop unrolling)
+// DuringAD is the same as above, but also runs the enzyme opt and autodiff passes.
+// PostAD will run all opts, including size increasing opts.
+#[derive(Debug, Eq, PartialEq)]
+pub(crate) enum AutodiffStage {
+    PreAD,
+    DuringAD,
+    PostAD,
+}
+
 pub(crate) unsafe fn llvm_optimize(
     cgcx: &CodegenContext<LlvmCodegenBackend>,
     dcx: DiagCtxtHandle<'_>,
     module: &ModuleCodegen<ModuleLlvm>,
     config: &ModuleConfig,
     opt_level: config::OptLevel,
     opt_stage: llvm::OptStage,
-    skip_size_increasing_opts: bool,
+    autodiff_stage: AutodiffStage,
 ) -> Result<(), FatalError> {
     // Enzyme:
     // The whole point of compiler based AD is to differentiate optimized IR instead of unoptimized
@@ -550,7 +560,7 @@ pub(crate) unsafe fn llvm_optimize(
     let unroll_loops;
     let vectorize_slp;
     let vectorize_loop;
-    let run_enzyme;
+    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).
@@ -559,17 +569,15 @@ pub(crate) unsafe fn llvm_optimize(
     // 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 skip_size_increasing_opts && cfg!(llvm_enzyme) {
+    if cfg!(llvm_enzyme) && autodiff_stage != AutodiffStage::PostAD {
         unroll_loops = false;
         vectorize_slp = false;
         vectorize_loop = false;
-        run_enzyme = true;
     } else {
         unroll_loops =
             opt_level != config::OptLevel::Size && opt_level != config::OptLevel::SizeMin;
         vectorize_slp = config.vectorize_slp;
         vectorize_loop = config.vectorize_loop;
-        run_enzyme = false;
     }
     trace!(?unroll_loops, ?vectorize_slp, ?vectorize_loop, ?run_enzyme);
     let using_thin_buffers = opt_stage == llvm::OptStage::PreLinkThinLTO || config.bitcode_needed();
@@ -691,18 +699,14 @@ pub(crate) unsafe fn optimize(
             _ => llvm::OptStage::PreLinkNoLTO,
         };
 
-        // If we know that we will later run AD, then we disable vectorization and loop unrolling
-        let skip_size_increasing_opts = cfg!(llvm_enzyme);
+        // 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 };
         return unsafe {
-            llvm_optimize(
-                cgcx,
-                dcx,
-                module,
-                config,
-                opt_level,
-                opt_stage,
-                skip_size_increasing_opts,
-            )
+            llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, autodiff_stage)
         };
     }
     Ok(())

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -285,7 +285,7 @@ pub(crate) fn differentiate<'ll>(
     module: &'ll ModuleCodegen<ModuleLlvm>,
     cgcx: &CodegenContext<LlvmCodegenBackend>,
     diff_items: Vec<AutoDiffItem>,
-    config: &ModuleConfig,
+    _config: &ModuleConfig,
 ) -> Result<(), FatalError> {
     for item in &diff_items {
         trace!("{}", item);
@@ -320,29 +320,29 @@ 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,
-            )?
-        };
-    }
+    //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/Cargo.toml

@@ -14,3 +14,6 @@ libc = "0.2.73"
 # pinned `cc` in `rustc_codegen_ssa` if you update `cc` here.
 cc = "=1.2.7"
 # tidy-alphabetical-end
+
+[lints.rust]
+unexpected_cfgs = { level = "warn", check-cfg = ['cfg(llvm_enzyme)'] }

compiler/rustc_llvm/build.rs

@@ -193,6 +193,25 @@ 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!");
+        }
+    }
+
+    // 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);
     }

compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp

@@ -688,8 +688,11 @@ struct LLVMRustSanitizerOptions {
   bool SanitizeKernelAddressRecover;
 };
 
+
 // This symbol won't be available or used when Enzyme is not enabled
-extern "C" void registerEnzyme(llvm::PassBuilder &PB) __attribute__((weak));
+#ifdef ENZYME
+extern "C" void registerEnzyme(llvm::PassBuilder &PB);
+#endif
 
 extern "C" LLVMRustResult LLVMRustOptimize(
     LLVMModuleRef ModuleRef, LLVMTargetMachineRef TMRef,
@@ -1015,6 +1018,7 @@ extern "C" LLVMRustResult LLVMRustOptimize(
   }
 
   // now load "-enzyme" pass:
+#ifdef ENZYME
   if (RunEnzyme) {
     registerEnzyme(PB);
     if (auto Err = PB.parsePassPipeline(MPM, "enzyme")) {
@@ -1023,6 +1027,7 @@ extern "C" LLVMRustResult LLVMRustOptimize(
       return LLVMRustResult::Failure;
     }
   }
+#endif
 
   // Upgrade all calls to old intrinsics first.
   for (Module::iterator I = TheModule->begin(), E = TheModule->end(); I != E;)

tests/codegen/autodiff.rs

@@ -15,9 +15,9 @@ fn square(x: &f64) -> f64 {
 // CHECK-NEXT:invertstart:
 // CHECK-NEXT:  %_0 = fmul double %x.0.val, %x.0.val
 // CHECK-NEXT:  %0 = fadd fast double %x.0.val, %x.0.val
-// CHECK-NEXT:  %1 = load double, ptr %"x'", align 8, !alias.scope !17816, !noalias !17819
+// CHECK-NEXT:  %1 = load double, ptr %"x'", align 8
 // CHECK-NEXT:  %2 = fadd fast double %1, %0
-// CHECK-NEXT:  store double %2, ptr %"x'", align 8, !alias.scope !17816, !noalias !17819
+// CHECK-NEXT:  store double %2, ptr %"x'", align 8
 // CHECK-NEXT:  ret double %_0
 // CHECK-NEXT:}