I will clean this up, I promise!

ZuseZ4 · ZuseZ4 · commit e3ba2d2e329b · 2024-01-24T22:39:19.000-05:00
diff --git a/compiler/rustc_ast/src/expand/autodiff_attrs.rs b/compiler/rustc_ast/src/expand/autodiff_attrs.rs
@@ -90,10 +90,13 @@ pub struct AutoDiffAttrs {
     pub input_activity: Vec<DiffActivity>,
 }
 
-fn name(x: &NestedMetaItem) -> String {
+fn first_ident(x: &NestedMetaItem) -> rustc_span::symbol::Ident {
     let segments = &x.meta_item().unwrap().path.segments;
     assert!(segments.len() == 1);
-    segments[0].ident.name.to_string()
+    segments[0].ident
+}
+fn name(x: &NestedMetaItem) -> String {
+    first_ident(x).name.to_string()
 }
 
 impl AutoDiffAttrs{
@@ -143,6 +146,13 @@ impl AutoDiffAttrs {
             input_activity: Vec::new(),
         }
     }
+    pub fn source() -> Self {
+        AutoDiffAttrs {
+            mode: DiffMode::Source,
+            ret_activity: DiffActivity::None,
+            input_activity: Vec::new(),
+        }
+    }
 
     pub fn is_active(&self) -> bool {
         match self.mode {
diff --git a/compiler/rustc_builtin_macros/src/autodiff.rs b/compiler/rustc_builtin_macros/src/autodiff.rs
@@ -1,5 +1,6 @@
 #![allow(unused_imports)]
 #![allow(unused_variables)]
+#![allow(unused_mut)]
 //use crate::util::check_builtin_macro_attribute;
 //use crate::util::check_autodiff;
 
@@ -9,12 +10,20 @@ use rustc_ast::ptr::P;
 use rustc_ast::{BindingAnnotation, ByRef};
 use rustc_ast::{self as ast, FnHeader, FnSig, Generics, StmtKind, NestedMetaItem, MetaItemKind};
 use rustc_ast::{Fn, ItemKind, Stmt, TyKind, Unsafe, PatKind};
+use rustc_ast::tokenstream::*;
 use rustc_expand::base::{Annotatable, ExtCtxt};
 use rustc_span::symbol::{kw, sym, Ident};
 use rustc_span::Span;
 use thin_vec::{thin_vec, ThinVec};
 use rustc_span::Symbol;
 use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode};
+use rustc_ast::token::{Token, TokenKind};
+
+fn first_ident(x: &NestedMetaItem) -> rustc_span::symbol::Ident {
+    let segments = &x.meta_item().unwrap().path.segments;
+    assert!(segments.len() == 1);
+    segments[0].ident
+}
 
 pub fn expand(
     ecx: &mut ExtCtxt<'_>,
@@ -33,7 +42,8 @@ pub fn expand(
             return vec![item];
         }
     };
-    let orig_item: P<ast::Item> = item.clone().expect_item();
+    let mut orig_item: P<ast::Item> = item.clone().expect_item();
+    //dbg!(&orig_item.tokens);
     let primal = orig_item.ident.clone();
 
     // Allow using `#[autodiff(...)]` only on a Fn
@@ -47,6 +57,25 @@ pub fn expand(
             .emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
         return vec![item];
     };
+    // create TokenStream from vec elemtents:
+    // meta_item doesn't have a .tokens field
+    let ts: Vec<Token> = meta_item_vec.clone()[1..].iter().map(|x| {
+        let val = first_ident(x);
+        let t = Token::from_ast_ident(val);
+        t
+    }).collect();
+    let comma: Token = Token::new(TokenKind::Comma, Span::default());
+    let mut ts: Vec<TokenTree> = vec![];
+    for t in meta_item_vec.clone()[1..].iter() {
+        let val = first_ident(t);
+        let t = Token::from_ast_ident(val);
+        ts.push(TokenTree::Token(t, Spacing::Joint));
+        ts.push(TokenTree::Token(comma.clone(), Spacing::Alone));
+    }
+    dbg!(&ts);
+    let ts: TokenStream = TokenStream::from_iter(ts);
+    dbg!(&ts);
+
     let x: AutoDiffAttrs = AutoDiffAttrs::from_ast(&meta_item_vec, has_ret);
     dbg!(&x);
     //let span = ecx.with_def_site_ctxt(sig_span);
@@ -66,7 +95,29 @@ pub fn expand(
         generics: Generics::default(),
         body: Some(d_body),
     }));
-    let d_fn = ecx.item(span, d_ident, rustc_ast::AttrVec::default(), asdf);
+    let mut tmp = P(ast::NormalAttr::from_ident(Ident::with_dummy_span(sym::autodiff_into)));
+    let mut attr: ast::Attribute = ast::Attribute {
+        kind: ast::AttrKind::Normal(tmp.clone()),
+        id: ast::AttrId::from_u32(0),
+        style: ast::AttrStyle::Outer,
+        span: span,
+    };
+    orig_item.attrs.push(attr);
+
+    // Now update for d_fn
+    tmp.item.args = rustc_ast::AttrArgs::Delimited(rustc_ast::DelimArgs {
+        dspan: DelimSpan::dummy(),
+        delim: rustc_ast::token::Delimiter::Parenthesis,
+        tokens: ts,
+    });
+    let mut attr2: ast::Attribute = ast::Attribute {
+        kind: ast::AttrKind::Normal(tmp),
+        id: ast::AttrId::from_u32(0),
+        style: ast::AttrStyle::Outer,
+        span: span,
+    };
+    let attr_vec: rustc_ast::AttrVec = thin_vec![attr2];
+    let d_fn = ecx.item(span, d_ident, attr_vec, asdf);
 
     let orig_annotatable = Annotatable::Item(orig_item.clone());
     let d_annotatable = Annotatable::Item(d_fn);
diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs
@@ -269,6 +269,7 @@ impl WriteBackendMethods for LlvmCodegenBackend {
         config: &ModuleConfig,
     ) -> Result<(), FatalError> {
         dbg!("cg_llvm autodiff");
+        dbg!("Differentiating {} functions", diff_fncs.len());
         unsafe { back::write::differentiate(module, cgcx, diff_fncs, typetrees, config) }
     }
 
diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -382,6 +382,7 @@ fn generate_lto_work<B: ExtraBackendMethods>(
     import_only_modules: Vec<(SerializedModule<B::ModuleBuffer>, WorkProduct)>,
 ) -> Vec<(WorkItem<B>, u64)> {
     let _prof_timer = cgcx.prof.generic_activity("codegen_generate_lto_work");
+    dbg!("Differentiating {} functions", autodiff.len());
 
     if !needs_fat_lto.is_empty() {
         assert!(needs_thin_lto.is_empty());
diff --git a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
@@ -693,16 +693,15 @@ 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 = tcx.get_attrs(id, sym::autodiff_into);
 
     let attrs = attrs
         .into_iter()
-        .filter(|attr| attr.name_or_empty() == sym::autodiff)
-        //.filter(|attr| attr.name_or_empty() == sym::autodiff_into)
+        .filter(|attr| attr.name_or_empty() == sym::autodiff_into)
         .collect::<Vec<_>>();
-    if attrs.len() > 0 {
-      dbg!("autodiff_attrs len = > 0: {}", attrs.len());
+
+    if !attrs.is_empty() {
+        dbg!("autodiff_attrs amount = {}", attrs.len());
     }
 
     // check for exactly one autodiff attribute on extern block
@@ -723,18 +722,12 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
     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(),
-        };
-    }
+    if list.len() == 0 { return AutoDiffAttrs::source(); }
 
     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
+    let (mode, list) = match list.split_first() {
+        Some((NestedMetaItem::MetaItem(MetaItem { path: ref p1, kind: MetaItemKind::Word, .. }), list)) => {
+            (p1.segments.first().unwrap().ident, list)
         }
         _ => {
             tcx.sess
@@ -749,7 +742,6 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
     // 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,
         _ => {
@@ -763,9 +755,9 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
     };
 
     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
+    let (ret_symbol, list) = match list.split_last() {
+        Some((NestedMetaItem::MetaItem(MetaItem { path: ref p1, kind: MetaItemKind::Word, .. }), list)) => {
+            (p1.segments.first().unwrap().ident, list)
         }
         _ => {
             tcx.sess
@@ -792,7 +784,7 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
 
     let msg_arg_activity = "autodiff attribute must contain the return activity";
     let mut arg_activities: Vec<DiffActivity> = vec![];
-    for arg in &list[2..] {
+    for arg in list {
         let arg_symbol = match arg {
             NestedMetaItem::MetaItem(MetaItem {
                 path: ref p2, kind: MetaItemKind::Word, ..
@@ -846,6 +838,7 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> AutoDiffAttrs {
         if ret_activity == DiffActivity::Duplicated
             || ret_activity == DiffActivity::DuplicatedNoNeed
         {
+            dbg!("ret_activity = {:?}", ret_activity);
             tcx.sess
                 .struct_span_err(
                     attr.span, msg_rev_incompatible_arg,
diff --git a/compiler/rustc_feature/src/builtin_attrs.rs b/compiler/rustc_feature/src/builtin_attrs.rs
@@ -361,6 +361,13 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         large_assignments, experimental!(move_size_limit)
     ),
 
+    // Autodiff
+    ungated!(
+        autodiff_into, Normal,
+        template!(Word, List: r#""...""#),
+        DuplicatesOk,
+    ),
+
     // Entry point:
     gated!(unix_sigpipe, Normal, template!(Word, NameValueStr: "inherit|sig_ign|sig_dfl"), ErrorFollowing, experimental!(unix_sigpipe)),
     ungated!(start, Normal, template!(Word), WarnFollowing),
diff --git a/compiler/rustc_monomorphize/src/partitioning.rs b/compiler/rustc_monomorphize/src/partitioning.rs
@@ -1184,7 +1184,7 @@ fn collect_and_partition_mono_items(tcx: TyCtxt<'_>, (): ()) -> (&DefIdSet, &[Au
             println!("item: {:?}", item);
             let source = usage_map.used_map.get(&item).unwrap()
                 .into_iter()
-                .filter_map(|item| match *item {
+                .find_map(|item| match *item {
                     MonoItem::Fn(ref instance_s) => {
                         let source_id = instance_s.def_id();
                         println!("source_id_inner: {:?}", source_id);
@@ -1205,8 +1205,8 @@ fn collect_and_partition_mono_items(tcx: TyCtxt<'_>, (): ()) -> (&DefIdSet, &[Au
                         None
                     }
                     _ => None,
-                })
-                .next();
+                });
+                //.next();
             println!("source: {:?}", source);
 
             source.map(|inst| {
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
@@ -440,6 +440,7 @@ symbols! {
         augmented_assignments,
         auto_traits,
         autodiff,
+        autodiff_into,
         automatically_derived,
         avx,
         avx512_target_feature,

Original file line number	Diff line number	Diff line change
`@@ -269,6 +269,7 @@ impl WriteBackendMethods for LlvmCodegenBackend {`
`269`	`269`	`config: &ModuleConfig,`
`270`	`270`	`) -> Result<(), FatalError> {`
`271`	`271`	`dbg!("cg_llvm autodiff");`
	`272`	`+ dbg!("Differentiating {} functions", diff_fncs.len());`
`272`	`273`	`unsafe { back::write::differentiate(module, cgcx, diff_fncs, typetrees, config) }`
`273`	`274`	`}`
`274`	`275`