Use libloading to open Enzyme at runtime

Author: ZuseZ4

compiler/rustc_codegen_llvm/Cargo.toml

@@ -14,6 +14,7 @@ bitflags = "2.4.1"
 gimli = "0.31"
 itertools = "0.12"
 libc = "0.2"
+libloading = "0.8.8"
 measureme = "12.0.1"
 object = { version = "0.37.0", default-features = false, features = ["std", "read"] }
 rustc-demangle = "0.1.21"

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -522,31 +522,33 @@ fn thin_lto(
     }
 }
 
-fn enable_autodiff_settings(ad: &[config::AutoDiff]) {
+fn enable_autodiff_settings(cgcx: &CodegenContext<LlvmCodegenBackend>, ad: &[config::AutoDiff]) {
+    use rustc_codegen_ssa::back::write::EnzymeWrapper;
+    let enzyme = rustc_codegen_ssa::back::write::EnzymeWrapper::current(cgcx);
     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(
@@ -601,7 +603,7 @@ pub(crate) fn run_pass_manager(
     };
 
     if enable_ad {
-        enable_autodiff_settings(&config.autodiff);
+        enable_autodiff_settings(&cgcx, &config.autodiff);
     }
 
     unsafe {

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -665,6 +665,18 @@ pub(crate) unsafe fn llvm_optimize(
 
     let llvm_plugins = config.llvm_plugins.join(",");
 
+    let enzyme_fn = if consider_ad {
+        let wrapper = rustc_codegen_ssa::back::write::EnzymeWrapper::current(cgcx);
+        wrapper.lock().unwrap().registerEnzymeAndPassPipeline
+    } else {
+        //dbg!(run_enzyme);
+        //dbg!(consider_ad);
+        std::ptr::null()
+    };
+
+    dbg!(&enzyme_fn);
+
+
     let result = unsafe {
         llvm::LLVMRustOptimize(
             module.module_llvm.llmod(),
@@ -684,7 +696,8 @@ pub(crate) unsafe fn llvm_optimize(
             vectorize_loop,
             config.no_builtins,
             config.emit_lifetime_markers,
-            run_enzyme,
+            enzyme_fn,
+            //run_enzyme,
             print_before_enzyme,
             print_after_enzyme,
             print_passes,