add minimal ast based macro parser

ZuseZ4 · ZuseZ4 · commit 670054202ad1 · 2024-01-15T18:02:02.000-05:00
diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
@@ -2574,6 +2574,12 @@ impl FnRetTy {
             FnRetTy::Ty(ty) => ty.span,
         }
     }
+    pub fn has_ret(&self) -> bool {
+        match self {
+            FnRetTy::Default(_) => false,
+            FnRetTy::Ty(_) => true,
+        }
+    }
 }
 
 #[derive(Clone, Copy, PartialEq, Encodable, Decodable, Debug)]
diff --git a/compiler/rustc_ast/src/expand/autodiff_attrs.rs b/compiler/rustc_ast/src/expand/autodiff_attrs.rs
@@ -1,7 +1,11 @@
-use super::typetree::TypeTree;
-use std::str::FromStr;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};//, StableOrd};
 use crate::HashStableContext;
+use crate::expand::typetree::TypeTree;
+use thin_vec::ThinVec;
+//use rustc_expand::base::{Annotatable, ExtCtxt};
+use std::str::FromStr;
+
+use crate::NestedMetaItem;
 
 #[allow(dead_code)]
 #[derive(Clone, Copy, Eq, PartialEq, Encodable, Decodable, Debug)]
@@ -52,12 +56,22 @@ impl<CTX: HashStableContext> HashStable<CTX> for DiffActivity {
     }
 }
 
+impl FromStr for DiffMode {
+    type Err = ();
 
+    fn from_str(s: &str) -> Result<DiffMode, ()> { match s {
+            "Inactive" => Ok(DiffMode::Inactive),
+            "Source" => Ok(DiffMode::Source),
+            "Forward" => Ok(DiffMode::Forward),
+            "Reverse" => Ok(DiffMode::Reverse),
+            _ => Err(()),
+        }
+    }
+}
 impl FromStr for DiffActivity {
     type Err = ();
 
-    fn from_str(s: &str) -> Result<DiffActivity, ()> {
-        match s {
+    fn from_str(s: &str) -> Result<DiffActivity, ()> { match s {
             "None" => Ok(DiffActivity::None),
             "Active" => Ok(DiffActivity::Active),
             "Const" => Ok(DiffActivity::Const),
@@ -76,6 +90,43 @@ pub struct AutoDiffAttrs {
     pub input_activity: Vec<DiffActivity>,
 }
 
+fn name(x: &NestedMetaItem) -> String {
+    let segments = &x.meta_item().unwrap().path.segments;
+    assert!(segments.len() == 1);
+    segments[0].ident.name.to_string()
+}
+
+impl AutoDiffAttrs{
+    pub fn has_ret_activity(&self) -> bool {
+        match self.ret_activity {
+            DiffActivity::None => false,
+            _ => true,
+        }
+    }
+    pub fn from_ast(meta_item: &ThinVec<NestedMetaItem>, has_ret: bool) -> Self {
+        let mode = name(&meta_item[1]);
+        let mode = DiffMode::from_str(&mode).unwrap();
+        let activities: Vec<DiffActivity> = meta_item[2..].iter().map(|x| {
+            let activity_str = name(&x);
+            DiffActivity::from_str(&activity_str).unwrap()
+        }).collect();
+
+        // If a return type exist, we need to split the last activity,
+        // otherwise we return None as placeholder.
+        let (ret_activity, input_activity) = if has_ret {
+            activities.split_last().unwrap()
+        } else {
+            (&DiffActivity::None, activities.as_slice())
+        };
+
+        AutoDiffAttrs {
+            mode,
+            ret_activity: *ret_activity,
+            input_activity: input_activity.to_vec(),
+        }
+    }
+}
+
 impl<CTX: HashStableContext> HashStable<CTX> for AutoDiffAttrs {
     fn hash_stable(&self, hcx: &mut CTX, hasher: &mut StableHasher) {
         self.mode.hash_stable(hcx, hasher);
diff --git a/compiler/rustc_ast/src/lib.rs b/compiler/rustc_ast/src/lib.rs
@@ -50,6 +50,7 @@ pub mod ptr;
 pub mod token;
 pub mod tokenstream;
 pub mod visit;
+//pub mod autodiff_attrs;
 
 pub use self::ast::*;
 pub use self::ast_traits::{AstDeref, AstNodeWrapper, HasAttrs, HasNodeId, HasSpan, HasTokens};
diff --git a/compiler/rustc_builtin_macros/src/autodiff.rs b/compiler/rustc_builtin_macros/src/autodiff.rs
@@ -1,17 +1,19 @@
-#![allow(unused)]
-
-use crate::errors;
+#![allow(unused_imports)]
 //use crate::util::check_builtin_macro_attribute;
 //use crate::util::check_autodiff;
 
+use std::string::String;
+use crate::errors;
 use rustc_ast::ptr::P;
-use rustc_ast::{self as ast, FnHeader, FnSig, Generics, StmtKind};
-use rustc_ast::{Fn, ItemKind, Stmt, TyKind, Unsafe};
+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_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};
 
 pub fn expand(
     ecx: &mut ExtCtxt<'_>,
@@ -20,69 +22,183 @@ pub fn expand(
     item: Annotatable,
 ) -> Vec<Annotatable> {
     //check_builtin_macro_attribute(ecx, meta_item, sym::alloc_error_handler);
-    //check_builtin_macro_attribute(ecx, meta_item, sym::autodiff);
 
-    dbg!(&meta_item);
+    let meta_item_vec: ThinVec<NestedMetaItem> = match meta_item.kind {
+        ast::MetaItemKind::List(ref vec) => vec.clone(),
+        _ => {
+            ecx.sess
+                .dcx()
+                .emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
+            return vec![item];
+        }
+    };
     let input = item.clone();
     let orig_item: P<ast::Item> = item.clone().expect_item();
     let mut d_item: P<ast::Item> = item.clone().expect_item();
+    let primal = orig_item.ident.clone();
 
-    // Allow using `#[autodiff(...)]` on a Fn
-    let (fn_item, _ty_span)  = if let Annotatable::Item(item) = &item
+    // Allow using `#[autodiff(...)]` only on a Fn
+    let (fn_item, has_ret, sig, sig_span) = if let Annotatable::Item(item) = &item
         && let ItemKind::Fn(box ast::Fn { sig, .. }) = &item.kind
     {
-        dbg!(&item);
-        (item, ecx.with_def_site_ctxt(sig.span))
+        (item, sig.decl.output.has_ret(), sig, ecx.with_def_site_ctxt(sig.span))
     } else {
         ecx.sess
             .dcx()
             .emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
         return vec![input];
     };
-    let _x: &ItemKind = &fn_item.kind;
-    d_item.ident.name =
-        Symbol::intern(format!("d_{}", fn_item.ident.name).as_str());
+    let x: AutoDiffAttrs = AutoDiffAttrs::from_ast(&meta_item_vec, has_ret);
+    dbg!(&x);
+    let span = ecx.with_def_site_ctxt(fn_item.span);
+
+    let (d_decl, old_names, new_args) = gen_enzyme_decl(ecx, &sig.decl, &x, span, sig_span);
+    let d_body = gen_enzyme_body(ecx, primal, &old_names, &new_args, span, sig_span);
+    let meta_item_name = meta_item_vec[0].meta_item().unwrap();
+    d_item.ident = meta_item_name.path.segments[0].ident;
+    // update d_item
+    if let ItemKind::Fn(box ast::Fn { sig, body, .. }) = &mut d_item.kind {
+        *sig.decl = d_decl;
+        *body = Some(d_body);
+    } else {
+        ecx.sess
+            .dcx()
+            .emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
+        return vec![input];
+    }
+
     let orig_annotatable = Annotatable::Item(orig_item.clone());
     let d_annotatable = Annotatable::Item(d_item.clone());
     return vec![orig_annotatable, d_annotatable];
 }
 
-// #[rustc_std_internal_symbol]
-// unsafe fn __rg_oom(size: usize, align: usize) -> ! {
-//     handler(core::alloc::Layout::from_size_align_unchecked(size, align))
-// }
-//fn generate_handler(cx: &ExtCtxt<'_>, handler: Ident, span: Span, sig_span: Span) -> Stmt {
-//    let usize = cx.path_ident(span, Ident::new(sym::usize, span));
-//    let ty_usize = cx.ty_path(usize);
-//    let size = Ident::from_str_and_span("size", span);
-//    let align = Ident::from_str_and_span("align", span);
-//
-//    let layout_new = cx.std_path(&[sym::alloc, sym::Layout, sym::from_size_align_unchecked]);
-//    let layout_new = cx.expr_path(cx.path(span, layout_new));
-//    let layout = cx.expr_call(
-//        span,
-//        layout_new,
-//        thin_vec![cx.expr_ident(span, size), cx.expr_ident(span, align)],
-//    );
-//
-//    let call = cx.expr_call_ident(sig_span, handler, thin_vec![layout]);
-//
-//    let never = ast::FnRetTy::Ty(cx.ty(span, TyKind::Never));
-//    let params = thin_vec![cx.param(span, size, ty_usize.clone()), cx.param(span, align, ty_usize)];
-//    let decl = cx.fn_decl(params, never);
-//    let header = FnHeader { unsafety: Unsafe::Yes(span), ..FnHeader::default() };
-//    let sig = FnSig { decl, header, span: span };
-//
-//    let body = Some(cx.block_expr(call));
-//    let kind = ItemKind::Fn(Box::new(Fn {
-//        defaultness: ast::Defaultness::Final,
-//        sig,
-//        generics: Generics::default(),
-//        body,
-//    }));
-//
-//    let attrs = thin_vec![cx.attr_word(sym::rustc_std_internal_symbol, span)];
-//
-//    let item = cx.item(span, Ident::from_str_and_span("__rg_oom", span), attrs, kind);
-//    cx.stmt_item(sig_span, item)
-//}
+// shadow arguments must be mutable references or ptrs, because Enzyme will write into them.
+fn assure_mut_ref(ty: &ast::Ty) -> ast::Ty {
+    let mut ty = ty.clone();
+    match ty.kind {
+        TyKind::Ptr(ref mut mut_ty) => {
+            mut_ty.mutbl = ast::Mutability::Mut;
+        }
+        TyKind::Ref(_, ref mut mut_ty) => {
+            mut_ty.mutbl = ast::Mutability::Mut;
+        }
+        _ => {
+            panic!("unsupported type: {:?}", ty);
+        }
+    }
+    ty
+}
+
+
+// The body of our generated functions will consist of three black_Box calls.
+// The first will call the primal function with the original arguments.
+// The second will just take the shadow arguments.
+// The third will (unsafely) call std::mem::zeroed(), to match the return type of the new function
+// (whatever that might be). This way we surpress rustc from optimizing anyt argument away.
+fn gen_enzyme_body(ecx: &ExtCtxt<'_>, primal: Ident, old_names: &[String], new_names: &[String], span: Span, sig_span: Span) -> P<ast::Block> {
+    let blackbox_path = ecx.std_path(&[Symbol::intern("hint"), Symbol::intern("black_box")]);
+    let zeroed_path = ecx.std_path(&[Symbol::intern("mem"), Symbol::intern("zeroed")]);
+
+    let primal_call_expr = ecx.expr_path(ecx.path_ident(span, primal));
+    let blackbox_call_expr = ecx.expr_path(ecx.path(span, blackbox_path));
+    let zeroed_call_expr = ecx.expr_path(ecx.path(span, zeroed_path));
+
+    let mem_zeroed_call: Stmt = ecx.stmt_expr(ecx.expr_call(
+        span,
+        zeroed_call_expr.clone(),
+        thin_vec![],
+    ));
+    let unsafe_block_with_zeroed_call: P<ast::Expr> = ecx.expr_block(P(ast::Block {
+        stmts: thin_vec![mem_zeroed_call],
+        id: ast::DUMMY_NODE_ID,
+        rules: ast::BlockCheckMode::Unsafe(ast::UserProvided),
+        span: sig_span,
+        tokens: None,
+        could_be_bare_literal: false,
+    }));
+    // create ::core::hint::black_box(array(arr));
+    let _primal_call = ecx.expr_call(
+        span,
+        primal_call_expr.clone(),
+        old_names.iter().map(|name| {
+            ecx.expr_path(ecx.path_ident(span, Ident::from_str(name)))
+        }).collect(),
+    );
+
+    // create ::core::hint::black_box(grad_arr, tang_y));
+    let black_box1 = ecx.expr_call(
+        sig_span,
+        blackbox_call_expr.clone(),
+        new_names.iter().map(|arg| {
+            ecx.expr_path(ecx.path_ident(span, Ident::from_str(arg)))
+        }).collect(),
+    );
+
+    // create ::core::hint::black_box(unsafe { ::core::mem::zeroed() })
+    let black_box2 = ecx.expr_call(
+        sig_span,
+        blackbox_call_expr.clone(),
+        thin_vec![unsafe_block_with_zeroed_call.clone()],
+    );
+
+    let mut body = ecx.block(span, ThinVec::new());
+    body.stmts.push(ecx.stmt_expr(black_box1));
+    body.stmts.push(ecx.stmt_expr(black_box2));
+    body
+}
+
+// Generate the new function declaration. Const arguments are kept as is. Duplicated arguments must
+// be pointers or references. Those receive a shadow argument, which is a mutable reference/pointer.
+// Active arguments must be scalars. Their shadow argument is added to the return type (and will be
+// zero-initialized by Enzyme). Active arguments are not handled yet.
+// Each argument of the primal function (and the return type if existing) must be annotated with an
+// activity.
+fn gen_enzyme_decl(_ecx: &ExtCtxt<'_>, decl: &ast::FnDecl, x: &AutoDiffAttrs, _span: Span, _sig_span: Span)
+        -> (ast::FnDecl, Vec<String>, Vec<String>) {
+    assert!(decl.inputs.len() == x.input_activity.len());
+    assert!(decl.output.has_ret() == x.has_ret_activity());
+    let mut d_decl = decl.clone();
+    let mut d_inputs = Vec::new();
+    let mut new_inputs = Vec::new();
+    let mut old_names = Vec::new();
+    for (arg, activity) in decl.inputs.iter().zip(x.input_activity.iter()) {
+        dbg!(&arg);
+        d_inputs.push(arg.clone());
+        match activity {
+            DiffActivity::Duplicated => {
+                let mut shadow_arg = arg.clone();
+                shadow_arg.ty = P(assure_mut_ref(&arg.ty));
+                // adjust name depending on mode
+                let old_name = if let PatKind::Ident(_, ident, _) = shadow_arg.pat.kind {
+                    ident.name
+                } else {
+                    dbg!(&shadow_arg.pat);
+                    panic!("not an ident?");
+                };
+                old_names.push(old_name.to_string());
+                let name: String = match x.mode {
+                    DiffMode::Reverse => format!("d{}", old_name),
+                    DiffMode::Forward => format!("b{}", old_name),
+                    _ => panic!("unsupported mode: {}", old_name),
+                };
+                dbg!(&name);
+                new_inputs.push(name.clone());
+                shadow_arg.pat = P(ast::Pat {
+                    // TODO: Check id
+                    id: ast::DUMMY_NODE_ID,
+                    kind: PatKind::Ident(BindingAnnotation::NONE,
+                        Ident::from_str_and_span(&name, shadow_arg.pat.span),
+                        None,
+                    ),
+                    span: shadow_arg.pat.span,
+                    tokens: shadow_arg.pat.tokens.clone(),
+                });
+
+                d_inputs.push(shadow_arg);
+            }
+            _ => {},
+        }
+    }
+    d_decl.inputs = d_inputs.into();
+    (d_decl, old_names, new_inputs)
+}

Original file line number	Diff line number	Diff line change
`@@ -2574,6 +2574,12 @@ impl FnRetTy {`
`2574`	`2574`	`FnRetTy::Ty(ty) => ty.span,`
`2575`	`2575`	`}`
`2576`	`2576`	`}`
	`2577`	`+ pub fn has_ret(&self) -> bool {`
	`2578`	`+ match self {`
	`2579`	`+ FnRetTy::Default(_) => false,`
	`2580`	`+ FnRetTy::Ty(_) => true,`
	`2581`	`+ }`
	`2582`	`+ }`
`2577`	`2583`	`}`
`2578`	`2584`
`2579`	`2585`	`#[derive(Clone, Copy, PartialEq, Encodable, Decodable, Debug)]`