EnzymeWrapper::get_instance to return MutexGuard, EnzymeWrapper::get_or_init takes only sysroot, not a whole session or context

Author: sgasho

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -529,33 +529,34 @@ fn thin_lto(
 }
 
 fn enable_autodiff_settings(cgcx: &CodegenContext<LlvmCodegenBackend>, ad: &[config::AutoDiff]) {
-    use std::sync::Mutex;
-    let enzyme: &'static Mutex<llvm::EnzymeWrapper> = llvm::EnzymeWrapper::init(cgcx);
+    // 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);
 
     for val in ad {
         // We intentionally don't use a wildcard, to not forget handling anything new.
         match val {
             config::AutoDiff::PrintPerf => {
-                enzyme.lock().unwrap().set_print_perf(true);
+                enzyme.set_print_perf(true);
             }
             config::AutoDiff::PrintAA => {
-                enzyme.lock().unwrap().set_print_activity(true);
+                enzyme.set_print_activity(true);
             }
             config::AutoDiff::PrintTA => {
-                enzyme.lock().unwrap().set_print_type(true);
+                enzyme.set_print_type(true);
             }
             config::AutoDiff::PrintTAFn(fun) => {
-                enzyme.lock().unwrap().set_print_type(true); // Enable general type printing
-                enzyme.lock().unwrap().set_print_type_fun(&fun); // Set specific function to analyze
+                enzyme.set_print_type(true); // Enable general type printing
+                enzyme.set_print_type_fun(&fun); // Set specific function to analyze
             }
             config::AutoDiff::Inline => {
-                enzyme.lock().unwrap().set_inline(true);
+                enzyme.set_inline(true);
             }
             config::AutoDiff::LooseTypes => {
-                enzyme.lock().unwrap().set_loose_types(true);
+                enzyme.set_loose_types(true);
             }
             config::AutoDiff::PrintSteps => {
-                enzyme.lock().unwrap().set_print(true);
+                enzyme.set_print(true);
             }
             // We handle this in the PassWrapper.cpp
             config::AutoDiff::PrintPasses => {}
@@ -574,9 +575,9 @@ fn enable_autodiff_settings(cgcx: &CodegenContext<LlvmCodegenBackend>, ad: &[con
         }
     }
     // This helps with handling enums for now.
-    enzyme.lock().unwrap().set_strict_aliasing(false);
+    enzyme.set_strict_aliasing(false);
     // FIXME(ZuseZ4): Test this, since it was added a long time ago.
-    enzyme.lock().unwrap().set_rust_rules(true);
+    enzyme.set_rust_rules(true);
 }
 
 pub(crate) fn run_pass_manager(

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -2,7 +2,7 @@ use std::ffi::{CStr, CString};
 use std::io::{self, Write};
 use std::path::{Path, PathBuf};
 use std::ptr::null_mut;
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 use std::{fs, slice, str};
 
 use libc::{c_char, c_int, c_void, size_t};
@@ -727,8 +727,9 @@ pub(crate) unsafe fn llvm_optimize(
     let llvm_plugins = config.llvm_plugins.join(",");
 
     let enzyme_fn = if consider_ad {
-        let wrapper: &'static Mutex<llvm::EnzymeWrapper> = llvm::EnzymeWrapper::init(cgcx);
-        wrapper.lock().unwrap().registerEnzymeAndPassPipeline
+        // Initialize Enzyme if not already done (idempotent due to OnceLock)
+        let wrapper = llvm::EnzymeWrapper::get_or_init(&cgcx.sysroot);
+        wrapper.registerEnzymeAndPassPipeline
     } else {
         std::ptr::null()
     };

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1099,7 +1099,13 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         // vs. copying a struct with mixed types requires different derivative handling.
         // The TypeTree tells Enzyme exactly what memory layout to expect.
         if let Some(tt) = tt {
-            crate::typetree::add_tt(self.cx().llmod, self.cx().llcx, memcpy, tt);
+            crate::typetree::add_tt(
+                self.cx().llmod,
+                self.cx().llcx,
+                memcpy,
+                tt,
+                &self.cx().tcx.sess.opts.sysroot,
+            );
         }
     }
 

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -375,7 +375,13 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
     );
 
     if !fnc_tree.args.is_empty() || !fnc_tree.ret.0.is_empty() {
-        crate::typetree::add_tt(cx.llmod, cx.llcx, fn_to_diff, fnc_tree);
+        crate::typetree::add_tt(
+            cx.llmod,
+            cx.llcx,
+            fn_to_diff,
+            fnc_tree,
+            &builder.cx().tcx.sess.opts.sysroot,
+        );
     }
 
     let call = builder.call(enzyme_ty, None, None, ad_fn, &args, None, None);

compiler/rustc_codegen_llvm/src/lib.rs

@@ -41,7 +41,7 @@ use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
 use rustc_middle::ty::TyCtxt;
 use rustc_middle::util::Providers;
 use rustc_session::Session;
-use rustc_session::config::{OptLevel, OutputFilenames, PrintKind, PrintRequest};
+use rustc_session::config::{AutoDiff, OptLevel, OutputFilenames, PrintKind, PrintRequest};
 use rustc_span::Symbol;
 use rustc_target::spec::{RelocModel, TlsModel};
 
@@ -240,6 +240,13 @@ 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));
+            }
+        }
     }
 
     fn provide(&self, providers: &mut Providers) {

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -97,10 +97,8 @@ pub(crate) use self::Enzyme_AD::*;
 #[cfg(feature = "llvm_enzyme")]
 pub(crate) mod Enzyme_AD {
     use std::ffi::{c_char, c_void};
-    use std::sync::{Mutex, OnceLock};
+    use std::sync::{Mutex, MutexGuard, OnceLock};
 
-    use rustc_codegen_ssa::back::write::CodegenContext;
-    use rustc_codegen_ssa::traits::WriteBackendMethods;
     use rustc_middle::bug;
     use rustc_session::config::{Sysroot, host_tuple};
     use rustc_session::filesearch;
@@ -202,18 +200,29 @@ pub(crate) mod Enzyme_AD {
     static ENZYME_INSTANCE: OnceLock<Mutex<EnzymeWrapper>> = OnceLock::new();
 
     impl EnzymeWrapper {
-        pub(crate) fn init<'a, B: WriteBackendMethods>(
-            cgcx: &'a CodegenContext<B>,
-        ) -> &'static Mutex<Self> {
-            ENZYME_INSTANCE.get_or_init(|| {
-                Self::call_dynamic(cgcx)
-                    .unwrap_or_else(|e| bug!("failed to load Enzyme: {e}"))
-                    .into()
-            })
+        /// Initialize EnzymeWrapper with the given sysroot if not already initialized.
+        /// Safe to call multiple times - subsequent calls are no-ops due to OnceLock.
+        pub(crate) fn get_or_init(
+            sysroot: &rustc_session::config::Sysroot,
+        ) -> MutexGuard<'static, Self> {
+            ENZYME_INSTANCE
+                .get_or_init(|| {
+                    Self::call_dynamic(sysroot)
+                        .unwrap_or_else(|e| bug!("failed to load Enzyme: {e}"))
+                        .into()
+                })
+                .lock()
+                .unwrap()
         }
 
-        pub(crate) fn get_instance() -> &'static Mutex<Self> {
-            ENZYME_INSTANCE.get().expect("EnzymeWrapper not initialized")
+        /// Get the EnzymeWrapper instance. Panics if not initialized.
+        /// Call get_or_init with a sysroot first.
+        pub(crate) fn get_instance() -> MutexGuard<'static, Self> {
+            ENZYME_INSTANCE
+                .get()
+                .expect("EnzymeWrapper not initialized. Call get_or_init with sysroot first.")
+                .lock()
+                .unwrap()
         }
 
         pub(crate) fn new_type_tree(&self) -> CTypeTreeRef {
@@ -350,10 +359,10 @@ pub(crate) mod Enzyme_AD {
         }
 
         #[allow(non_snake_case)]
-        fn call_dynamic<'a, B: WriteBackendMethods>(
-            cgcx: &'a CodegenContext<B>,
+        fn call_dynamic(
+            sysroot: &rustc_session::config::Sysroot,
         ) -> Result<Self, Box<dyn std::error::Error>> {
-            let enzyme_path = Self::get_enzyme_path(&cgcx.sysroot)?;
+            let enzyme_path = Self::get_enzyme_path(sysroot)?;
             let lib = unsafe { libloading::Library::new(enzyme_path)? };
 
             load_ptrs_by_symbols_fn!(
@@ -587,19 +596,19 @@ pub(crate) mod Fallback_AD {
 impl TypeTree {
     pub(crate) fn new() -> TypeTree {
         let wrapper = EnzymeWrapper::get_instance();
-        let inner = wrapper.lock().unwrap().new_type_tree();
+        let inner = wrapper.new_type_tree();
         TypeTree { inner }
     }
 
     pub(crate) fn from_type(t: CConcreteType, ctx: &Context) -> TypeTree {
         let wrapper = EnzymeWrapper::get_instance();
-        let inner = wrapper.lock().unwrap().new_type_tree_ct(t, ctx);
+        let inner = wrapper.new_type_tree_ct(t, ctx);
         TypeTree { inner }
     }
 
     pub(crate) fn merge(self, other: Self) -> Self {
         let wrapper = EnzymeWrapper::get_instance();
-        wrapper.lock().unwrap().merge_type_tree(self.inner, other.inner);
+        wrapper.merge_type_tree(self.inner, other.inner);
         drop(other);
         self
     }
@@ -614,7 +623,7 @@ impl TypeTree {
     ) -> Self {
         let layout = std::ffi::CString::new(layout).unwrap();
         let wrapper = EnzymeWrapper::get_instance();
-        wrapper.lock().unwrap().shift_indicies_eq(
+        wrapper.shift_indicies_eq(
             self.inner,
             layout.as_ptr(),
             offset as i64,
@@ -627,36 +636,30 @@ impl TypeTree {
 
     pub(crate) fn insert(&mut self, indices: &[i64], ct: CConcreteType, ctx: &Context) {
         let wrapper = EnzymeWrapper::get_instance();
-        wrapper.lock().unwrap().tree_insert_eq(
-            self.inner,
-            indices.as_ptr(),
-            indices.len(),
-            ct,
-            ctx,
-        );
+        wrapper.tree_insert_eq(self.inner, indices.as_ptr(), indices.len(), ct, ctx);
     }
 }
 
 impl Clone for TypeTree {
     fn clone(&self) -> Self {
         let wrapper = EnzymeWrapper::get_instance();
-        let inner = wrapper.lock().unwrap().new_type_tree_tr(self.inner);
+        let inner = wrapper.new_type_tree_tr(self.inner);
         TypeTree { inner }
     }
 }
 
 impl std::fmt::Display for TypeTree {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let wrapper = EnzymeWrapper::get_instance();
-        let ptr = wrapper.lock().unwrap().tree_to_string(self.inner);
+        let ptr = wrapper.tree_to_string(self.inner);
         let cstr = unsafe { std::ffi::CStr::from_ptr(ptr) };
         match cstr.to_str() {
             Ok(x) => write!(f, "{}", x)?,
             Err(err) => write!(f, "could not parse: {}", err)?,
         }
 
         // delete C string pointer
-        wrapper.lock().unwrap().tree_to_string_free(ptr);
+        wrapper.tree_to_string_free(ptr);
 
         Ok(())
     }
@@ -671,6 +674,6 @@ impl std::fmt::Debug for TypeTree {
 impl Drop for TypeTree {
     fn drop(&mut self) {
         let wrapper = EnzymeWrapper::get_instance();
-        wrapper.lock().unwrap().free_type_tree(self.inner)
+        wrapper.free_type_tree(self.inner)
     }
 }

compiler/rustc_codegen_llvm/src/typetree.rs

@@ -63,6 +63,7 @@ pub(crate) fn add_tt<'ll>(
     llcx: &'ll llvm::Context,
     fn_def: &'ll Value,
     tt: FncTree,
+    sysroot: &rustc_session::config::Sysroot,
 ) {
     let inputs = tt.args;
     let ret_tt: RustTypeTree = tt.ret;
@@ -75,7 +76,7 @@ pub(crate) fn add_tt<'ll>(
     let attr_name = "enzyme_type";
     let c_attr_name = CString::new(attr_name).unwrap();
 
-    let enzyme_wrapper = EnzymeWrapper::get_instance().lock().unwrap();
+    let enzyme_wrapper = EnzymeWrapper::get_or_init(sysroot);
 
     for (i, input) in inputs.iter().enumerate() {
         unsafe {
@@ -120,6 +121,7 @@ pub(crate) fn add_tt<'ll>(
     _llcx: &'ll llvm::Context,
     _fn_def: &'ll Value,
     _tt: FncTree,
+    _sysroot: &rustc_session::config::Sysroot,
 ) {
     unimplemented!()
 }

compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp

@@ -888,7 +888,7 @@ extern "C" LLVMRustResult LLVMRustOptimize(
   }
 
   // now load "-enzyme" pass:
-#ifdef ENZYME
+  // With dlopen, ENZYME macro may not be defined, so check EnzymePtr directly
   if (EnzymePtr) {
 
     if (PrintBeforeEnzyme) {
@@ -910,7 +910,6 @@ extern "C" LLVMRustResult LLVMRustOptimize(
       MPM.addPass(PrintModulePass(outs(), Banner, true, false));
     }
   }
-#endif
   if (PrintPasses) {
     // Print all passes from the PM:
     std::string Pipeline;