autodiff: typetree recursive depth query from enzyme with fallback

Signed-off-by: Karan Janthe <[email protected]>

Author: KMJ-007

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -127,7 +127,6 @@ pub(crate) mod Enzyme_AD {
         );
         pub(crate) fn EnzymeTypeTreeToString(arg1: CTypeTreeRef) -> *const c_char;
         pub(crate) fn EnzymeTypeTreeToStringFree(arg1: *const c_char);
-        pub(crate) fn EnzymeGetMaxTypeDepth() -> ::std::os::raw::c_uint;
     }
 
     unsafe extern "C" {

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -1847,3 +1847,15 @@ extern "C" void LLVMRustSetNoSanitizeHWAddress(LLVMValueRef Global) {
   MD.NoHWAddress = true;
   GV.setSanitizerMetadata(MD);
 }
+
+#ifdef ENZYME
+extern "C" {
+extern llvm::cl::opt<unsigned> EnzymeMaxTypeDepth;
+}
+
+extern "C" size_t LLVMRustEnzymeGetMaxTypeDepth() { return EnzymeMaxTypeDepth; }
+#else
+extern "C" size_t LLVMRustEnzymeGetMaxTypeDepth() {
+  return 6; // Default fallback depth
+}
+#endif

compiler/rustc_middle/src/ty/mod.rs

@@ -2266,16 +2266,16 @@ fn typetree_from_ty_inner<'tcx>(
     #[cfg(llvm_enzyme)]
     {
         unsafe extern "C" {
-            fn EnzymeGetMaxTypeDepth() -> ::std::os::raw::c_uint;
+            fn LLVMRustEnzymeGetMaxTypeDepth() -> usize;
         }
-        let max_depth = unsafe { EnzymeGetMaxTypeDepth() } as usize;
-        if depth > max_depth {
+        let max_depth = unsafe { LLVMRustEnzymeGetMaxTypeDepth() };
+        if depth >= max_depth {
             return TypeTree::new();
         }
     }
 
     #[cfg(not(llvm_enzyme))]
-    if depth > 6 {
+    if depth >= 6 {
         return TypeTree::new();
     }