update rustc_middle, rustc_monomorphize, and (partly) cg_ssa

Author: ZuseZ4

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_ssa/src/assert_module_sources.rs

@@ -46,7 +46,7 @@ pub fn assert_module_sources(tcx: TyCtxt<'_>, set_reuse: &dyn Fn(&mut CguReuseTr
         }
 
         let available_cgus =
-            tcx.collect_and_partition_mono_items(()).1.iter().map(|cgu| cgu.name()).collect();
+            tcx.collect_and_partition_mono_items(()).2.iter().map(|cgu| cgu.name()).collect();
 
         let mut ams = AssertModuleSource {
             tcx,

compiler/rustc_codegen_ssa/src/back/symbol_export.rs

@@ -331,7 +331,7 @@ fn exported_symbols_provider_local(
         // external linkage is enough for monomorphization to be linked to.
         let need_visibility = tcx.sess.target.dynamic_linking && !tcx.sess.target.only_cdylib;
 
-        let (_, cgus) = tcx.collect_and_partition_mono_items(());
+        let (_, _, cgus) = tcx.collect_and_partition_mono_items(());
 
         for (mono_item, data) in cgus.iter().flat_map(|cgu| cgu.items().iter()) {
             if data.linkage != Linkage::External {

compiler/rustc_codegen_ssa/src/base.rs

@@ -594,7 +594,7 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
 
     // Run the monomorphization collector and partition the collected items into
     // codegen units.
-    let codegen_units = tcx.collect_and_partition_mono_items(()).1;
+    let codegen_units = tcx.collect_and_partition_mono_items(()).2;
 
     // Force all codegen_unit queries so they are already either red or green
     // when compile_codegen_unit accesses them. We are not able to re-execute
@@ -980,7 +980,7 @@ pub fn provide(providers: &mut Providers) {
             config::OptLevel::SizeMin => config::OptLevel::Default,
         };
 
-        let (defids, _) = tcx.collect_and_partition_mono_items(cratenum);
+        let (defids, _, _) = tcx.collect_and_partition_mono_items(cratenum);
 
         let any_for_speed = defids.items().any(|id| {
             let CodegenFnAttrs { optimize, .. } = tcx.codegen_fn_attrs(*id);

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -1,4 +1,5 @@
-use rustc_ast::{ast, attr, MetaItemKind, NestedMetaItem};
+use rustc_ast::{ast, attr, MetaItem, MetaItemKind, NestedMetaItem};
+use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode};
 use rustc_attr::{list_contains_name, InlineAttr, InstructionSetAttr, OptimizeAttr};
 use rustc_errors::struct_span_err;
 use rustc_hir as hir;
@@ -14,6 +15,8 @@ use rustc_span::symbol::Ident;
 use rustc_span::{sym, Span};
 use rustc_target::spec::{abi, SanitizerSet};
 
+use std::str::FromStr;
+
 use crate::errors;
 use crate::target_features::from_target_feature;
 use crate::{
@@ -689,6 +692,174 @@ fn check_link_name_xor_ordinal(
     }
 }
 
+fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
+    //let attrs = tcx.get_attrs(id, sym::autodiff_into);
+    let attrs = tcx.get_attrs(id, sym::autodiff);
+
+    let attrs = attrs
+        .into_iter()
+        .filter(|attr| attr.name_or_empty() == sym::autodiff)
+        //.filter(|attr| attr.name_or_empty() == sym::autodiff_into)
+        .collect::<Vec<_>>();
+    if attrs.len() > 0 {
+      dbg!("autodiff_attrs len = > 0: {}", attrs.len());
+    }
+
+    // check for exactly one autodiff attribute on extern block
+    let msg_once = "autodiff attribute can only be applied once";
+    let attr = match &attrs[..] {
+        &[] => return AutoDiffAttrs::inactive(),
+        &[elm] => elm,
+        x => {
+            tcx.sess
+                .struct_span_err(x[1].span, msg_once)
+                .span_label(x[1].span, "more than one")
+                .emit();
+
+            return AutoDiffAttrs::inactive();
+        }
+    };
+
+    let list = attr.meta_item_list().unwrap_or_default();
+
+    // empty autodiff attribute macros (i.e. `#[autodiff]`) are used to mark source functions
+    if list.len() == 0 {
+        return AutoDiffAttrs {
+            mode: DiffMode::Source,
+            ret_activity: DiffActivity::None,
+            input_activity: Vec::new(),
+        };
+    }
+
+    let msg_ad_mode = "autodiff attribute must contain autodiff mode";
+    let mode = match &list[0] {
+        NestedMetaItem::MetaItem(MetaItem { path: ref p2, kind: MetaItemKind::Word, .. }) => {
+            p2.segments.first().unwrap().ident
+        }
+        _ => {
+            tcx.sess
+                .struct_span_err(attr.span, msg_ad_mode)
+                .span_label(attr.span, "empty argument list")
+                .emit();
+
+            return AutoDiffAttrs::inactive();
+        }
+    };
+
+    // parse mode
+    let msg_mode = "mode should be either forward or reverse";
+    let mode = match mode.as_str() {
+        //map(|x| x.as_str()) {
+        "Forward" => DiffMode::Forward,
+        "Reverse" => DiffMode::Reverse,
+        _ => {
+            tcx.sess
+                .struct_span_err(attr.span, msg_mode)
+                .span_label(attr.span, "invalid mode")
+                .emit();
+
+            return AutoDiffAttrs::inactive();
+        }
+    };
+
+    let msg_ret_activity = "autodiff attribute must contain the return activity";
+    let ret_symbol = match &list[1] {
+        NestedMetaItem::MetaItem(MetaItem { path: ref p2, kind: MetaItemKind::Word, .. }) => {
+            p2.segments.first().unwrap().ident
+        }
+        _ => {
+            tcx.sess
+                .struct_span_err(attr.span, msg_ret_activity)
+                .span_label(attr.span, "missing return activity")
+                .emit();
+
+            return AutoDiffAttrs::inactive();
+        }
+    };
+
+    let msg_unknown_ret_activity = "unknown return activity";
+    let ret_activity = match DiffActivity::from_str(ret_symbol.as_str()) {
+        Ok(x) => x,
+        Err(_) => {
+            tcx.sess
+                .struct_span_err(attr.span, msg_unknown_ret_activity)
+                .span_label(attr.span, "invalid return activity")
+                .emit();
+
+            return AutoDiffAttrs::inactive();
+        }
+    };
+
+    let msg_arg_activity = "autodiff attribute must contain the return activity";
+    let mut arg_activities: Vec<DiffActivity> = vec![];
+    for arg in &list[2..] {
+        let arg_symbol = match arg {
+            NestedMetaItem::MetaItem(MetaItem {
+                path: ref p2, kind: MetaItemKind::Word, ..
+            }) => p2.segments.first().unwrap().ident,
+            _ => {
+                tcx.sess
+                    .struct_span_err(
+                        attr.span, msg_arg_activity,
+                    )
+                    .span_label(attr.span, "missing return activity")
+                    .emit();
+
+                return AutoDiffAttrs::inactive();
+            }
+        };
+
+        match DiffActivity::from_str(arg_symbol.as_str()) {
+            Ok(arg_activity) => arg_activities.push(arg_activity),
+            Err(_) => {
+                tcx.sess
+                    .struct_span_err(attr.span, msg_unknown_ret_activity)
+                    .span_label(attr.span, "invalid input activity")
+                    .emit();
+
+                return AutoDiffAttrs::inactive();
+            }
+        }
+    }
+
+    let msg_fwd_incompatible_ret = "Forward Mode is incompatible with Active ret";
+    let msg_fwd_incompatible_arg = "Forward Mode is incompatible with Active ret";
+    let msg_rev_incompatible_arg = "Reverse Mode is only compatible with Active, None, or Const ret";
+    if mode == DiffMode::Forward {
+        if ret_activity == DiffActivity::Active {
+            tcx.sess
+                .struct_span_err(attr.span, msg_fwd_incompatible_ret)
+                .span_label(attr.span, "invalid return activity")
+                .emit();
+            return AutoDiffAttrs::inactive();
+        }
+        if arg_activities.iter().filter(|&x| *x == DiffActivity::Active).count() > 0 {
+            tcx.sess
+                .struct_span_err(attr.span, msg_fwd_incompatible_arg)
+                .span_label(attr.span, "invalid input activity")
+                .emit();
+            return AutoDiffAttrs::inactive();
+        }
+    }
+
+    if mode == DiffMode::Reverse {
+        if ret_activity == DiffActivity::Duplicated
+            || ret_activity == DiffActivity::DuplicatedNoNeed
+        {
+            tcx.sess
+                .struct_span_err(
+                    attr.span, msg_rev_incompatible_arg,
+                )
+                .span_label(attr.span, "invalid return activity")
+                .emit();
+            return AutoDiffAttrs::inactive();
+        }
+    }
+
+    AutoDiffAttrs { mode, ret_activity, input_activity: arg_activities }
+}
+
 pub fn provide(providers: &mut Providers) {
-    *providers = Providers { codegen_fn_attrs, should_inherit_track_caller, ..*providers };
+    *providers =
+        Providers { codegen_fn_attrs, should_inherit_track_caller, autodiff_attrs, ..*providers };
 }

compiler/rustc_middle/src/arena.rs

@@ -88,6 +88,7 @@ macro_rules! arena_types {
             [] upvars_mentioned: rustc_data_structures::fx::FxIndexMap<rustc_hir::HirId, rustc_hir::Upvar>,
             [] object_safety_violations: rustc_middle::traits::ObjectSafetyViolation,
             [] codegen_unit: rustc_middle::mir::mono::CodegenUnit<'tcx>,
+            [] autodiff_item: rustc_ast::expand::autodiff_attrs::AutoDiffItem,
             [decode] attribute: rustc_ast::Attribute,
             [] name_set: rustc_data_structures::unord::UnordSet<rustc_span::symbol::Symbol>,
             [] ordered_name_set: rustc_data_structures::fx::FxIndexSet<rustc_span::symbol::Symbol>,

compiler/rustc_middle/src/query/erase.rs

@@ -200,6 +200,10 @@ impl<T0, T1> EraseType for (&'_ T0, &'_ [T1]) {
     type Result = [u8; size_of::<(&'static (), &'static [()])>()];
 }
 
+impl<T0, T1, T2> EraseType for (&'_ T0, &'_ [T1], &'_ [T2]) {
+    type Result = [u8; size_of::<(&'static (), &'static [()], &'static [()])>()];
+}
+
 macro_rules! trivial {
     ($($ty:ty),+ $(,)?) => {
         $(

compiler/rustc_middle/src/query/mod.rs

@@ -55,6 +55,7 @@ use crate::ty::{GenericArg, GenericArgsRef};
 use rustc_arena::TypedArena;
 use rustc_ast as ast;
 use rustc_ast::expand::{allocator::AllocatorKind, StrippedCfgItem};
+use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem};
 use rustc_attr as attr;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::{FxHashMap, FxIndexMap, FxIndexSet};
@@ -1227,6 +1228,13 @@ rustc_queries! {
         separate_provide_extern
     }
 
+    /// The list autodiff extern functions in current crate
+    query autodiff_attrs(def_id: DefId) -> &'tcx AutoDiffAttrs {
+        desc { |tcx| "computing autodiff attributes of `{}`", tcx.def_path_str(def_id) }
+        arena_cache
+        cache_on_disk_if { def_id.is_local() }
+    }
+
     query asm_target_features(def_id: DefId) -> &'tcx FxIndexSet<Symbol> {
         desc { |tcx| "computing target features for inline asm of `{}`", tcx.def_path_str(def_id) }
     }
@@ -1880,7 +1888,7 @@ rustc_queries! {
         separate_provide_extern
     }
 
-    query collect_and_partition_mono_items(_: ()) -> (&'tcx DefIdSet, &'tcx [CodegenUnit<'tcx>]) {
+    query collect_and_partition_mono_items(_: ()) -> (&'tcx DefIdSet, &'tcx [AutoDiffItem], &'tcx [CodegenUnit<'tcx>]) {
         eval_always
         desc { "collect_and_partition_mono_items" }
     }

compiler/rustc_middle/src/ty/mod.rs

@@ -195,6 +195,8 @@ pub struct ResolverAstLowering {
     pub node_id_to_def_id: NodeMap<LocalDefId>,
     pub def_id_to_node_id: IndexVec<LocalDefId, ast::NodeId>,
 
+    /// Mapping of autodiff function IDs
+    pub autodiff_map: FxHashMap<LocalDefId, LocalDefId>,
     pub trait_map: NodeMap<Vec<hir::TraitCandidate>>,
     /// List functions and methods for which lifetime elision was successful.
     pub lifetime_elision_allowed: FxHashSet<ast::NodeId>,

compiler/rustc_monomorphize/Cargo.toml

@@ -5,6 +5,7 @@ edition = "2021"
 
 [dependencies]
 # tidy-alphabetical-start
+rustc_ast = { path = "../rustc_ast" }
 rustc_data_structures = { path = "../rustc_data_structures" }
 rustc_errors = { path = "../rustc_errors" }
 rustc_fluent_macro = { path = "../rustc_fluent_macro" }
@@ -13,6 +14,7 @@ rustc_macros = { path = "../rustc_macros" }
 rustc_middle = { path = "../rustc_middle" }
 rustc_session = { path = "../rustc_session" }
 rustc_span = { path = "../rustc_span" }
+rustc_symbol_mangling = { path = "../rustc_symbol_mangling" }
 rustc_target = { path = "../rustc_target" }
 serde = "1"
 serde_json = "1"