basic type tree added

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

Author: KMJ-007

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -11,6 +11,7 @@ mod llvm_enzyme {
         AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ret_activity,
         valid_ty_for_activity,
     };
+    use rustc_ast::expand::typetree::{TypeTree, Type, Kind};
     use rustc_ast::ptr::P;
     use rustc_ast::token::{Lit, LitKind, Token, TokenKind};
     use rustc_ast::tokenstream::*;
@@ -25,7 +26,6 @@ mod llvm_enzyme {
     use tracing::{debug, trace};
 
     use crate::errors;
-    use crate::expand::typetree::TypeTree;
 
     pub(crate) fn outer_normal_attr(
         kind: &P<rustc_ast::NormalAttr>,
@@ -325,10 +325,9 @@ mod llvm_enzyme {
         }
         let span = ecx.with_def_site_ctxt(expand_span);
 
-        // Prepare placeholder type trees for now
-        let num_args = sig.decl.inputs.len();
-        let inputs = vec![TypeTree::new(); num_args];
-        let output = TypeTree::new();
+        // Construct real type trees from function signature
+        let (inputs, output) = construct_typetree_from_fnsig(&sig);
+        
         // Use the new into_item method to construct the AutoDiffItem
         let autodiff_item = x.clone().into_item(
             primal.to_string(),
@@ -1059,4 +1058,105 @@ mod llvm_enzyme {
     }
 }
 
+#[cfg(llvm_enzyme)]
+fn construct_typetree_from_ty(ty: &ast::Ty) -> TypeTree {
+    match &ty.kind {
+        TyKind::Path(..) => {
+            // Handle basic types like f32, f64, i32, etc.
+            // For now, we'll use a simple heuristic based on the path
+            // In a full implementation, this would need to be more sophisticated
+            TypeTree(vec![Type {
+                offset: 0,
+                size: 8, // Default size, should be computed properly
+                kind: Kind::Float, // Default to float, should be determined from type
+                child: TypeTree::new(),
+            }])
+        }
+        TyKind::Ptr(ptr_ty) => {
+            TypeTree(vec![Type {
+                offset: 0,
+                size: 8, // Pointer size
+                kind: Kind::Pointer,
+                child: construct_typetree_from_ty(&ptr_ty.ty),
+            }])
+        }
+        TyKind::Ref(_, ref_ty) => {
+            TypeTree(vec![Type {
+                offset: 0,
+                size: 8, // Reference size
+                kind: Kind::Pointer,
+                child: construct_typetree_from_ty(&ref_ty.ty),
+            }])
+        }
+        TyKind::Slice(slice_ty) => {
+            TypeTree(vec![Type {
+                offset: 0,
+                size: 16, // Slice is (ptr, len)
+                kind: Kind::Pointer,
+                child: construct_typetree_from_ty(&slice_ty.ty),
+            }])
+        }
+        TyKind::Array(array_ty) => {
+            // For arrays, we need to handle the size properly
+            let elem_ty = construct_typetree_from_ty(&array_ty.ty);
+            TypeTree(vec![Type {
+                offset: 0,
+                size: 8, // Array size depends on element type and count
+                kind: Kind::Pointer,
+                child: elem_ty,
+            }])
+        }
+        TyKind::Tup(tuple_types) => {
+            let mut types = Vec::new();
+            let mut offset = 0;
+            for (i, tuple_ty) in tuple_types.iter().enumerate() {
+                let elem_ty = construct_typetree_from_ty(tuple_ty);
+                // For tuples, we need to handle alignment and padding
+                // This is a simplified version
+                types.push(Type {
+                    offset: offset as isize,
+                    size: 8, // Should be computed based on actual type
+                    kind: Kind::Float, // Default
+                    child: elem_ty,
+                });
+                offset += 8; // Simplified alignment
+            }
+            TypeTree(types)
+        }
+        _ => {
+            // Default case for unknown types
+            TypeTree(vec![Type {
+                offset: 0,
+                size: 8,
+                kind: Kind::Float,
+                child: TypeTree::new(),
+            }])
+        }
+    }
+}
+
+#[cfg(llvm_enzyme)]
+fn construct_typetree_from_fnsig(sig: &ast::FnSig) -> (Vec<TypeTree>, TypeTree) {
+    // Construct type trees for input arguments
+    let inputs: Vec<TypeTree> = sig.decl.inputs.iter()
+        .map(|param| construct_typetree_from_ty(&param.ty))
+        .collect();
+    
+    // Construct type tree for return type
+    let output = match &sig.decl.output {
+        FnRetTy::Default(span) => {
+            // Unit type ()
+            TypeTree(vec![Type {
+                offset: 0,
+                size: 0,
+                kind: Kind::Integer,
+                child: TypeTree::new(),
+            }])
+        }
+        FnRetTy::Ty(ty) => construct_typetree_from_ty(ty),
+    };
+    
+    (inputs, output)
+}
+
 pub(crate) use llvm_enzyme::{expand_forward, expand_reverse};

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -1,6 +1,7 @@
 use std::ptr;
 
 use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivity, DiffMode};
+use rustc_ast::expand::typetree::{FncTree, TypeTree};
 use rustc_codegen_ssa::ModuleCodegen;
 use rustc_codegen_ssa::back::write::ModuleConfig;
 use rustc_codegen_ssa::common::TypeKind;
@@ -16,6 +17,7 @@ use crate::declare::declare_simple_fn;
 use crate::errors::{AutoDiffWithoutEnable, LlvmError};
 use crate::llvm::AttributePlace::Function;
 use crate::llvm::{Metadata, True};
+use crate::typetree::to_enzyme_typetree;
 use crate::value::Value;
 use crate::{CodegenContext, LlvmCodegenBackend, ModuleLlvm, attributes, llvm};
 
@@ -294,6 +296,7 @@ fn generate_enzyme_call<'ll>(
     fn_to_diff: &'ll Value,
     outer_fn: &'ll Value,
     attrs: AutoDiffAttrs,
+    fnc_tree: Option<FncTree>,
 ) {
     // We have to pick the name depending on whether we want forward or reverse mode autodiff.
     let mut ad_name: String = match attrs.mode {
@@ -361,6 +364,15 @@ fn generate_enzyme_call<'ll>(
         let attr = llvm::AttributeKind::NoInline.create_attr(cx.llcx);
         attributes::apply_to_llfn(ad_fn, Function, &[attr]);
 
+        // TODO(KMJ-007): Add type tree metadata if available
+        // This requires adding CreateTypeTreeAttribute to LLVM bindings
+        // if let Some(tree) = fnc_tree {
+        //     let data_layout = cx.data_layout();
+        //     let enzyme_tree = to_enzyme_typetree(tree, data_layout, cx.llcx);
+        //     let tt_attr = llvm::CreateTypeTreeAttribute(cx.llcx, enzyme_tree);
+        //     attributes::apply_to_llfn(ad_fn, Function, &[tt_attr]);
+        // }
+
         // We add a made-up attribute just such that we can recognize it after AD to update
         // (no)-inline attributes. We'll then also remove this attribute.
         let enzyme_marker_attr = llvm::CreateAttrString(cx.llcx, "enzyme_marker");
@@ -505,7 +517,17 @@ pub(crate) fn differentiate<'ll>(
             ));
         };
 
-        generate_enzyme_call(&cx, fn_def, fn_target, item.attrs.clone());
+        // Construct function type tree from item's type trees
+        let fnc_tree = if !item.inputs.is_empty() || !item.output.0.is_empty() {
+            Some(FncTree {
+                inputs: item.inputs.clone(),
+                output: item.output.clone(),
+            })
+        } else {
+            None
+        };
+
+        generate_enzyme_call(&cx, fn_def, fn_target, item.attrs.clone(), fnc_tree);
     }
 
     // FIXME(ZuseZ4): support SanitizeHWAddress and prevent illegal/unsupported opts

compiler/rustc_codegen_llvm/src/lib.rs

@@ -74,6 +74,7 @@ mod llvm_util;
 mod mono_item;
 mod type_;
 mod type_of;
+mod typetree;
 mod va_arg;
 mod value;
 

compiler/rustc_codegen_llvm/src/typetree.rs

@@ -0,0 +1,33 @@
+use crate::llvm;
+use rustc_ast::expand::typetree::{Kind, TypeTree};
+
+pub fn to_enzyme_typetree(
+    tree: TypeTree,
+    llvm_data_layout: &str,
+    llcx: &llvm::Context,
+) -> llvm::TypeTree {
+    tree.0.iter().fold(llvm::TypeTree::new(), |obj, x| {
+        let scalar = match x.kind {
+            Kind::Integer => llvm::CConcreteType::DT_Integer,
+            Kind::Float => llvm::CConcreteType::DT_Float,
+            Kind::Double => llvm::CConcreteType::DT_Double,
+            Kind::Pointer => llvm::CConcreteType::DT_Pointer,
+            _ => panic!("Unknown kind {:?}", x.kind),
+        };
+
+        let tt = llvm::TypeTree::from_type(scalar, llcx).only(-1);
+
+        let tt = if !x.child.0.is_empty() {
+            let inner_tt = to_enzyme_typetree(x.child.clone(), llvm_data_layout, llcx);
+            tt.merge(inner_tt.only(-1))
+        } else {
+            tt
+        };
+
+        if x.offset != -1 {
+            obj.merge(tt.shift(llvm_data_layout, 0, x.size as isize, x.offset as usize))
+        } else {
+            obj.merge(tt)
+        }
+    })
+}