continue rebasing

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -16,7 +16,7 @@ use crate::type_::Type;
 use crate::typetree::to_enzyme_typetree;
 use crate::LlvmCodegenBackend;
 use crate::ModuleLlvm;
-use crate::{base, DiffTypeTree};
+use crate::DiffTypeTree;
 use llvm::{
     enzyme_rust_forward_diff, enzyme_rust_reverse_diff, BasicBlock, CreateEnzymeLogic,
     CreateTypeAnalysis, EnzymeLogicRef, EnzymeTypeAnalysisRef, LLVMAddFunction,

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -409,7 +409,7 @@ fn add_unused_functions(cx: &CodegenCx<'_, '_>) {
 /// All items participating in code generation together with (instrumented)
 /// items inlined into them.
 fn codegenned_and_inlined_items(tcx: TyCtxt<'_>) -> DefIdSet {
-    let (items, cgus) = tcx.collect_and_partition_mono_items(());
+    let (items, _, cgus) = tcx.collect_and_partition_mono_items(());
     let mut visited = DefIdSet::default();
     let mut result = items.clone();
 

compiler/rustc_codegen_llvm/src/lib.rs

@@ -40,13 +40,12 @@ use rustc_codegen_ssa::back::write::{
 use rustc_codegen_ssa::traits::*;
 use rustc_codegen_ssa::ModuleCodegen;
 use rustc_codegen_ssa::{CodegenResults, CompiledModule};
-use rustc_data_structures::fx::FxIndexMap;
+use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
 use rustc_errors::{DiagnosticMessage, ErrorGuaranteed, FatalError, Handler, SubdiagnosticMessage};
 use rustc_fluent_macro::fluent_messages;
 use rustc_metadata::EncodedMetadata;
 use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
 use rustc_middle::middle::autodiff_attrs::AutoDiffItem;
-use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::TyCtxt;
 use rustc_middle::util::Providers;
 use rustc_session::config::{OptLevel, OutputFilenames, PrintKind, PrintRequest};

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1,7 +1,6 @@
 #![allow(non_camel_case_types)]
 #![allow(non_upper_case_globals)]
 
-use rustc_codegen_ssa::coverageinfo::map as coverage_map;
 use rustc_middle::middle::autodiff_attrs::DiffActivity;
 
 use super::debuginfo::{

compiler/rustc_monomorphize/src/partitioning.rs

@@ -92,6 +92,11 @@
 //! source-level module, functions from the same module will be available for
 //! inlining, even when they are not marked `#[inline]`.
 
+// Manuel, fixing rebase
+use rustc_symbol_mangling::symbol_name_for_instance_in_crate;
+//use crate::ty::ParamEnv;
+use rustc_middle::ty::ParamEnv;
+
 use std::cmp;
 use std::collections::hash_map::Entry;
 use std::fs::{self, File};
@@ -1138,6 +1143,50 @@ fn collect_and_partition_mono_items(tcx: TyCtxt<'_>, (): ()) -> (&DefIdSet, &[Au
         })
         .collect();
 
+
+    let autodiff_items = items
+        .iter()
+        .filter_map(|item| match *item {
+            MonoItem::Fn(ref instance) => Some((item, instance)),
+            _ => None,
+        })
+        .filter_map(|(item, instance)| {
+            let target_id = instance.def_id();
+            let target_attrs = tcx.autodiff_attrs(target_id);
+            if !target_attrs.apply_autodiff() {
+                return None;
+            }
+
+            let target_symbol =
+                symbol_name_for_instance_in_crate(tcx, instance.clone(), LOCAL_CRATE);
+            let range = usage_map.index.get(&item).unwrap();
+
+            let source = usage_map.targets[range.clone()]
+                .into_iter()
+                .filter_map(|item| match *item {
+                    MonoItem::Fn(ref instance_s) => {
+                        let source_id = instance_s.def_id();
+
+                        if tcx.autodiff_attrs(source_id).is_active() {
+                            return Some(instance_s);
+                        }
+
+                        None
+                    }
+                    _ => None,
+                })
+                .next();
+
+            source.map(|inst| {
+                let (inputs, output) = fnc_typetrees(inst.ty(tcx, ParamEnv::empty()), tcx);
+                let symb = symbol_name_for_instance_in_crate(tcx, inst.clone(), LOCAL_CRATE);
+
+                target_attrs.clone().into_item(symb, target_symbol, inputs, output)
+            })
+        });
+
+    let autodiff_items = tcx.arena.alloc_from_iter(autodiff_items);
+
     // Output monomorphization stats per def_id
     if let SwitchWithOptPath::Enabled(ref path) = tcx.sess.opts.unstable_opts.dump_mono_stats {
         if let Err(err) =
@@ -1198,7 +1247,163 @@ fn collect_and_partition_mono_items(tcx: TyCtxt<'_>, (): ()) -> (&DefIdSet, &[Au
         }
     }
 
-    (tcx.arena.alloc(mono_items), codegen_units)
+    (tcx.arena.alloc(mono_items), autodiff_items, codegen_units)
+}
+use rustc_middle::ty::{self, Adt, ParamEnvAnd, Ty};
+use rustc_target::abi::FieldsShape;
+use std::iter;
+
+pub fn typetree_empty() -> TypeTree {
+    TypeTree(vec![])
+}
+
+pub fn typetree_from_ty<'a>(ty: Ty<'a>, tcx: TyCtxt<'a>, depth: usize) -> TypeTree {
+    if ty.is_unsafe_ptr() || ty.is_ref() || ty.is_box() {
+        if ty.is_fn_ptr() {
+            unimplemented!("what to do whith fn ptr?");
+        }
+
+        let inner_ty = ty.builtin_deref(true).unwrap().ty;
+        let child = typetree_from_ty(inner_ty, tcx, depth + 1);
+
+        let tt = Type { offset: -1, kind: Kind::Pointer, size: 8, child };
+        //println!("{:depth$} add indirection {:?}", "", tt);
+
+        return TypeTree(vec![tt]);
+    }
+
+    if ty.is_scalar() {
+        assert!(!ty.is_any_ptr());
+
+        let (kind, size) = if ty.is_integral() {
+            (Kind::Integer, 8)
+        } else {
+            assert!(ty.is_floating_point());
+            match ty {
+                x if x == tcx.types.f32 => (Kind::Float, 4),
+                x if x == tcx.types.f64 => (Kind::Double, 8),
+                _ => panic!("floatTy scalar that is neither f32 nor f64"),
+            }
+        };
+
+        return TypeTree(vec![Type { offset: -1, child: typetree_empty(), kind, size }]);
+    }
+
+    let param_env_and = ParamEnvAnd { param_env: ParamEnv::empty(), value: ty };
+
+    let layout = tcx.layout_of(param_env_and);
+    assert!(layout.is_ok());
+
+    let layout = layout.unwrap().layout;
+    let fields = layout.fields();
+    let max_size = layout.size();
+
+    if ty.is_adt() {
+        let adt_def = ty.ty_adt_def().unwrap();
+        let substs = match ty.kind() {
+            Adt(_, subst_ref) => subst_ref,
+            _ => panic!(""),
+        };
+
+        if adt_def.is_struct() {
+            let (offsets, _memory_index) = match fields {
+                FieldsShape::Arbitrary { offsets: o, memory_index: m } => (o, m),
+                _ => panic!(""),
+            };
+            //println!("{:depth$} combine fields", "");
+
+            let fields = adt_def.all_fields();
+            let fields = fields
+                .into_iter()
+                .zip(offsets.into_iter())
+                .filter_map(|(field, offset)| {
+                    let field_ty: Ty<'_> = field.ty(tcx, substs);
+                    let field_ty: Ty<'_> =
+                        tcx.normalize_erasing_regions(ParamEnv::empty(), field_ty);
+
+                    if field_ty.is_phantom_data() {
+                        return None;
+                    }
+
+                    let mut child = typetree_from_ty(field_ty, tcx, depth + 1).0;
+
+                    for c in &mut child {
+                        if c.offset == -1 {
+                            c.offset = offset.bytes() as isize
+                        } else {
+                            c.offset += offset.bytes() as isize;
+                        }
+                    }
+
+                    //inner_tt.offset = offset;
+
+                    //println!("{:depth$} -> {:?}", "", child);
+
+                    Some(child)
+                })
+                .flatten()
+                .collect::<Vec<Type>>();
+
+            let ret_tt = TypeTree(fields);
+            //println!("{:depth$} into {:?}", "", ret_tt);
+            return ret_tt;
+        } else {
+            unimplemented!("adt that isn't a struct");
+        }
+    }
+
+    if ty.is_array() {
+        let (stride, count) = match fields {
+            FieldsShape::Array { stride: s, count: c } => (s, c),
+            _ => panic!(""),
+        };
+        let byte_stride = stride.bytes_usize();
+        let byte_max_size = max_size.bytes_usize();
+
+        assert!(byte_stride * *count as usize == byte_max_size);
+        assert!(*count > 0); // return empty TT for empty?
+        let sub_ty = ty.builtin_index().unwrap();
+        let subtt = typetree_from_ty(sub_ty, tcx, depth + 1);
+
+        // calculate size of subtree
+        let param_env_and = ParamEnvAnd { param_env: ParamEnv::empty(), value: sub_ty };
+        let size = tcx.layout_of(param_env_and).unwrap().size.bytes() as usize;
+        let tt = TypeTree(
+            iter::repeat(subtt)
+                .take(*count as usize)
+                .enumerate()
+                .map(|(idx, x)| x.0.into_iter().map(move |x| x.add_offset((idx * size) as isize)))
+                .flatten()
+                .collect(),
+        );
+
+        //println!("{:depth$} repeated array into {:?}", "", tt);
+
+        return tt;
+    }
+
+    if ty.is_slice() {
+        let sub_ty = ty.builtin_index().unwrap();
+        let subtt = typetree_from_ty(sub_ty, tcx, depth + 1);
+
+        return subtt;
+    }
+
+    //println!("Warning: create empty typetree for {}", ty);
+    typetree_empty()
+}
+
+pub fn fnc_typetrees<'tcx>(fn_ty: Ty<'tcx>, tcx: TyCtxt<'tcx>) -> (Vec<TypeTree>, TypeTree) {
+    let fnc_binder: ty::Binder<'_, ty::FnSig<'_>> = fn_ty.fn_sig(tcx);
+
+    // TODO: verify.
+    let x: ty::FnSig<'_> = fnc_binder.skip_binder();
+
+    let inputs = x.inputs().into_iter().map(|x| typetree_from_ty(*x, tcx, 0)).collect();
+
+    let output = typetree_from_ty(x.output(), tcx, 0);
+
+    (inputs, output)
 }
 
 /// Outputs stats about instantiation counts and estimated size, per `MonoItem`'s