autodiff: generate TypeTree metadata for extractvalue

compiler/rustc_codegen_llvm/src/asm.rs

@@ -365,7 +365,7 @@ impl<'ll, 'tcx> AsmBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                     let value = if output_types.len() == 1 {
                         result
                     } else {
-                        self.extract_value(result, op_idx[&idx] as u64)
+                        self.extract_value(result, op_idx[&idx] as u64, None)
                     };
                     let value =
                         llvm_fixup_output(self, value, reg.reg_class(), &place.layout, instance);

compiler/rustc_codegen_llvm/src/builder.rs

@@ -2,7 +2,7 @@ use std::borrow::{Borrow, Cow};
 use std::iter;
 use std::ops::Deref;
 
-use rustc_ast::expand::typetree::FncTree;
+use rustc_ast::expand::typetree::{FncTree, TypeTree};
 pub(crate) mod autodiff;
 pub(crate) mod gpu_offload;
 
@@ -636,7 +636,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         let name = format!("llvm.{}{oop_str}.with.overflow", if signed { 's' } else { 'u' });
 
         let res = self.call_intrinsic(name, &[self.type_ix(width)], &[lhs, rhs]);
-        (self.extract_value(res, 0), self.extract_value(res, 1))
+        (self.extract_value(res, 0, None), self.extract_value(res, 1, None))
     }
 
     fn from_immediate(&mut self, val: Self::Value) -> Self::Value {
@@ -1262,9 +1262,19 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         }
     }
 
-    fn extract_value(&mut self, agg_val: &'ll Value, idx: u64) -> &'ll Value {
+    fn extract_value(&mut self, agg_val: &'ll Value, idx: u64, tt: Option<TypeTree>) -> &'ll Value {
         assert_eq!(idx as c_uint as u64, idx);
-        unsafe { llvm::LLVMBuildExtractValue(self.llbuilder, agg_val, idx as c_uint, UNNAMED) }
+        let extract =
+            unsafe { llvm::LLVMBuildExtractValue(self.llbuilder, agg_val, idx as c_uint, UNNAMED) };
+        if let Some(tt) = tt {
+            let fnc_tree = FncTree { args: vec![], ret: tt };
+            crate::typetree::add_tt(
+                self,
+                extract,
+                fnc_tree,
+            );
+        }
+        extract
     }
 
     fn insert_value(&mut self, agg_val: &'ll Value, elt: &'ll Value, idx: u64) -> &'ll Value {
@@ -1284,7 +1294,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         unsafe {
             llvm::LLVMSetCleanup(landing_pad, llvm::TRUE);
         }
-        (self.extract_value(landing_pad, 0), self.extract_value(landing_pad, 1))
+        (self.extract_value(landing_pad, 0, None), self.extract_value(landing_pad, 1, None))
     }
 
     fn filter_landing_pad(&mut self, pers_fn: &'ll Value) {
@@ -1385,8 +1395,8 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
                 llvm::FALSE, // SingleThreaded
             );
             llvm::LLVMSetWeak(value, weak.to_llvm_bool());
-            let val = self.extract_value(value, 0);
-            let success = self.extract_value(value, 1);
+            let val = self.extract_value(value, 0, None);
+            let success = self.extract_value(value, 1, None);
             (val, success)
         }
     }

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -162,9 +162,9 @@ impl<'ll> OffloadKernelDims<'ll> {
             builder: &mut Builder<'_, 'll, 'tcx>,
             arr: &'ll Value,
         ) -> &'ll Value {
-            let x = builder.extract_value(arr, 0);
-            let y = builder.extract_value(arr, 1);
-            let z = builder.extract_value(arr, 2);
+            let x = builder.extract_value(arr, 0, None);
+            let y = builder.extract_value(arr, 1, None);
+            let z = builder.extract_value(arr, 2, None);
 
             let xy = builder.mul(x, y);
             builder.mul(xy, z)

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -791,6 +791,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 self.extract_value(
                     args[0].immediate(),
                     fn_args.const_at(2).to_leaf().to_i32() as u64,
+                    None,
                 )
             }
 
@@ -1074,7 +1075,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
             &[],
             &[llvtable, vtable_byte_offset, typeid],
         );
-        self.extract_value(type_checked_load, 0)
+        self.extract_value(type_checked_load, 0, None)
     }
 
     fn va_start(&mut self, va_list: &'ll Value) {
@@ -1186,7 +1187,7 @@ fn autocast<'ll>(
                 iter::zip(bx.struct_element_types(src_ty), bx.struct_element_types(dest_ty))
                     .enumerate()
             {
-                let elt = bx.extract_value(val, idx as u64);
+                let elt = bx.extract_value(val, idx as u64, None);
                 let casted_elt = autocast(bx, elt, src_element_ty, dest_element_ty);
                 ret = bx.insert_value(ret, casted_elt, idx as u64);
             }
@@ -1647,7 +1648,7 @@ fn codegen_gnu_try<'ll, 'tcx>(
         let vals = bx.landing_pad(lpad_ty, bx.eh_personality(), 1);
         let tydesc = bx.const_null(bx.type_ptr());
         bx.add_clause(vals, tydesc);
-        let ptr = bx.extract_value(vals, 0);
+        let ptr = bx.extract_value(vals, 0, None);
         let catch_ty = bx.type_func(&[bx.type_ptr(), bx.type_ptr()], bx.type_void());
         bx.call(catch_ty, None, None, catch_func, &[data, ptr], None, None);
         bx.ret(bx.const_bool(true));
@@ -1937,8 +1938,9 @@ fn get_args_from_tuple<'ll, 'tcx>(
                         let field = tuple_place.project_field(bx, tuple_index);
                         let llvm_ty = field.layout.llvm_type(bx.cx);
                         let pair_val = bx.load(llvm_ty, field.val.llval, field.val.align);
-                        result.push(bx.extract_value(pair_val, 0));
-                        result.push(bx.extract_value(pair_val, 1));
+
+                        result.push(bx.extract_value(pair_val, 0, None));
+                        result.push(bx.extract_value(pair_val, 1, None));
                         tuple_index += 1;
                     }
                     PassMode::Indirect { .. } => {

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -67,6 +67,8 @@ unsafe extern "C" {
     ) -> Option<&Value>;
 
     pub(crate) safe fn LLVMRustIsCall(V: &Value) -> bool;
+    pub(crate) safe fn LLVMRustSupportsEnzymeMD(V: &Value) -> bool;
+    pub(crate) fn LLVMRustSetEnzymeTypeMD(v: &Value, md: &Value);
 }
 
 unsafe extern "C" {
@@ -97,7 +99,7 @@ pub(crate) mod Enzyme_AD {
     use rustc_session::filesearch;
 
     use super::{CConcreteType, CTypeTreeRef, Context};
-    use crate::llvm::{EnzymeTypeTree, LLVMRustVersionMajor};
+    use crate::llvm::{EnzymeTypeTree, LLVMRustVersionMajor, Value};
 
     type EnzymeSetCLBoolFn = unsafe extern "C" fn(*mut c_void, u8);
     type EnzymeSetCLStringFn = unsafe extern "C" fn(*mut c_void, *const c_char);
@@ -115,6 +117,7 @@ pub(crate) mod Enzyme_AD {
         unsafe extern "C" fn(CTypeTreeRef, *const i64, usize, CConcreteType, &Context);
     type EnzymeTypeTreeToStringFn = unsafe extern "C" fn(CTypeTreeRef) -> *const c_char;
     type EnzymeTypeTreeToStringFreeFn = unsafe extern "C" fn(*const c_char);
+    type EnzymeTypeTreeToMDFn = unsafe extern "C" fn(CTypeTreeRef, &Context) -> Option<&Value>;
 
     #[allow(non_snake_case)]
     pub(crate) struct EnzymeWrapper {
@@ -129,6 +132,7 @@ pub(crate) mod Enzyme_AD {
         EnzymeTypeTreeInsertEq: EnzymeTypeTreeInsertEqFn,
         EnzymeTypeTreeToString: EnzymeTypeTreeToStringFn,
         EnzymeTypeTreeToStringFree: EnzymeTypeTreeToStringFreeFn,
+        EnzymeTypeTreeToMD: EnzymeTypeTreeToMDFn,
 
         EnzymePrintPerf: *mut c_void,
         EnzymePrintActivity: *mut c_void,
@@ -302,6 +306,14 @@ pub(crate) mod Enzyme_AD {
             unsafe { (self.EnzymeTypeTreeToStringFree)(ch) }
         }
 
+        pub(crate) fn tree_to_md<'a>(
+            &'a self,
+            tree: *mut EnzymeTypeTree,
+            ctx: &'a Context,
+        ) -> Option<&'a Value> {
+            unsafe { (self.EnzymeTypeTreeToMD)(tree, ctx) }
+        }
+
         pub(crate) fn get_max_type_depth(&self) -> usize {
             unsafe { std::ptr::read::<u32>(self.EnzymeMaxTypeDepth as *const u32) as usize }
         }
@@ -385,6 +397,7 @@ pub(crate) mod Enzyme_AD {
                 EnzymeTypeTreeInsertEq: EnzymeTypeTreeInsertEqFn,
                 EnzymeTypeTreeToString: EnzymeTypeTreeToStringFn,
                 EnzymeTypeTreeToStringFree: EnzymeTypeTreeToStringFreeFn,
+                EnzymeTypeTreeToMD: EnzymeTypeTreeToMDFn,
                 EnzymeSetCLBool: EnzymeSetCLBoolFn,
                 EnzymeSetCLString: EnzymeSetCLStringFn,
             );
@@ -416,6 +429,7 @@ pub(crate) mod Enzyme_AD {
                 EnzymeTypeTreeInsertEq,
                 EnzymeTypeTreeToString,
                 EnzymeTypeTreeToStringFree,
+                EnzymeTypeTreeToMD,
                 EnzymePrintPerf,
                 EnzymePrintActivity,
                 EnzymePrintType,

compiler/rustc_codegen_llvm/src/typetree.rs

@@ -73,6 +73,7 @@ fn process_typetree_recursive(
 enum TTLocation {
     Definition,
     Callsite,
+    Intrinsic,
 }
 
 #[cfg_attr(not(feature = "llvm_enzyme"), allow(unused))]
@@ -106,14 +107,19 @@ pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt:
     let attr_name = "enzyme_type";
     let c_attr_name = CString::new(attr_name).unwrap();
 
-    let tt_location: TTLocation =
-        if llvm::LLVMRustIsCall(fn_def) { TTLocation::Callsite } else { TTLocation::Definition };
+    let tt_location: TTLocation = if llvm::LLVMRustIsCall(fn_def) {
+        TTLocation::Callsite
+    } else if llvm::LLVMRustSupportsEnzymeMD(fn_def) {
+        TTLocation::Intrinsic
+    } else {
+        TTLocation::Definition
+    };
 
     let mut offset = 0;
     for (i, input) in inputs.iter().enumerate() {
         let (enzyme_tt, extra_ints) = to_enzyme_typetree(&input, llvm_data_layout, llcx);
 
-        // This scope is just a visual reminder that we *must* drop the enzyme_wrapper before
+        // This scope is a simple solution since we *must* drop the enzyme_wrapper before
         // we drop any typetrees (mainly enzyme_tt and extra_ints). Drop calls can not accept
         // arguments like an enzyme_wrapper, so the typetree drop impl has to call get_instance
         // on the static enzyme instance, which is behind a Mutex. Therefore we'd deadlock if we
@@ -133,6 +139,12 @@ pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt:
                 TTLocation::Callsite => {
                     attributes::apply_to_callsite(fn_def, arg_pos, &[attr]);
                 }
+                TTLocation::Intrinsic => {
+                    let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx);
+                    unsafe {
+                        llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap());
+                    }
+                }
             }
             enzyme_wrapper.tree_to_string_free(c_str.as_ptr());
             for v in &extra_ints {
@@ -141,6 +153,12 @@ pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt:
                 let int_attr = llvm::CreateAttrStringValueFromCStr(llcx, &c_attr_name, &c_str);
                 let arg_pos = llvm::AttributePlace::Argument(i as u32 + offset);
                 match tt_location {
+                    TTLocation::Intrinsic => {
+                        let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx);
+                        unsafe {
+                            llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap());
+                        }
+                    }
                     TTLocation::Definition => {
                         attributes::apply_to_llfn(fn_def, arg_pos, &[int_attr]);
                     }
@@ -152,8 +170,8 @@ pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt:
             }
         }
     }
-    // We will only fail this if Rust types got lowered to LLVM in a way that we didn't predict.
-    // Error, so we can learn from our mistakes.
+    // We will fail here if Rust types got lowered to LLVM in a way that we didn't predict.
+    // We Error, so we can learn from our mistakes.
     if matches!(tt_location, TTLocation::Definition) {
         let expected = offset as usize + inputs.len();
         let actual = llvm::count_params(fn_def) as usize;
@@ -180,6 +198,12 @@ pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt:
             TTLocation::Callsite => {
                 attributes::apply_to_callsite(fn_def, arg_pos, &[ret_attr]);
             }
+            TTLocation::Intrinsic => {
+                let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx);
+                unsafe {
+                    llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap());
+                }
+            }
         }
         enzyme_wrapper.tree_to_string_free(c_str.as_ptr());
     }

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -2360,8 +2360,8 @@ pub fn store_cast<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     if let Some(offset_from_start) = cast.rest_offset {
         assert_eq!(cast.prefix.len(), 1);
         assert_eq!(cast.rest.unit.size, cast.rest.total);
-        let first = bx.extract_value(value, 0);
-        let second = bx.extract_value(value, 1);
+        let first = bx.extract_value(value, 0, None);
+        let second = bx.extract_value(value, 1, None);
         bx.store(first, ptr, align);
         let second_ptr = bx.inbounds_ptradd(ptr, bx.const_usize(offset_from_start.bytes()));
         bx.store(second, second_ptr, align.restrict_for_offset(offset_from_start));

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -5,21 +5,60 @@ use rustc_abi as abi;
 use rustc_abi::{
     Align, BackendRepr, FIRST_VARIANT, FieldIdx, Primitive, Size, TagEncoding, VariantIdx, Variants,
 };
+use rustc_ast::expand::typetree::TypeTree;
 use rustc_hir::LangItem;
 use rustc_middle::mir::interpret::{Pointer, Scalar, alloc_range};
 use rustc_middle::mir::{self, ConstValue};
 use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
+use rustc_middle::ty::typetree::typetree_from_ty;
 use rustc_middle::ty::{self, Ty};
 use rustc_middle::{bug, span_bug};
 use rustc_session::config::{AnnotateMoves, DebugInfo, OptLevel};
+use rustc_span::sym;
 use tracing::{debug, instrument};
 
 use super::place::{PlaceRef, PlaceValue};
 use super::rvalue::transmute_scalar;
 use super::{FunctionCx, LocalRef};
-use crate::MemFlags;
 use crate::common::IntPredicate;
 use crate::traits::*;
+use crate::{MemFlags, TyCtxt};
+
+fn option_ptr_like_scalar_pair_tts<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Option<TypeTree> {
+    let ty::Adt(def, args) = ty.kind() else {
+        return None;
+    };
+
+    if !tcx.is_lang_item(def.did(), LangItem::Option) {
+        return None;
+    }
+
+    let inner = args.type_at(0);
+    if !(inner.is_ref() || inner.is_box() || nonnull_inner_ty(tcx, inner).is_some()) {
+        return None;
+    }
+
+    let tt = typetree_from_ty(tcx, inner);
+    //let some_layout = layout.for_variant(bx.cx(), VariantIdx::from_u32(1));
+    //let payload_layout = some_layout.field(bx.cx(), 0);
+    // this will be a slice
+    //let payload_ty = payload_layout.ty;
+    //let tt = rustc_middle::ty::typetree_from_ty(bx.tcx(), field0_ty.unwrap());
+    if tt == TypeTree::new() {
+        return None;
+    }
+    Some(tt)
+}
+
+fn nonnull_inner_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Option<Ty<'tcx>> {
+    if let ty::Adt(def, args) = ty.kind()
+        && tcx.is_diagnostic_item(sym::NonNull, def.did())
+    {
+        return Some(args.type_at(0));
+    }
+
+    None
+}
 
 /// The representation of a Rust value. The enum variant is in fact
 /// uniquely determined by the value's type, but is kept as a
@@ -349,8 +388,10 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
             debug!("Operand::from_immediate_or_packed_pair: unpacking {:?} @ {:?}", llval, layout);
 
             // Deconstruct the immediate aggregate.
-            let a_llval = bx.extract_value(llval, 0);
-            let b_llval = bx.extract_value(llval, 1);
+            let f1 = option_ptr_like_scalar_pair_tts(bx.tcx(), layout.ty);
+            let f2 = if f1.is_none() { None } else { Some(TypeTree::int(8)) };
+            let a_llval = bx.extract_value(llval, 0, f1);
+            let b_llval = bx.extract_value(llval, 1, f2);
             OperandValue::Pair(a_llval, b_llval)
         } else {
             OperandValue::Immediate(llval)

compiler/rustc_codegen_ssa/src/mir/statement.rs

@@ -2,6 +2,10 @@ use rustc_middle::mir::{self, NonDivergingIntrinsic, StmtDebugInfo};
 use rustc_middle::{bug, span_bug, ty};
 use tracing::instrument;
 
+use rustc_ast::expand::typetree::{TypeTree, FncTree};
+use rustc_middle::ty::typetree::typetree_from_ty;
+
+
 use super::{FunctionCx, LocalRef};
 use crate::mir::retag;
 use crate::traits::*;
@@ -111,7 +115,14 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 let dst = dst_val.immediate();
                 let src = src_val.immediate();
 
-                bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty(), None);
+                let tt = typetree_from_ty(bx.tcx(), pointee).add_indirection();
+                dbg!(&tt);
+                let fnc_tree = FncTree {
+                    args: vec![tt.clone(), tt],
+                    ret: TypeTree::new(),
+                };
+
+                bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty(), Some(fnc_tree));
             }
             mir::StatementKind::FakeRead(..)
             | mir::StatementKind::AscribeUserType(..)