addressing a ton of feedback, thank you

ZuseZ4 · ZuseZ4 · commit 950d91b7e5eb · 2025-01-26T20:15:26.000-05:00
diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs
@@ -256,9 +256,9 @@ fn generate_enzyme_call<'ll>(
             trace!("no dbg info");
         }
         // Dump module:
-        llvm::LLVMDumpModule(cx.llmod);
+        //llvm::LLVMDumpModule(cx.llmod);
         // now print the last instruction:
-        llvm::LLVMDumpValue(last_inst);
+        //llvm::LLVMDumpValue(last_inst);
 
         // Now that we copied the metadata, get rid of dummy code.
         llvm::LLVMRustEraseInstBefore(entry, last_inst);
diff --git a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
@@ -18,6 +18,7 @@ use rustc_middle::middle::codegen_fn_attrs::{
 };
 use rustc_middle::mir::mono::Linkage;
 use rustc_middle::query::Providers;
+use rustc_middle::span_bug;
 use rustc_middle::ty::{self as ty, TyCtxt};
 use rustc_session::parse::feature_err;
 use rustc_session::{Session, lint};
@@ -872,114 +873,89 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
 
     // check for exactly one autodiff attribute on placeholder functions.
     // There should only be one, since we generate a new placeholder per ad macro.
-    // TODO: re-enable this. We should fix that rustc_autodiff isn't applied multiple times to the
-    // source function.
-    let msg_once = "cg_ssa: implementation bug. Autodiff attribute can only be applied once";
+    // 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.len() {
         0 => return AutoDiffAttrs::error(),
         1 => attrs.get(0).unwrap(),
         _ => {
             attrs.get(0).unwrap()
-            //tcx.dcx().struct_span_err(attrs[1].span, msg_once).with_note("more than one").emit();
-            //return AutoDiffAttrs::error();
+            //FIXME(ZuseZ4): re-enable this check
+            //span_bug!(attrs[1].span, "cg_ssa: rustc_autodiff should only exist once per source");
         }
     };
 
     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 {
+    if list.is_empty() {
         return AutoDiffAttrs::source();
     }
 
     let [mode, input_activities @ .., ret_activity] = &list[..] else {
-        tcx.dcx()
-            .struct_span_err(attr.span, msg_once)
-            .with_note("Implementation bug in autodiff_attrs. Please report this!")
-            .emit();
-        return AutoDiffAttrs::error();
+        span_bug!(attr.span, "rustc_autodiff attribute must contain mode and activities");
     };
     let mode = if let MetaItemInner::MetaItem(MetaItem { path: ref p1, .. }) = mode {
         p1.segments.first().unwrap().ident
     } else {
-        let msg = "autodiff attribute must contain autodiff mode";
-        tcx.dcx().struct_span_err(attr.span, msg).with_note("empty argument list").emit();
-        return AutoDiffAttrs::error();
+        span_bug!(attr.span, "rustc_autodiff attribute must contain mode");
     };
 
     // parse mode
-    let msg_mode = "mode should be either forward or reverse";
     let mode = match mode.as_str() {
         "Forward" => DiffMode::Forward,
         "Reverse" => DiffMode::Reverse,
         "ForwardFirst" => DiffMode::ForwardFirst,
         "ReverseFirst" => DiffMode::ReverseFirst,
         _ => {
-            tcx.dcx().struct_span_err(attr.span, msg_mode).with_note("invalid mode").emit();
-            return AutoDiffAttrs::error();
+            span_bug!(attr.span, "rustc_autodiff attribute contains invalid mode");
         }
     };
 
     // First read the ret symbol from the attribute
     let ret_symbol = if let MetaItemInner::MetaItem(MetaItem { path: ref p1, .. }) = ret_activity {
         p1.segments.first().unwrap().ident
     } else {
-        let msg = "autodiff attribute must contain the return activity";
-        tcx.dcx().struct_span_err(attr.span, msg).with_note("missing return activity").emit();
-        return AutoDiffAttrs::error();
+        span_bug!(attr.span, "rustc_autodiff attribute must contain the return activity");
     };
 
     // Then parse it into an actual DiffActivity
-    let msg_unknown_ret_activity = "unknown return activity";
-    let ret_activity = match DiffActivity::from_str(ret_symbol.as_str()) {
-        Ok(x) => x,
-        Err(_) => {
-            tcx.dcx()
-                .struct_span_err(attr.span, msg_unknown_ret_activity)
-                .with_note("invalid return activity")
-                .emit();
-            return AutoDiffAttrs::error();
-        }
+    let Ok(ret_activity) = DiffActivity::from_str(ret_symbol.as_str()) else {
+        span_bug!(attr.span, "invalid return activity");
     };
 
     // Now parse all the intermediate (input) activities
-    let msg_arg_activity = "autodiff attribute must contain the return activity";
     let mut arg_activities: Vec<DiffActivity> = vec![];
     for arg in input_activities {
         let arg_symbol = if let MetaItemInner::MetaItem(MetaItem { path: ref p2, .. }) = arg {
-            p2.segments.first().unwrap().ident
+            match p2.segments.first() {
+                Some(x) => x.ident,
+                None => {
+                    span_bug!(
+                        attr.span,
+                        "rustc_autodiff attribute must contain the input activity"
+                    );
+                }
+            }
         } else {
-            tcx.dcx()
-                .struct_span_err(attr.span, msg_arg_activity)
-                .with_note("Implementation bug, please report this!")
-                .emit();
-            return AutoDiffAttrs::error();
+            span_bug!(attr.span, "rustc_autodiff attribute must contain the input activity");
         };
 
         match DiffActivity::from_str(arg_symbol.as_str()) {
             Ok(arg_activity) => arg_activities.push(arg_activity),
             Err(_) => {
-                tcx.dcx()
-                    .struct_span_err(attr.span, msg_unknown_ret_activity)
-                    .with_note("invalid input activity")
-                    .emit();
-                return AutoDiffAttrs::error();
+                span_bug!(attr.span, "invalid input activity");
             }
         }
     }
 
-    let mut msg = "".to_string();
     for &input in &arg_activities {
         if !valid_input_activity(mode, input) {
-            msg = format!("Invalid input activity {} for {} mode", input, mode);
+            span_bug!(attr.span, "Invalid input activity {} for {} mode", input, mode);
         }
     }
     if !valid_ret_activity(mode, ret_activity) {
-        msg = format!("Invalid return activity {} for {} mode", ret_activity, mode);
-    }
-    if msg != "".to_string() {
-        tcx.dcx().struct_span_err(attr.span, msg).with_note("invalid activity").emit();
-        return AutoDiffAttrs::error();
+        span_bug!(attr.span, "Invalid return activity {} for {} mode", ret_activity, mode);
     }
 
     AutoDiffAttrs { mode, ret_activity, input_activity: arg_activities }
diff --git a/compiler/rustc_codegen_ssa/src/traits/write.rs b/compiler/rustc_codegen_ssa/src/traits/write.rs
@@ -13,7 +13,6 @@ pub trait WriteBackendMethods: 'static + Sized + Clone {
     type ModuleBuffer: ModuleBufferMethods;
     type ThinData: Send + Sync;
     type ThinBuffer: ThinBufferMethods;
-    //type TypeTree: Clone;
 
     /// Merge all modules into main_module and returning it
     fn run_link(
@@ -38,7 +37,6 @@ pub trait WriteBackendMethods: 'static + Sized + Clone {
     ) -> Result<(Vec<LtoModuleCodegen<Self>>, Vec<WorkProduct>), FatalError>;
     fn print_pass_timings(&self);
     fn print_statistics(&self);
-    // does enzyme prep work, should do ad too.
     unsafe fn optimize(
         cgcx: &CodegenContext<Self>,
         dcx: DiagCtxtHandle<'_>,
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
@@ -1393,7 +1393,7 @@ rustc_queries! {
         feedable
     }
 
-    /// The list autodiff extern functions in current crate
+    /// 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
diff --git a/compiler/rustc_monomorphize/src/partitioning.rs b/compiler/rustc_monomorphize/src/partitioning.rs
@@ -255,9 +255,13 @@ where
         );
 
         // We can't differentiate something that got inlined.
-        let autodiff_active = match characteristic_def_id {
-            Some(def_id) => cx.tcx.autodiff_attrs(def_id).is_active(),
-            None => false,
+        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
         };
 
         if !autodiff_active && visibility == Visibility::Hidden && can_be_internalized {
@@ -1128,9 +1132,8 @@ pub(crate) fn adjust_activity_to_abi<'tcx>(
     fn_ty: Ty<'tcx>,
     da: &mut Vec<DiffActivity>,
 ) {
-    if !fn_ty.is_fn() {
-        // Error?
-        return;
+    if !matches!(fn_ty.kind(), ty::FnDef(..)) {
+        bug!("expected fn def for autodiff, got {:?}", fn_ty);
     }
     let fnc_binder: ty::Binder<'_, ty::FnSig<'_>> = fn_ty.fn_sig(tcx);
 
@@ -1149,7 +1152,7 @@ pub(crate) fn adjust_activity_to_abi<'tcx>(
             }
             let inner_ty = ty.builtin_deref(true).unwrap();
             if inner_ty.is_slice() {
-                // We know that the lenght will be passed as extra arg.
+                // We know that the length will be passed as extra arg.
                 if !da.is_empty() {
                     // We are looking at a slice. The length of that slice will become an
                     // extra integer on llvm level. Integers are always const.
@@ -1161,12 +1164,11 @@ pub(crate) fn adjust_activity_to_abi<'tcx>(
                         | DiffActivity::DuplicatedOnly
                         | DiffActivity::Duplicated => DiffActivity::FakeActivitySize,
                         DiffActivity::Const => DiffActivity::Const,
-                        _ => panic!("unexpected activity for ptr/ref"),
+                        _ => bug!("unexpected activity for ptr/ref"),
                     };
                     new_activities.push(activity);
                     new_positions.push(i + 1);
                 }
-                trace!("ABI MATCHING!");
                 continue;
             }
         }
@@ -1245,7 +1247,7 @@ fn collect_and_partition_mono_items(
         })
         .collect();
 
-    let autodiff_items2: Vec<_> = items
+    let autodiff_mono_items: Vec<_> = items
         .iter()
         .filter_map(|item| match *item {
             MonoItem::Fn(ref instance) => Some((item, instance)),
@@ -1254,7 +1256,7 @@ fn collect_and_partition_mono_items(
         .collect();
     let mut autodiff_items: Vec<AutoDiffItem> = vec![];
 
-    for (item, instance) in autodiff_items2 {
+    for (item, instance) in autodiff_mono_items {
         let target_id = instance.def_id();
         let target_attrs: &AutoDiffAttrs = tcx.autodiff_attrs(target_id);
         let mut input_activities: Vec<DiffActivity> = target_attrs.input_activity.clone();
@@ -1283,18 +1285,15 @@ fn collect_and_partition_mono_items(
             None => continue,
         };
 
-        println!("source_id: {:?}", inst.def_id());
+        debug!("source_id: {:?}", inst.def_id());
         let fn_ty = inst.ty(tcx, ty::TypingEnv::fully_monomorphized());
-        //let fn_ty = inst.ty(tcx, ParamEnv::empty());ty::TypingEnv<'tcx>::fully_monomorphized()
         assert!(fn_ty.is_fn());
         adjust_activity_to_abi(tcx, fn_ty, &mut input_activities);
-        //let (inputs, output) = (fnc_tree.args, fnc_tree.ret);
         let symb = symbol_name_for_instance_in_crate(tcx, inst.clone(), LOCAL_CRATE);
 
         let mut new_target_attrs = target_attrs.clone();
         new_target_attrs.input_activity = input_activities;
         let itm = new_target_attrs.into_item(symb, target_symbol);
-        //let itm = new_target_attrs.into_item(symb, target_symbol, inputs, output);
         autodiff_items.push(itm);
     }
     let autodiff_items = tcx.arena.alloc_from_iter(autodiff_items);
@@ -1359,7 +1358,7 @@ fn collect_and_partition_mono_items(
         }
     }
 
-    if autodiff_items.len() > 0 {
+    if !autodiff_items.is_empty() {
         trace!("AUTODIFF ITEMS EXIST");
         for item in &mut *autodiff_items {
             trace!("{}", &item);
diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs
@@ -195,7 +195,7 @@ pub enum CoverageLevel {
 //    }
 //}
 
-/// The different settings that the `-Z ad` flag can have.
+/// The different settings that the `-Z autodiff` flag can have.
 #[derive(Clone, Copy, PartialEq, Hash, Debug)]
 pub enum AutoDiff {
     /// Print TypeAnalysis information
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
@@ -398,7 +398,7 @@ mod desc {
     pub(crate) const parse_list: &str = "a space-separated list of strings";
     pub(crate) const parse_list_with_polarity: &str =
         "a comma-separated list of strings, with elements beginning with + or -";
-    pub(crate) const parse_autodiff: &str = "various values";
+    pub(crate) const parse_autodiff: &str = "a comma separated list of settings: `Print`, `PrintTA`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfterOpts`, `PrintModAfterEnzyme`, `LooseTypes`, `NoModOptAfter`, `EnableFncOpt`, `NoVecUnroll`, `Inline`";
     pub(crate) const parse_comma_list: &str = "a comma-separated list of strings";
     pub(crate) const parse_opt_comma_list: &str = parse_comma_list;
     pub(crate) const parse_number: &str = "a number";
@@ -1051,7 +1051,10 @@ pub mod parse {
                 "EnableFncOpt" => AutoDiff::EnableFncOpt,
                 "NoVecUnroll" => AutoDiff::NoVecUnroll,
                 "Inline" => AutoDiff::Inline,
-                _ => return false,
+                _ => {
+                    // FIXME(ZuseZ4): print an error saying which value is not recognized
+                    return false;
+                }
             };
             slot.push(variant);
         }
@@ -1767,7 +1770,21 @@ options! {
     assume_incomplete_release: bool = (false, parse_bool, [TRACKED],
         "make cfg(version) treat the current version as incomplete (default: no)"),
     autodiff: Vec<crate::config::AutoDiff> = (Vec::new(), parse_autodiff, [TRACKED],
-        "a list autodiff flags to enable (comma separated)"),
+        "a list of optional autodiff flags to enable
+         Optional extra settings:
+         `=PrintTA`
+         `=PrintAA`
+         `=PrintPerf`
+         `=Print`
+         `=PrintModBefore`
+         `=PrintModAfterOpts`
+         `=PrintModAfterEnzyme`
+         `=LooseTypes`
+         `=NoModOptAfter`
+         `=EnableFncOpt`
+         `=NoVecUnroll`
+         `=Inline`
+         Multiple options can be combined with commas."),
     #[rustc_lint_opt_deny_field_access("use `Session::binary_dep_depinfo` instead of this field")]
     binary_dep_depinfo: bool = (false, parse_bool, [TRACKED],
         "include artifacts (sysroot, crate dependencies) used during compilation in dep-info \
diff --git a/src/doc/unstable-book/src/compiler-flags/autodiff.md b/src/doc/unstable-book/src/compiler-flags/autodiff.md
@@ -0,0 +1,9 @@
+# `autodiff`
+
+The tracking issue for this feature is: [#124509](https://github.com/rust-lang/rust/issues/124509).
+
+------------------------
+
+This feature allows you to differentiate functions using automatic differentiation.
+Set the `-Zautodiff=<options>` compiler flag to adjust the behaviour of the autodiff feature.
+

Original file line number	Diff line number	Diff line change
`@@ -1393,7 +1393,7 @@ rustc_queries! {`
`1393`	`1393`	`feedable`
`1394`	`1394`	`}`
`1395`	`1395`
`1396`		`- /// The list autodiff extern functions in current crate`
	`1396`	`+ /// List of autodiff extern functions in the current crate.`
`1397`	`1397`	`query autodiff_attrs(def_id: DefId) -> &'tcx AutoDiffAttrs {`
`1398`	`1398`	desc { \|tcx\| "computing autodiff attributes of `{}`", tcx.def_path_str(def_id) }
`1399`	`1399`	`arena_cache`