re-add flags after refactoring

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -586,6 +586,29 @@ fn thin_lto(
     }
 }
 
+fn enable_autodiff_settings(ad: &[config::AutoDiff], module: &mut ModuleCodegen<ModuleLlvm>) {
+
+    for &val in ad {
+        match val {
+            config::AutoDiff::PrintModBefore => { unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) }; },
+            config::AutoDiff::PrintPerf => {llvm::set_print_perf(true);},
+            config::AutoDiff::PrintAA => {llvm::set_print_activity(true);},
+            config::AutoDiff::PrintTA => {llvm::set_print_type(true);},
+            config::AutoDiff::Inline => {llvm::set_inline(true);},
+            config::AutoDiff::LooseTypes => {llvm::set_loose_types(false);},
+            config::AutoDiff::PrintSteps => {llvm::set_print(true);},
+            // We handle this below
+            config::AutoDiff::PrintModAfter => {}
+            // This is required and already checked
+            config::AutoDiff::Enable => {}
+        }
+    }
+    // This helps with handling enums for now.
+    llvm::set_strict_aliasing(false);
+    // FIXME(ZuseZ4): Test this, since it was added a long time ago.
+    llvm::set_rust_rules(true);
+}
+
 pub(crate) fn run_pass_manager(
     cgcx: &CodegenContext<LlvmCodegenBackend>,
     dcx: DiagCtxtHandle<'_>,
@@ -604,10 +627,6 @@ pub(crate) fn run_pass_manager(
     let opt_stage = if thin { llvm::OptStage::ThinLTO } else { llvm::OptStage::FatLTO };
     let opt_level = config.opt_level.unwrap_or(config::OptLevel::No);
 
-    // If this rustc version was build with enzyme/autodiff enabled, and if users applied the
-    // `#[autodiff]` macro at least once, then we will later call llvm_optimize a second time.
-    debug!("running llvm pm opt pipeline");
-
     // The PostAD behavior is the same that we would have if no autodiff was used.
     // It will run the default optimization pipeline. If AD is enabled we select
     // the DuringAD stage, which will disable vectorization and loop unrolling, and
@@ -618,25 +637,27 @@ pub(crate) fn run_pass_manager(
     let stage =
         if enable_ad { write::AutodiffStage::DuringAD } else { write::AutodiffStage::PostAD };
 
-    // If Enzyme fails to differentiate the code, then this module will have everything that's
-    // needed to reproduce the bug.
-    if config.autodiff.contains(&config::AutoDiff::PrintModBefore) {
-        unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
+    if enable_ad {
+        enable_autodiff_settings(&config.autodiff, module);
     }
+
     unsafe {
         write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, stage)?;
     }
+
     if cfg!(llvm_enzyme) && enable_ad {
         let opt_stage = llvm::OptStage::FatLTO;
         let stage = write::AutodiffStage::PostAD;
         unsafe {
             write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, stage)?;
         }
+
+        // This is the final IR, so people should be able to inspect the optimized autodiff output.
+        if config.autodiff.contains(&config::AutoDiff::PrintModAfter) {
+            unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
+        }
     }
-    // This is the final IR, so people should be able to inspect the optimized autodiff output.
-    if config.autodiff.contains(&config::AutoDiff::PrintModAfter) {
-        unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
-    }
+
     debug!("lto done");
     Ok(())
 }

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -1,7 +1,7 @@
 #![allow(non_camel_case_types)]
 #![expect(dead_code)]
 
-use libc::{c_char, c_uint};
+use libc::{c_char, c_uint, c_void};
 
 use super::ffi::{BasicBlock, Metadata, Module, Type, Value};
 use crate::llvm::Bool;
@@ -35,3 +35,71 @@ pub enum LLVMRustVerifierFailureAction {
     LLVMPrintMessageAction = 1,
     LLVMReturnStatusAction = 2,
 }
+
+
+#[cfg(not(llvm_enzyme))]
+pub fn EnzymeSetCLBool(arg1: *mut ::std::os::raw::c_void, arg2: u8) {
+    unimplemented!()
+}
+#[cfg(not(llvm_enzyme))]
+pub fn EnzymeSetCLInteger(arg1: *mut ::std::os::raw::c_void, arg2: i64) {
+    unimplemented!()
+}
+
+#[cfg(llvm_enzyme)]
+extern "C" {
+    pub fn EnzymeSetCLBool(arg1: *mut ::std::os::raw::c_void, arg2: u8);
+    pub fn EnzymeSetCLInteger(arg1: *mut ::std::os::raw::c_void, arg2: i64);
+}
+
+#[cfg(llvm_enzyme)]
+extern "C" {
+    static mut EnzymePrintPerf: c_void;
+    static mut EnzymePrintActivity: c_void;
+    static mut EnzymePrintType: 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;
+}
+pub fn set_print_perf(print: bool) {
+    unsafe {
+        EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintPerf), print as u8);
+    }
+}
+pub fn set_print_activity(print: bool) {
+    unsafe {
+        EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintActivity), print as u8);
+    }
+}
+pub fn set_print_type(print: bool) {
+    unsafe {
+        EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintType), print as u8);
+    }
+}
+pub fn set_print(print: bool) {
+    unsafe {
+        EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrint), print as u8);
+    }
+}
+pub fn set_strict_aliasing(strict: bool) {
+    unsafe {
+        EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymeStrictAliasing), strict as u8);
+    }
+}
+pub fn set_loose_types(loose: bool) {
+    unsafe {
+        EnzymeSetCLBool(std::ptr::addr_of_mut!(looseTypeAnalysis), loose as u8);
+    }
+}
+pub fn set_inline(val: bool) {
+    unsafe {
+        EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymeInline), val as u8);
+    }
+}
+pub fn set_rust_rules(val: bool) {
+    unsafe {
+        EnzymeSetCLBool(std::ptr::addr_of_mut!(RustTypeRules), val as u8);
+    }
+}

compiler/rustc_session/src/config.rs

@@ -217,8 +217,6 @@ pub enum AutoDiff {
     /// Enzyme's loose type debug helper (can cause incorrect gradients!!)
     /// Usable in cases where Enzyme errors with `can not deduce type of X`.
     LooseTypes,
-    /// See Enzyme core documentation. FIXME(ZuseZ4): Clarify usages
-    RuntimeActivity,
     /// Runs Enzyme's aggressive inlining
     Inline,
 }

compiler/rustc_session/src/options.rs

@@ -707,7 +707,7 @@ mod desc {
     pub(crate) const parse_list: &str = "a space-separated list of strings";
     pub(crate) const parse_list_with_polarity: &str =
         "a comma-separated list of strings, with elements beginning with + or -";
-    pub(crate) const parse_autodiff: &str = "a comma separated list of settings: `Enable`, `PrintSteps`, `PrintTA`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfter`, `LooseTypes`, `RuntimeActivity`, `Inline`";
+    pub(crate) const parse_autodiff: &str = "a comma separated list of settings: `Enable`, `PrintSteps`, `PrintTA`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfter`, `LooseTypes`, `Inline`";
     pub(crate) const parse_comma_list: &str = "a comma-separated list of strings";
     pub(crate) const parse_opt_comma_list: &str = parse_comma_list;
     pub(crate) const parse_number: &str = "a number";
@@ -1356,7 +1356,6 @@ pub mod parse {
                 "PrintModBefore" => AutoDiff::PrintModBefore,
                 "PrintModAfter" => AutoDiff::PrintModAfter,
                 "LooseTypes" => AutoDiff::LooseTypes,
-                "RuntimeActivity" => AutoDiff::RuntimeActivity,
                 "Inline" => AutoDiff::Inline,
                 _ => {
                     // FIXME(ZuseZ4): print an error saying which value is not recognized
@@ -2090,7 +2089,6 @@ options! {
         `=PrintModBefore`
         `=PrintModAfter`
         `=LooseTypes`
-        '=RuntimeActivity`
         `=Inline`
         Multiple options can be combined with commas."),
     #[rustc_lint_opt_deny_field_access("use `Session::binary_dep_depinfo` instead of this field")]