Merge pull request #49 from EnzymeAD/more-helper

More helper
EnzymeAD · Sep 20, 2023 · f70a33d · f70a33d
2 parents 6becc11 + 731bbd1
commit f70a33d
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 7 deletions.
diff --git a/compiler/rustc_codegen_llvm/src/back/write.rs b/compiler/rustc_codegen_llvm/src/back/write.rs
@@ -673,12 +673,16 @@ pub(crate) unsafe fn enzyme_ad(
         CreateTypeAnalysis(logic_ref, std::ptr::null_mut(), std::ptr::null_mut(), 0);
 
     llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymeStrictAliasing), 0);
+
     if std::env::var("ENZYME_PRINT_TA").is_ok() {
       llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymePrintType), 1);
     }
     if std::env::var("ENZYME_PRINT_AA").is_ok() {
         llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymePrintActivity), 1);
     }
+    if std::env::var("ENZYME_PRINT_PERF").is_ok() {
+      llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymePrintPerf), 1);
+    }
     if std::env::var("ENZYME_PRINT").is_ok() {
         llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymePrint), 1);
     }
@@ -733,22 +737,29 @@ pub(crate) unsafe fn differentiate(
     _cgcx: &CodegenContext<LlvmCodegenBackend>,
     diff_items: Vec<AutoDiffItem>,
     _typetrees: FxHashMap<String, DiffTypeTree>,
-    config: &ModuleConfig,
+    _config: &ModuleConfig,
 ) -> Result<(), FatalError> {
     let llmod = module.module_llvm.llmod();
     let llcx = &module.module_llvm.llcx;
 
-    if config.enzyme_print_activity {
-        llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymePrintActivity), 1);
-        llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymePrintType), 1);
-        llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymePrint), 1);
-    }
-
     llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymeStrictAliasing), 0);
 
     if std::env::var("ENZYME_PRINT_MOD").is_ok() {
         unsafe {LLVMDumpModule(llmod);}
     }
+    if std::env::var("ENZYME_TT_DEPTH").is_ok() {
+        let depth = std::env::var("ENZYME_TT_DEPTH").unwrap();
+        let depth = depth.parse::<i64>().unwrap();
+        assert!(depth >= 1);
+        llvm::EnzymeSetCLInteger(std::ptr::addr_of_mut!(llvm::EnzymeMaxTypeDepth), depth);
+    }
+    if std::env::var("ENZYME_TT_WIDTH").is_ok() {
+        let width = std::env::var("ENZYME_TT_WIDTH").unwrap();
+        let width = width.parse::<i64>().unwrap();
+        assert!(width >= 1);
+        llvm::EnzymeSetCLInteger(std::ptr::addr_of_mut!(llvm::MaxTypeOffset), width);
+    }
+
     for item in diff_items {
         let res = enzyme_ad(llmod, llcx, item);
         assert!(res.is_ok());

diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -2766,6 +2766,11 @@ extern "C" {
 }
 
 extern "C" {
+    pub static mut MaxIntOffset: c_void;
+    pub static mut MaxTypeOffset: c_void;
+    pub static mut EnzymeMaxTypeDepth: c_void;
+
+    pub static mut EnzymePrintPerf: c_void;
     pub static mut EnzymePrintActivity: c_void;
     pub static mut EnzymePrintType: c_void;
     pub static mut EnzymePrint: c_void;

diff --git a/src/bootstrap/llvm.rs b/src/bootstrap/llvm.rs
@@ -806,6 +806,9 @@ impl Step for Enzyme {
 
         builder.update_submodule(&Path::new("src").join("tools").join("enzyme"));
         let mut cfg = cmake::Config::new(builder.src.join("src/tools/enzyme/enzyme/"));
+        // TODO: Find a nicer way to use Enzyme Debug builds
+        //cfg.profile("Debug");
+        //cfg.define("CMAKE_BUILD_TYPE", "Debug");
         configure_cmake(builder, target, &mut cfg, true, LdFlags::default(), &[]);
 
         // Re-use the same flags as llvm to control the level of debug information