move to OnceLock, compiles but fails all tests and SIGSEGV

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -523,30 +523,32 @@ fn thin_lto(
 }
 
 fn enable_autodiff_settings(ad: &[config::AutoDiff]) {
+    use crate::llvm::EnzymeWrapper;
+    let enzyme = EnzymeWrapper::current();
     for val in ad {
         // We intentionally don't use a wildcard, to not forget handling anything new.
         match val {
             config::AutoDiff::PrintPerf => {
-                llvm::set_print_perf(true);
+                enzyme.lock().unwrap().set_print_perf(true);
             }
             config::AutoDiff::PrintAA => {
-                llvm::set_print_activity(true);
+                enzyme.lock().unwrap().set_print_activity(true);
             }
             config::AutoDiff::PrintTA => {
-                llvm::set_print_type(true);
+                enzyme.lock().unwrap().set_print_type(true);
             }
             config::AutoDiff::PrintTAFn(fun) => {
-                llvm::set_print_type(true); // Enable general type printing
-                llvm::set_print_type_fun(&fun); // Set specific function to analyze
+                enzyme.lock().unwrap().set_print_type(true); // Enable general type printing
+                enzyme.lock().unwrap().set_print_type_fun(&fun); // Set specific function to analyze
             }
             config::AutoDiff::Inline => {
-                llvm::set_inline(true);
+                enzyme.lock().unwrap().set_inline(true);
             }
             config::AutoDiff::LooseTypes => {
-                llvm::set_loose_types(true);
+                enzyme.lock().unwrap().set_loose_types(true);
             }
             config::AutoDiff::PrintSteps => {
-                llvm::set_print(true);
+                enzyme.lock().unwrap().set_print(true);
             }
             // We handle this in the PassWrapper.cpp
             config::AutoDiff::PrintPasses => {}
@@ -563,9 +565,9 @@ fn enable_autodiff_settings(ad: &[config::AutoDiff]) {
         }
     }
     // This helps with handling enums for now.
-    llvm::set_strict_aliasing(false);
+    enzyme.lock().unwrap().set_strict_aliasing(false);
     // FIXME(ZuseZ4): Test this, since it was added a long time ago.
-    llvm::set_rust_rules(true);
+    enzyme.lock().unwrap().set_rust_rules(true);
 }
 
 pub(crate) fn run_pass_manager(

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -665,21 +665,17 @@ pub(crate) unsafe fn llvm_optimize(
 
     let llvm_plugins = config.llvm_plugins.join(",");
 
-    fn call_dynamic() -> Result<*const c_void, Box<dyn std::error::Error>> {
-        unsafe {
-            let lib = libloading::Library::new("/home/manuel/prog/rust/build/x86_64-unknown-linux-gnu/enzyme/lib/libEnzyme-21.so")?;
-            let func: libloading::Symbol<'_, c_void> = lib.get(b"registerEnzymeAndPassPipeline")?;
-            let func = func.try_as_raw_ptr().unwrap();
-            dbg!(func);
-            Ok(func as *const c_void)
-        }
-    }
-    let enzyme_fn = if cfg!(llvm_enzyme) && run_enzyme {
-        call_dynamic().unwrap_or(std::ptr::null())
+    let enzyme_fn = if run_enzyme {
+        let wrapper = crate::llvm::EnzymeWrapper::current();
+        wrapper.lock().unwrap().registerEnzymeAndPassPipeline
     } else {
-       std::ptr::null()
+        //dbg!(run_enzyme);
+        //dbg!(consider_ad);
+        std::ptr::null()
     };
 
+    dbg!(&enzyme_fn);
+
 
     let result = unsafe {
         llvm::LLVMRustOptimize(

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -51,6 +51,9 @@ unsafe extern "C" {
     pub(crate) fn LLVMGetNamedFunction(M: &Module, Name: *const c_char) -> Option<&Value>;
 }
 
+
+
+
 #[repr(C)]
 #[derive(Copy, Clone, PartialEq)]
 pub(crate) enum LLVMRustVerifierFailureAction {
@@ -59,113 +62,201 @@ pub(crate) enum LLVMRustVerifierFailureAction {
     LLVMReturnStatusAction = 2,
 }
 
-#[cfg(not(llvm_enzyme))]
+#[cfg(llvm_enzyme)]
 pub(crate) use self::Enzyme_AD::*;
 
-#[cfg(not(llvm_enzyme))]
+//#[cfg(llvm_enzyme)]
 pub(crate) mod Enzyme_AD {
-    use std::ffi::{CString, c_char};
+    use std::ffi::CString;
+    //use std::ffi::{CString, c_char};
 
     use libc::c_void;
 
-    unsafe extern "C" {
-        pub(crate) fn EnzymeSetCLBool(arg1: *mut ::std::os::raw::c_void, arg2: u8);
-        pub(crate) fn EnzymeSetCLString(arg1: *mut ::std::os::raw::c_void, arg2: *const c_char);
-    }
-    unsafe extern "C" {
-        static mut EnzymePrintPerf: c_void;
-        static mut EnzymePrintActivity: c_void;
-        static mut EnzymePrintType: c_void;
-        static mut EnzymeFunctionToAnalyze: c_void;
-        static mut EnzymePrint: c_void;
-        static mut EnzymeStrictAliasing: c_void;
-        static mut looseTypeAnalysis: c_void;
-        static mut EnzymeInline: c_void;
-        static mut RustTypeRules: c_void;
+type SetFlag = unsafe extern "C" fn(*mut c_void, u8);
+
+#[derive(Debug)]
+pub(crate) struct EnzymeFns {
+    pub set_cl: SetFlag,
+}
+
+    #[derive(Debug)]
+    pub(crate) struct EnzymeWrapper {
+        EnzymePrintPerf: *mut c_void,
+        EnzymePrintActivity: *mut c_void,
+        EnzymePrintType: *mut c_void,
+        EnzymeFunctionToAnalyze: *mut c_void,
+        EnzymePrint: *mut c_void,
+        EnzymeStrictAliasing: *mut c_void,
+        looseTypeAnalysis: *mut c_void,
+        EnzymeInline: *mut c_void,
+        RustTypeRules: *mut c_void,
+
+        EnzymeSetCLBool: EnzymeFns,
+        EnzymeSetCLString: EnzymeFns,
+        pub registerEnzymeAndPassPipeline: *const c_void,
     }
-    pub(crate) fn set_print_perf(print: bool) {
-        unsafe {
-            EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintPerf), print as u8);
+        fn call_dynamic() -> Result<EnzymeWrapper, Box<dyn std::error::Error>> {
+            fn load_ptr(lib: &libloading::Library, bytes: &[u8]) -> Result<*mut c_void, Box<dyn std::error::Error>> {
+                // Safety: symbol lookup from a loaded shared object.
+                unsafe {
+                    let s: libloading::Symbol<'_, *mut c_void> = lib.get(bytes)?;
+                    let s = s.try_as_raw_ptr().unwrap();
+                    Ok(s as *mut c_void)
+                }
+            }
+            dbg!("starting");
+            dbg!("Loading Enzyme");
+            let lib = unsafe {libloading::Library::new("/home/manuel/prog/rust/build/x86_64-unknown-linux-gnu/enzyme/lib/libEnzyme-21.so")?};
+            dbg!("second");
+            let EnzymeSetCLBool: libloading::Symbol<'_, SetFlag> = unsafe{lib.get(b"EnzymeSetCLBool")?};
+            dbg!("third");
+            let registerEnzymeAndPassPipeline =
+                load_ptr(&lib, b"registerEnzymeAndPassPipeline").unwrap() as *const c_void;
+            dbg!("fourth");
+            //let EnzymeSetCLBool: libloading::Symbol<'_, unsafe extern "C" fn(&mut c_void, u8) -> ()> = unsafe{lib.get(b"registerEnzymeAndPassPipeline")?};
+            //let EnzymeSetCLBool = unsafe {EnzymeSetCLBool.try_as_raw_ptr().unwrap()};
+            let EnzymeSetCLString: libloading::Symbol<'_, SetFlag> = unsafe{ lib.get(b"EnzymeSetCLString")?};
+            dbg!("done");
+            //let EnzymeSetCLString = unsafe {EnzymeSetCLString.try_as_raw_ptr().unwrap()};
+
+            let EnzymePrintPerf = load_ptr(&lib, b"EnzymePrintPerf").unwrap();
+            let EnzymePrintActivity = load_ptr(&lib, b"EnzymePrintActivity").unwrap();
+            let EnzymePrintType = load_ptr(&lib, b"EnzymePrintType").unwrap();
+            let EnzymeFunctionToAnalyze = load_ptr(&lib, b"EnzymeFunctionToAnalyze").unwrap();
+            let EnzymePrint = load_ptr(&lib, b"EnzymePrint").unwrap();
+
+            let EnzymeStrictAliasing = load_ptr(&lib, b"EnzymeStrictAliasing").unwrap();
+            let looseTypeAnalysis = load_ptr(&lib, b"looseTypeAnalysis").unwrap();
+            let EnzymeInline = load_ptr(&lib, b"EnzymeInline").unwrap();
+            let RustTypeRules = load_ptr(&lib, b"RustTypeRules").unwrap();
+
+            let wrap = EnzymeWrapper {
+                EnzymePrintPerf,
+                EnzymePrintActivity,
+                EnzymePrintType,
+                EnzymeFunctionToAnalyze,
+                EnzymePrint,
+                EnzymeStrictAliasing,
+                looseTypeAnalysis,
+                EnzymeInline,
+                RustTypeRules,
+                //EnzymeSetCLBool: EnzymeFns {set_cl: unsafe{*EnzymeSetCLBool}},
+                //EnzymeSetCLString: EnzymeFns {set_cl: unsafe{*EnzymeSetCLString}},
+                EnzymeSetCLBool: EnzymeFns {set_cl: *EnzymeSetCLBool},
+                EnzymeSetCLString: EnzymeFns {set_cl: *EnzymeSetCLString},
+                registerEnzymeAndPassPipeline,
+            };
+            dbg!(&wrap);
+            Ok(wrap)
         }
-    }
-    pub(crate) fn set_print_activity(print: bool) {
-        unsafe {
-            EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintActivity), print as u8);
+use std::sync::Mutex;
+unsafe impl Sync for EnzymeWrapper {}
+unsafe impl Send for EnzymeWrapper {}
+    impl EnzymeWrapper {
+        pub(crate) fn current() -> &'static Mutex<EnzymeWrapper> {
+            use std::sync::OnceLock;
+            static CELL: OnceLock<Mutex<EnzymeWrapper>> = OnceLock::new();
+            fn init_enzyme() -> Mutex<EnzymeWrapper> {
+                call_dynamic().unwrap().into()
+            }
+            CELL.get_or_init(|| init_enzyme())
         }
-    }
-    pub(crate) fn set_print_type(print: bool) {
-        unsafe {
-            EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintType), print as u8);
+        pub(crate) fn set_print_perf(&mut self, print: bool) {
+            unsafe {
+                //(self.EnzymeSetCLBool.set_cl)(self.EnzymePrintPerf, print as u8);
+                //(self.EnzymeSetCLBool)(std::ptr::addr_of_mut!(self.EnzymePrintPerf), print as u8);
+            }
         }
-    }
-    pub(crate) fn set_print_type_fun(fun_name: &str) {
-        let c_fun_name = CString::new(fun_name).unwrap();
-        unsafe {
-            EnzymeSetCLString(
-                std::ptr::addr_of_mut!(EnzymeFunctionToAnalyze),
-                c_fun_name.as_ptr() as *const c_char,
-            );
+
+        pub(crate) fn set_print_activity(&mut self, print: bool) {
+            unsafe {
+                //(self.EnzymeSetCLBool.set_cl)(self.EnzymePrintActivity, print as u8);
+                //(self.EnzymeSetCLBool)(std::ptr::addr_of_mut!(self.EnzymePrintActivity), print as u8);
+            }
         }
-    }
-    pub(crate) fn set_print(print: bool) {
-        unsafe {
-            EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrint), print as u8);
+
+        pub(crate) fn set_print_type(&mut self, print: bool) {
+            unsafe {
+               // (self.EnzymeSetCLBool.set_cl)(self.EnzymePrintType, print as u8);
+            }
         }
-    }
-    pub(crate) fn set_strict_aliasing(strict: bool) {
-        unsafe {
-            EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymeStrictAliasing), strict as u8);
+
+        pub(crate) fn set_print_type_fun(&mut self, fun_name: &str) {
+            let _c_fun_name = CString::new(fun_name).unwrap();
+            //unsafe {
+            //    (self.EnzymeSetCLString.set_cl)(
+            //        self.EnzymeFunctionToAnalyze,
+            //        c_fun_name.as_ptr() as *const c_char,
+            //    );
+            //}
         }
-    }
-    pub(crate) fn set_loose_types(loose: bool) {
-        unsafe {
-            EnzymeSetCLBool(std::ptr::addr_of_mut!(looseTypeAnalysis), loose as u8);
+
+        pub(crate) fn set_print(&mut self, print: bool) {
+            unsafe {
+                //(self.EnzymeSetCLBool.set_cl)(self.EnzymePrint, print as u8);
+            }
         }
-    }
-    pub(crate) fn set_inline(val: bool) {
-        unsafe {
-            EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymeInline), val as u8);
+
+        pub(crate) fn set_strict_aliasing(&mut self, strict: bool) {
+            unsafe {
+                //(self.EnzymeSetCLBool.set_cl)(self.EnzymeStrictAliasing, strict as u8);
+            }
         }
-    }
-    pub(crate) fn set_rust_rules(val: bool) {
-        unsafe {
-            EnzymeSetCLBool(std::ptr::addr_of_mut!(RustTypeRules), val as u8);
+
+        pub(crate) fn set_loose_types(&mut self, loose: bool) {
+            unsafe {
+                //(self.EnzymeSetCLBool.set_cl)(self.looseTypeAnalysis, loose as u8);
+            }
+        }
+
+        pub(crate) fn set_inline(&mut self, val: bool) {
+            unsafe {
+                //(self.EnzymeSetCLBool.set_cl)(self.EnzymeInline, val as u8);
+            }
+        }
+
+        pub(crate) fn set_rust_rules(&mut self, val: bool) {
+            unsafe {
+                //(self.EnzymeSetCLBool.set_cl)(self.RustTypeRules, val as u8);
+            }
         }
     }
+
+
 }
 
-#[cfg(llvm_enzyme)]
+#[cfg(not(llvm_enzyme))]
 pub(crate) use self::Fallback_AD::*;
 
-#[cfg(llvm_enzyme)]
+#[cfg(not(llvm_enzyme))]
 pub(crate) mod Fallback_AD {
     #![allow(unused_variables)]
 
     pub(crate) fn set_inline(val: bool) {
-        //unimplemented!()
+        unimplemented!()
     }
     pub(crate) fn set_print_perf(print: bool) {
-        //unimplemented!()
+        unimplemented!()
     }
     pub(crate) fn set_print_activity(print: bool) {
-        //unimplemented!()
+        unimplemented!()
     }
     pub(crate) fn set_print_type(print: bool) {
-        //unimplemented!()
+        unimplemented!()
     }
     pub(crate) fn set_print_type_fun(fun_name: &str) {
-        //unimplemented!()
+        unimplemented!()
     }
     pub(crate) fn set_print(print: bool) {
-        //unimplemented!()
+        unimplemented!()
     }
     pub(crate) fn set_strict_aliasing(strict: bool) {
-        //unimplemented!()
+        unimplemented!()
     }
     pub(crate) fn set_loose_types(loose: bool) {
-        //unimplemented!()
+        unimplemented!()
     }
     pub(crate) fn set_rust_rules(val: bool) {
-        //unimplemented!()
+        unimplemented!()
     }
 }