diff --git a/compiler/rustc_codegen_llvm/src/back/write.rs b/compiler/rustc_codegen_llvm/src/back/write.rs index 5e0eda30d68a2..923db9952dc58 100644 --- 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, diff_items: Vec, _typetrees: FxHashMap, - 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::().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::().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 index 173d58ebca2af..5a01478e0b406 100644 --- 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 index d4dc34594f31c..8dc4c3af667fb 100644 --- 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