autodiff: adding recursion max depth in typetree

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

Author: KMJ-007

compiler/rustc_middle/src/ty/mod.rs

@@ -2257,8 +2257,33 @@ 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 stack overflow
+/// from pathological deeply nested types. Combined with cycle detection.
+const MAX_TYPETREE_DEPTH: usize = 32;
+
 pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
-    if ty.is_scalar() {
+    let mut visited = Vec::new();
+    typetree_from_ty_impl(tcx, ty, 0, &mut visited)
+}
+
+fn typetree_from_ty_impl<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    ty: Ty<'tcx>,
+    depth: usize,
+    visited: &mut Vec<Ty<'tcx>>,
+) -> TypeTree {
+    if depth > MAX_TYPETREE_DEPTH {
+        return TypeTree::new();
+    }
+
+    if visited.contains(&ty) {
+        return TypeTree::new();
+    }
+
+    visited.push(ty);
+
+    let result = 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())
         } else if ty.is_floating_point() {
@@ -2267,116 +2292,118 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
                 x if x == tcx.types.f32 => (Kind::Float, 4),
                 x if x == tcx.types.f64 => (Kind::Double, 8),
                 x if x == tcx.types.f128 => (Kind::F128, 16),
-                _ => return TypeTree::new(),
+                _ => (Kind::Integer, 0),
             }
         } else {
-            return TypeTree::new();
+            (Kind::Integer, 0)
         };
 
-        return TypeTree(vec![Type { offset: -1, size, kind, child: TypeTree::new() }]);
-    }
-
-    if ty.is_ref() || ty.is_raw_ptr() || ty.is_box() {
-        let inner_ty = if let Some(inner) = ty.builtin_deref(true) {
-            inner
+        if size > 0 {
+            TypeTree(vec![Type { offset: -1, size, kind, child: TypeTree::new() }])
         } else {
-            return TypeTree::new();
-        };
-
-        let child = typetree_from_ty(tcx, inner_ty);
-        return TypeTree(vec![Type {
-            offset: -1,
-            size: tcx.data_layout.pointer_size().bytes_usize(),
-            kind: Kind::Pointer,
-            child,
-        }]);
-    }
-
-    if ty.is_array() {
+            TypeTree::new()
+        }
+    } else if ty.is_ref() || ty.is_raw_ptr() || ty.is_box() {
+        if let Some(inner_ty) = ty.builtin_deref(true) {
+            let child = typetree_from_ty_impl(tcx, inner_ty, depth + 1, visited);
+            TypeTree(vec![Type {
+                offset: -1,
+                size: tcx.data_layout.pointer_size().bytes_usize(),
+                kind: Kind::Pointer,
+                child,
+            }])
+        } else {
+            TypeTree::new()
+        }
+    } else if ty.is_array() {
         if let ty::Array(element_ty, len_const) = ty.kind() {
             let len = len_const.try_to_target_usize(tcx).unwrap_or(0);
             if len == 0 {
-                return TypeTree::new();
-            }
-
-            let element_tree = typetree_from_ty(tcx, *element_ty);
-
-            let element_layout = tcx
-                .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(*element_ty))
-                .ok()
-                .map(|layout| layout.size.bytes_usize())
-                .unwrap_or(0);
-
-            if element_layout == 0 {
-                return TypeTree::new();
+                TypeTree::new()
+            } else {
+                let element_tree = typetree_from_ty_impl(tcx, *element_ty, depth + 1, visited);
+                let element_layout = tcx
+                    .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(*element_ty))
+                    .ok()
+                    .map(|layout| layout.size.bytes_usize())
+                    .unwrap_or(0);
+
+                if element_layout == 0 {
+                    TypeTree::new()
+                } else {
+                    // For homogeneous arrays, use offset -1 instead of individual entries
+                    if element_tree.0.len() == 1 && element_tree.0[0].offset == -1 {
+                        TypeTree(vec![Type {
+                            offset: -1,
+                            size: element_tree.0[0].size,
+                            kind: element_tree.0[0].kind,
+                            child: element_tree.0[0].child.clone(),
+                        }])
+                    } else {
+                        let mut types = Vec::new();
+                        for i in 0..len {
+                            let base_offset = (i as usize * element_layout) as isize;
+
+                            for elem_type in &element_tree.0 {
+                                types.push(Type {
+                                    offset: if elem_type.offset == -1 {
+                                        base_offset
+                                    } else {
+                                        base_offset + elem_type.offset
+                                    },
+                                    size: elem_type.size,
+                                    kind: elem_type.kind,
+                                    child: elem_type.child.clone(),
+                                });
+                            }
+                        }
+                        TypeTree(types)
+                    }
+                }
             }
-
+        } else {
+            TypeTree::new()
+        }
+    } else if ty.is_slice() {
+        if let ty::Slice(element_ty) = ty.kind() {
+            typetree_from_ty_impl(tcx, *element_ty, depth + 1, visited)
+        } else {
+            TypeTree::new()
+        }
+    } else if let ty::Tuple(tuple_types) = ty.kind() {
+        if tuple_types.is_empty() {
+            TypeTree::new()
+        } else {
             let mut types = Vec::new();
-            for i in 0..len {
-                let base_offset = (i as usize * element_layout) as isize;
+            let mut current_offset = 0;
+
+            for tuple_ty in tuple_types.iter() {
+                let element_tree = typetree_from_ty_impl(tcx, tuple_ty, depth + 1, visited);
+                let element_layout = tcx
+                    .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(tuple_ty))
+                    .ok()
+                    .map(|layout| layout.size.bytes_usize())
+                    .unwrap_or(0);
 
                 for elem_type in &element_tree.0 {
                     types.push(Type {
                         offset: if elem_type.offset == -1 {
-                            base_offset
+                            current_offset as isize
                         } else {
-                            base_offset + elem_type.offset
+                            current_offset as isize + elem_type.offset
                         },
                         size: elem_type.size,
                         kind: elem_type.kind,
                         child: elem_type.child.clone(),
                     });
                 }
-            }
 
-            return TypeTree(types);
-        }
-    }
-
-    if ty.is_slice() {
-        if let ty::Slice(element_ty) = ty.kind() {
-            let element_tree = typetree_from_ty(tcx, *element_ty);
-            return element_tree;
-        }
-    }
-
-    if let ty::Tuple(tuple_types) = ty.kind() {
-        if tuple_types.is_empty() {
-            return TypeTree::new();
-        }
-
-        let mut types = Vec::new();
-        let mut current_offset = 0;
-
-        for tuple_ty in tuple_types.iter() {
-            let element_tree = typetree_from_ty(tcx, tuple_ty);
-
-            let element_layout = tcx
-                .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(tuple_ty))
-                .ok()
-                .map(|layout| layout.size.bytes_usize())
-                .unwrap_or(0);
-
-            for elem_type in &element_tree.0 {
-                types.push(Type {
-                    offset: if elem_type.offset == -1 {
-                        current_offset as isize
-                    } else {
-                        current_offset as isize + elem_type.offset
-                    },
-                    size: elem_type.size,
-                    kind: elem_type.kind,
-                    child: elem_type.child.clone(),
-                });
+                current_offset += element_layout;
             }
 
-            current_offset += element_layout;
+            TypeTree(types)
         }
-
-        return TypeTree(types);
-    }
-
-    if let ty::Adt(adt_def, args) = ty.kind() {
+    } else if let ty::Adt(adt_def, args) = ty.kind() {
         if adt_def.is_struct() {
             let struct_layout =
                 tcx.layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(ty));
@@ -2385,7 +2412,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, visited);
 
                     let field_offset = layout.fields.offset(field_idx).bytes_usize();
 
@@ -2403,10 +2430,17 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
                     }
                 }
 
-                return TypeTree(types);
+                TypeTree(types)
+            } else {
+                TypeTree::new()
             }
+        } else {
+            TypeTree::new()
         }
-    }
+    } else {
+        TypeTree::new()
+    };
 
-    TypeTree::new()
+    visited.pop();
+    result
 }

tests/run-make/autodiff/type-trees/mixed-struct-typetree/mixed.check

@@ -0,0 +1,2 @@
+; Check that mixed struct with large array compiles successfully
+CHECK: define{{.*}}"enzyme_type"="{[]:Float@float}"{{.*}}@test_mixed_struct{{.*}}"enzyme_type"="{[]:Pointer}"

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

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

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

@@ -0,0 +1,23 @@
+#![feature(autodiff)]
+
+use std::autodiff::autodiff_reverse;
+
+#[repr(C)]
+struct Container {
+    header: i64,
+    data: [f32; 1000],
+}
+
+#[autodiff_reverse(d_test, Duplicated, Active)]
+#[no_mangle]
+#[inline(never)]
+fn test_mixed_struct(container: &Container) -> f32 {
+    container.data[0] + container.data[999]
+}
+
+fn main() {
+    let container = Container { header: 42, data: [1.0; 1000] };
+    let mut d_container = Container { header: 0, data: [0.0; 1000] };
+    let result = d_test(&container, &mut d_container, 1.0);
+    std::hint::black_box(result);
+}

tests/run-make/autodiff/type-trees/recursion-depth-typetree/recursion.check

@@ -0,0 +1,8 @@
+; Check that recursive types compile successfully without infinite loops
+; The recursion depth limit should prevent stack overflow during TypeTree generation
+
+; Function should compile and have enzyme_type attributes
+CHECK: define{{.*}}"enzyme_type"="{[]:Float@float}"{{.*}}@test_recursion_depth{{.*}}"enzyme_type"="{[]:Pointer}"
+
+; Compilation should complete without hanging or crashing
+; The mere fact that this test runs means recursion limiting is working

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

@@ -0,0 +1,18 @@
+//@ needs-enzyme
+//@ ignore-cross-compile
+
+use run_make_support::{llvm_filecheck, rfs, rustc};
+
+fn main() {
+    // This test ensures that recursive types don't cause infinite loops
+    // The compiler should complete successfully due to recursion limits
+    rustc()
+        .input("test.rs")
+        .arg("-Zautodiff=Enable")
+        .arg("-Zautodiff=NoPostopt")
+        .opt_level("0")
+        .emit("llvm-ir")
+        .run();
+
+    llvm_filecheck().patterns("recursion.check").stdin_buf(rfs::read("test.ll")).run();
+}

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

@@ -0,0 +1,30 @@
+#![feature(autodiff)]
+
+use std::autodiff::autodiff_reverse;
+
+// Create mutually recursive types that would cause cycles
+#[repr(C)]
+struct NodeA {
+    value: f32,
+    b_ref: Option<Box<NodeB>>,
+}
+
+#[repr(C)]
+struct NodeB {
+    value: f64,
+    a_ref: Option<Box<NodeA>>, // Mutual recursion: A -> B -> A -> B...
+}
+
+#[autodiff_reverse(d_test, Duplicated, Active)]
+#[no_mangle]
+#[inline(never)]
+fn test_recursion_depth(node: &NodeA) -> f32 {
+    node.value
+}
+
+fn main() {
+    let node = NodeA { value: 1.0, b_ref: None };
+    let mut d_node = NodeA { value: 0.0, b_ref: None };
+    let result = d_test(&node, &mut d_node, 1.0);
+    std::hint::black_box(result);
+}