feat: dlopen Enzyme

Cargo.lock

@@ -3613,6 +3613,7 @@ dependencies = [
  "gimli 0.31.1",
  "itertools",
  "libc",
+ "libloading 0.9.0",
  "measureme",
  "object 0.37.3",
  "rustc-demangle",

compiler/rustc_codegen_llvm/Cargo.toml

@@ -14,6 +14,7 @@ bitflags = "2.4.1"
 gimli = "0.31"
 itertools = "0.12"
 libc = "0.2"
+libloading = { version = "0.9.0", optional = true }
 measureme = "12.0.1"
 object = { version = "0.37.0", default-features = false, features = ["std", "read"] }
 rustc-demangle = "0.1.21"
@@ -46,7 +47,6 @@ tracing = "0.1"
 [features]
 # tidy-alphabetical-start
 check_only = ["rustc_llvm/check_only"]
-llvm_enzyme = []
+llvm_enzyme = ["dep:libloading"]
 llvm_offload = []
 # tidy-alphabetical-end
-

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -528,31 +528,37 @@ fn thin_lto(
  }
 }
 
-fn enable_autodiff_settings(ad: &[config::AutoDiff]) {
+#[cfg(feature = "llvm_enzyme")]
+pub(crate) fn enable_autodiff_settings(
+ sysroot: &rustc_session::config::Sysroot,
+ ad: &[config::AutoDiff],
+) {
+ 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.
  match val {
  config::AutoDiff::PrintPerf => {
- llvm::set_print_perf(true);
+ enzyme.set_print_perf(true);
  }
  config::AutoDiff::PrintAA => {
- llvm::set_print_activity(true);
+ enzyme.set_print_activity(true);
  }
  config::AutoDiff::PrintTA => {
- llvm::set_print_type(true);
+ enzyme.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.set_print_type(true); // Enable general type printing
+ enzyme.set_print_type_fun(&fun); // Set specific function to analyze
  }
  config::AutoDiff::Inline => {
- llvm::set_inline(true);
+ enzyme.set_inline(true);
  }
  config::AutoDiff::LooseTypes => {
- llvm::set_loose_types(true);
+ enzyme.set_loose_types(true);
  }
  config::AutoDiff::PrintSteps => {
- llvm::set_print(true);
+ enzyme.set_print(true);
  }
  // We handle this in the PassWrapper.cpp
  config::AutoDiff::PrintPasses => {}
@@ -571,9 +577,9 @@ fn enable_autodiff_settings(ad: &[config::AutoDiff]) {
  }
  }
  // This helps with handling enums for now.
- llvm::set_strict_aliasing(false);
+ enzyme.set_strict_aliasing(false);
  // FIXME(ZuseZ4): Test this, since it was added a long time ago.
- llvm::set_rust_rules(true);
+ enzyme.set_rust_rules(true);
 }
 
 pub(crate) fn run_pass_manager(
@@ -607,10 +613,6 @@ pub(crate) fn run_pass_manager(
  if enable_ad { write::AutodiffStage::DuringAD } else { write::AutodiffStage::PostAD }
  };
 
- if enable_ad {
- enable_autodiff_settings(&config.autodiff);
- }
-
  unsafe {
  write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage);
  }

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -730,6 +730,13 @@ pub(crate) unsafe fn llvm_optimize(
 
  let llvm_plugins = config.llvm_plugins.join(",");
 
+ let enzyme_fn = if consider_ad {
+ let wrapper = llvm::EnzymeWrapper::get_or_init(&cgcx.sysroot);
+ wrapper.registerEnzymeAndPassPipeline
+ } else {
+ std::ptr::null()
+ };
+
  let result = unsafe {
  llvm::LLVMRustOptimize(
  module.module_llvm.llmod(),
@@ -749,7 +756,7 @@ pub(crate) unsafe fn llvm_optimize(
  vectorize_loop,
  config.no_builtins,
  config.emit_lifetime_markers,
- run_enzyme,
+ enzyme_fn,
  print_before_enzyme,
  print_after_enzyme,
  print_passes,

compiler/rustc_codegen_llvm/src/lib.rs

@@ -240,6 +240,18 @@ impl CodegenBackend for LlvmCodegenBackend {
 
  fn init(&self, sess: &Session) {
  llvm_util::init(sess); // Make sure llvm is inited
+
+ #[cfg(feature = "llvm_enzyme")]
+ {
+ use rustc_session::config::AutoDiff;
+ 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);
+ }
+ }
+ }
  }
 
  fn provide(&self, providers: &mut Providers) {