autodiff: recursive depth limit added in typetree

Author: KMJ-007

compiler/rustc_middle/src/ty/mod.rs

@@ -2257,7 +2257,25 @@ pub fn fnc_typetrees<'tcx>(tcx: TyCtxt<'tcx>, fn_ty: Ty<'tcx>) -> FncTree {
 
 /// Generate TypeTree for a specific type.
 /// This function analyzes a Rust type and creates appropriate TypeTree metadata.
+
+/// Maximum recursion depth for TypeTree generation to prevent infinite loops
+/// Set to 32 levels which should be sufficient for most practical type hierarchies
+/// while preventing stack overflow from pathological recursive types.
+const MAX_TYPETREE_DEPTH: usize = 32;
+
+fn typetree_from_ty_with_depth<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>, depth: usize) -> TypeTree {
+    if depth > MAX_TYPETREE_DEPTH {
+        return TypeTree::new();
+    }
+
+    typetree_from_ty_impl(tcx, ty, depth)
+}
+
 pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
+    typetree_from_ty_with_depth(tcx, ty, 0)
+}
+
+fn typetree_from_ty_impl<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>, depth: usize) -> TypeTree {
     if ty.is_scalar() {
         let (kind, size) = if ty.is_integral() || ty.is_char() || ty.is_bool() {
             (Kind::Integer, ty.primitive_size(tcx).bytes_usize())
@@ -2299,7 +2317,7 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
                 return TypeTree::new();
             }
 
-            let element_tree = typetree_from_ty(tcx, *element_ty);
+            let element_tree = typetree_from_ty_impl(tcx, *element_ty, depth + 1);
 
             let element_layout = tcx
                 .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(*element_ty))
@@ -2335,7 +2353,7 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
 
     if ty.is_slice() {
         if let ty::Slice(element_ty) = ty.kind() {
-            let element_tree = typetree_from_ty(tcx, *element_ty);
+            let element_tree = typetree_from_ty_impl(tcx, *element_ty, depth + 1);
             return element_tree;
         }
     }
@@ -2349,7 +2367,7 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
         let mut current_offset = 0;
 
         for tuple_ty in tuple_types.iter() {
-            let element_tree = typetree_from_ty(tcx, tuple_ty);
+            let element_tree = typetree_from_ty_impl(tcx, tuple_ty, depth + 1);
 
             let element_layout = tcx
                 .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(tuple_ty))
@@ -2385,7 +2403,7 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
 
                 for (field_idx, field_def) in adt_def.all_fields().enumerate() {
                     let field_ty = field_def.ty(tcx, args);
-                    let field_tree = typetree_from_ty(tcx, field_ty);
+                    let field_tree = typetree_from_ty_impl(tcx, field_ty, depth + 1);
 
                     let field_offset = layout.fields.offset(field_idx).bytes_usize();
 

tests/run-make/autodiff/type-trees/recursive-typetree/recursive.check

@@ -0,0 +1,4 @@
+; Check that recursive TypeTree generation doesn't infinite loop
+; The depth limit should prevent infinite recursion
+
+CHECK: define{{.*}}"enzyme_type"="{[]:Float@double}"{{.*}}@test_recursive{{.*}}"enzyme_type"="{[]:Pointer}"

tests/run-make/autodiff/type-trees/recursive-typetree/rmake.rs

@@ -0,0 +1,9 @@
+//@ needs-enzyme
+//@ ignore-cross-compile
+
+use run_make_support::{llvm_filecheck, rfs, rustc};
+
+fn main() {
+    rustc().input("test.rs").arg("-Zautodiff=Enable").emit("llvm-ir").run();
+    llvm_filecheck().patterns("recursive.check").stdin_buf(rfs::read("test.ll")).run();
+}

tests/run-make/autodiff/type-trees/recursive-typetree/test.rs

@@ -0,0 +1,21 @@
+#![feature(autodiff)]
+
+use std::autodiff::autodiff_reverse;
+
+#[repr(C)]
+struct Node {
+    value: f64,
+    next: Option<Box<Node>>,
+}
+
+#[autodiff_reverse(d_test, Duplicated, Active)]
+#[no_mangle]
+fn test_recursive(node: &Node) -> f64 {
+    node.value
+}
+
+fn main() {
+    let node = Node { value: 1.0, next: None };
+    let mut d_node = Node { value: 0.0, next: None };
+    let _result = d_test(&node, &mut d_node, 1.0);
+}