change enable_autodiff_settngs args + call enable_autodiff_settings from LlvmCodegenBackend::init

Author: sgasho

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -528,10 +528,13 @@ fn thin_lto(
     }
 }
 
-fn enable_autodiff_settings(cgcx: &CodegenContext<LlvmCodegenBackend>, ad: &[config::AutoDiff]) {
+pub(crate) fn enable_autodiff_settings(
+    sysroot: &rustc_session::config::Sysroot,
+    ad: &[config::AutoDiff],
+) {
     // Initialize Enzyme if not already done (idempotent due to OnceLock)
     // This ensures it works even if LlvmCodegenBackend::init() didn't run it
-    let mut enzyme = llvm::EnzymeWrapper::get_or_init(&cgcx.sysroot);
+    let mut enzyme = llvm::EnzymeWrapper::get_or_init(sysroot);
 
     for val in ad {
         // We intentionally don't use a wildcard, to not forget handling anything new.
@@ -612,9 +615,9 @@ pub(crate) fn run_pass_manager(
         if enable_ad { write::AutodiffStage::DuringAD } else { write::AutodiffStage::PostAD }
     };
 
-    if enable_ad {
-        enable_autodiff_settings(&cgcx, &config.autodiff);
-    }
+    // if enable_ad {
+    //     enable_autodiff_settings(&cgcx, &config.autodiff);
+    // }
 
     unsafe {
         write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage);

compiler/rustc_codegen_llvm/src/lib.rs

@@ -241,10 +241,14 @@ impl CodegenBackend for LlvmCodegenBackend {
     fn init(&self, sess: &Session) {
         llvm_util::init(sess); // Make sure llvm is inited
 
-        if sess.opts.unstable_opts.autodiff.contains(&AutoDiff::Enable) {
-            #[cfg(feature = "llvm_enzyme")]
-            {
-                drop(llvm::EnzymeWrapper::get_or_init(&sess.opts.sysroot));
+        #[cfg(feature = "llvm_enzyme")]
+        {
+            if sess.opts.unstable_opts.autodiff.contains(&AutoDiff::Enable) {
+                {
+                    use crate::back::lto::enable_autodiff_settings;
+
+                    enable_autodiff_settings(&sess.opts.sysroot, &sess.opts.unstable_opts.autodiff);
+                }
             }
         }
     }