Merge pull request #40 from EnzymeAD/nicer-debug

improve debugging

Author: ZuseZ4

README.md

@@ -35,9 +35,13 @@ cargo +enzyme test --examples
 ```
 
 ## Enzyme Config
-Instead of having to re-build LLVM and Enzyme every time, we can adjust the 
-Enzyme behavior on some higher level inside the following file or using the config parameter
-[compiler/rustc_codegen_llvm/src/back/write.rs](https://github.com/EnzymeAD/rust/blob/c7fbb43f83a224cfa290dde2a46a930dce0c5507/compiler/rustc_codegen_llvm/src/back/write.rs#L712)
+To help with debugging, Enzyme can be configured using three environment variables.
+```bash
+export ENZYME_PRINT_TA
+export ENZYME_PRINT_AA
+export ENZYME_PRINT
+```
+These will print TypeAnalysis, ActivityAnalysis and the primal + generated function respectively.
 
 We want to make sure that we have EnzymeStrictAliasing=0 all the time to make sure
 that Enzyme handles enums correctly. Other flags might be enabled for debug purpose.

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -671,9 +671,17 @@ pub(crate) unsafe fn enzyme_ad(
     let logic_ref: EnzymeLogicRef = CreateEnzymeLogic(opt as u8);
     let type_analysis: EnzymeTypeAnalysisRef =
         CreateTypeAnalysis(logic_ref, std::ptr::null_mut(), std::ptr::null_mut(), 0);
-    
+
     llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymeStrictAliasing), 0);
-    llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymePrintType), 1);
+    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").is_ok() {
+        llvm::EnzymeSetCLBool(std::ptr::addr_of_mut!(llvm::EnzymePrint), 1);
+    }
 
     let mut res: &Value = match item.attrs.mode {
         DiffMode::Forward => enzyme_rust_forward_diff(
@@ -742,7 +750,7 @@ pub(crate) unsafe fn differentiate(
         let res = enzyme_ad(llmod, llcx, item);
         assert!(res.is_ok());
     }
-    
+
     let mut f = LLVMGetFirstFunction(llmod);
     loop {
         if let Some(lf) = f {