wip

Author: ZuseZ4

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -79,6 +79,7 @@ pub struct AutoDiffItem {
     pub target: String,
     pub attrs: AutoDiffAttrs,
 }
+
 #[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
 pub struct AutoDiffAttrs {
     /// Conceptually either forward or reverse mode AD, as described in various autodiff papers and
@@ -231,7 +232,7 @@ impl AutoDiffAttrs {
         self.ret_activity == DiffActivity::ActiveOnly
     }
 
-    pub fn error() -> Self {
+    pub const fn error() -> Self {
         AutoDiffAttrs {
             mode: DiffMode::Error,
             ret_activity: DiffActivity::None,

compiler/rustc_codegen_ssa/src/base.rs

@@ -2,7 +2,6 @@ use std::cmp;
 use std::collections::BTreeSet;
 use std::time::{Duration, Instant};
 
-use rustc_middle::mir::mono::MonoItemPartitions;
 use itertools::Itertools;
 use rustc_abi::FIRST_VARIANT;
 use rustc_ast::expand::allocator::{ALLOCATOR_METHODS, AllocatorKind, global_fn_name};
@@ -19,7 +18,7 @@ use rustc_middle::middle::debugger_visualizer::{DebuggerVisualizerFile, Debugger
 use rustc_middle::middle::exported_symbols::SymbolExportKind;
 use rustc_middle::middle::{exported_symbols, lang_items};
 use rustc_middle::mir::BinOp;
-use rustc_middle::mir::mono::{CodegenUnit, CodegenUnitNameBuilder, MonoItem};
+use rustc_middle::mir::mono::{CodegenUnit, CodegenUnitNameBuilder, MonoItem, MonoItemPartitions};
 use rustc_middle::query::Providers;
 use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutOf, TyAndLayout};
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
@@ -620,7 +619,8 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
 
     // Run the monomorphization collector and partition the collected items into
     // codegen units.
-    let MonoItemPartitions {codegen_units, autodiff_items, ..} = tcx.collect_and_partition_mono_items(());
+    let MonoItemPartitions { codegen_units, autodiff_items, .. } =
+        tcx.collect_and_partition_mono_items(());
     let autodiff_fncs = autodiff_items.to_vec();
 
     // Force all codegen_unit queries so they are already either red or green

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -71,6 +71,13 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
         codegen_fn_attrs.flags |= CodegenFnAttrFlags::TRACK_CALLER;
     }
 
+    // If our rustc version supports autodiff/enzyme, then we call our handler
+    // to check for any `#[rustc_autodiff(...)]` attributes.
+    //if cfg!(llvm_enzyme) {
+    //    let ad = autodiff_attrs(tcx, did.into());
+    //    codegen_fn_attrs.autodiff_item = ad;
+    //}
+
     // When `no_builtins` is applied at the crate level, we should add the
     // `no-builtins` attribute to each function to ensure it takes effect in LTO.
     let crate_attrs = tcx.hir().attrs(rustc_hir::CRATE_HIR_ID);
@@ -867,7 +874,7 @@ impl<'a> MixedExportNameAndNoMangleState<'a> {
 /// being differentiated). The other form is #[rustc_autodiff(Mode, ActivityList)] on top of the
 /// placeholder functions. We wrote the rustc_autodiff attributes ourself, so this should never
 /// panic, unless we introduced a bug when parsing the autodiff macro.
-fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
+fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
     let attrs = tcx.get_attrs(id, sym::rustc_autodiff);
 
     let attrs =
@@ -878,7 +885,7 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
     // FIXME(ZuseZ4): re-enable this check. Currently we add multiple, which doesn't cause harm but
     // looks strange e.g. under cargo-expand.
     let attr = match &attrs[..] {
-        [] => return AutoDiffAttrs::error(),
+        [] => return None,
         [attr] => attr,
         // These two attributes are the same and unfortunately duplicated due to a previous bug.
         [attr, _attr2] => attr,
@@ -893,7 +900,7 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
 
     // empty autodiff attribute macros (i.e. `#[autodiff]`) are used to mark source functions
     if list.is_empty() {
-        return AutoDiffAttrs::source();
+        return Some(AutoDiffAttrs::source());
     }
 
     let [mode, input_activities @ .., ret_activity] = &list[..] else {
@@ -962,10 +969,10 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
         span_bug!(attr.span, "Invalid return activity {} for {} mode", ret_activity, mode);
     }
 
-    AutoDiffAttrs { mode, ret_activity, input_activity: arg_activities }
+    Some(AutoDiffAttrs { mode, ret_activity, input_activity: arg_activities })
 }
 
 pub(crate) fn provide(providers: &mut Providers) {
     *providers =
-        Providers { codegen_fn_attrs, should_inherit_track_caller, autodiff_attrs, ..*providers };
+        Providers { codegen_fn_attrs, should_inherit_track_caller, ..*providers };
 }

compiler/rustc_middle/src/middle/codegen_fn_attrs.rs

@@ -1,4 +1,5 @@
 use rustc_abi::Align;
+use rustc_ast::expand::autodiff_attrs::AutoDiffAttrs;
 use rustc_attr_parsing::{InlineAttr, InstructionSetAttr, OptimizeAttr};
 use rustc_macros::{HashStable, TyDecodable, TyEncodable};
 use rustc_span::Symbol;
@@ -52,6 +53,8 @@ pub struct CodegenFnAttrs {
     /// The `#[patchable_function_entry(...)]` attribute. Indicates how many nops should be around
     /// the function entry.
     pub patchable_function_entry: Option<PatchableFunctionEntry>,
+    /// For the `#[autodiff]` macros.
+    pub autodiff_item: Option<AutoDiffAttrs>,
 }
 
 #[derive(Copy, Clone, Debug, TyEncodable, TyDecodable, HashStable)]
@@ -160,6 +163,7 @@ impl CodegenFnAttrs {
             instruction_set: None,
             alignment: None,
             patchable_function_entry: None,
+            autodiff_item: None,
         }
     }
 

compiler/rustc_middle/src/query/mod.rs

@@ -13,7 +13,7 @@ use std::sync::Arc;
 use rustc_arena::TypedArena;
 use rustc_ast::expand::StrippedCfgItem;
 use rustc_ast::expand::allocator::AllocatorKind;
-use rustc_ast::expand::autodiff_attrs::AutoDiffAttrs;
+//use rustc_ast::expand::autodiff_attrs::AutoDiffAttrs;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
 use rustc_data_structures::sorted_map::SortedMap;
@@ -1394,11 +1394,11 @@ rustc_queries! {
     }
 
     /// List of autodiff extern functions in the 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 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) }

compiler/rustc_monomorphize/src/partitioning.rs

@@ -98,7 +98,7 @@ use std::fs::{self, File};
 use std::io::Write;
 use std::path::{Path, PathBuf};
 
-use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivity};
+use rustc_ast::expand::autodiff_attrs::{AutoDiffItem, DiffActivity};
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
 use rustc_data_structures::sync;
 use rustc_data_structures::unord::{UnordMap, UnordSet};
@@ -254,19 +254,20 @@ where
             always_export_generics,
         );
 
+        // TODO: This currently crashes compilation:
+        // thread 'rustc' panicked at compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs:242:1:
+        // DefId(2:45681 ~ core[1460]::iter::adapters::copied::Copied) does not have a "codegen_fn_attrs"
+        //
         // We can't differentiate something that got inlined.
-        let autodiff_active = if cfg!(llvm_enzyme) {
-            match characteristic_def_id {
-                Some(def_id) => cx.tcx.autodiff_attrs(def_id).is_active(),
-                None => false,
-            }
-        } else {
-            false
+        let autodiff_active = cfg!(llvm_enzyme) && {
+            characteristic_def_id
+                .and_then(|def_id| cx.tcx.codegen_fn_attrs(def_id).autodiff_item.as_ref())
+                .is_some_and(|ad| ad.is_active())
         };
 
-        if !autodiff_active && visibility == Visibility::Hidden && can_be_internalized {
-            internalization_candidates.insert(mono_item);
-        }
+        //if !autodiff_active && visibility == Visibility::Hidden && can_be_internalized {
+        //    internalization_candidates.insert(mono_item);
+        //}
         let size_estimate = mono_item.size_estimate(cx.tcx);
 
         cgu.items_mut().insert(mono_item, MonoItemData {
@@ -1254,7 +1255,11 @@ fn collect_and_partition_mono_items(tcx: TyCtxt<'_>, (): ()) -> MonoItemPartitio
 
     for (item, instance) in autodiff_mono_items {
         let target_id = instance.def_id();
-        let target_attrs: &AutoDiffAttrs = tcx.autodiff_attrs(target_id);
+        let cg_fn_attr = tcx.codegen_fn_attrs(target_id).autodiff_item.clone();
+        let Some(target_attrs) = cg_fn_attr else {
+            continue;
+        };
+        //let target_attrs: &AutoDiffAttrs = tcx.autodiff_attrs(target_id);
         let mut input_activities: Vec<DiffActivity> = target_attrs.input_activity.clone();
         if target_attrs.is_source() {
             trace!("source found: {:?}", target_id);
@@ -1269,7 +1274,8 @@ fn collect_and_partition_mono_items(tcx: TyCtxt<'_>, (): ()) -> MonoItemPartitio
             usage_map.used_map.get(&item).unwrap().into_iter().find_map(|item| match *item {
                 MonoItem::Fn(ref instance_s) => {
                     let source_id = instance_s.def_id();
-                    if tcx.autodiff_attrs(source_id).is_active() {
+                    //if tcx.autodiff_attrs(source_id).is_active() {
+                    if let Some(ad) = &tcx.codegen_fn_attrs(source_id).autodiff_item && ad.is_active() {
                         return Some(instance_s);
                     }
                     None