diff options
Diffstat (limited to 'compiler/rustc_builtin_macros/src')
| -rw-r--r-- | compiler/rustc_builtin_macros/src/autodiff.rs | 159 | ||||
| -rw-r--r-- | compiler/rustc_builtin_macros/src/cfg_eval.rs | 13 | ||||
| -rw-r--r-- | compiler/rustc_builtin_macros/src/define_opaque.rs | 10 | ||||
| -rw-r--r-- | compiler/rustc_builtin_macros/src/derive.rs | 2 | ||||
| -rw-r--r-- | compiler/rustc_builtin_macros/src/test.rs | 1 |
5 files changed, 112 insertions, 73 deletions
diff --git a/compiler/rustc_builtin_macros/src/autodiff.rs b/compiler/rustc_builtin_macros/src/autodiff.rs index dcd3c1ce8d9..5cd653c7945 100644 --- a/compiler/rustc_builtin_macros/src/autodiff.rs +++ b/compiler/rustc_builtin_macros/src/autodiff.rs @@ -26,6 +26,16 @@ mod llvm_enzyme { use crate::errors; + pub(crate) fn outer_normal_attr( + kind: &P<rustc_ast::NormalAttr>, + id: rustc_ast::AttrId, + span: Span, + ) -> rustc_ast::Attribute { + let style = rustc_ast::AttrStyle::Outer; + let kind = rustc_ast::AttrKind::Normal(kind.clone()); + rustc_ast::Attribute { kind, id, style, span } + } + // If we have a default `()` return type or explicitley `()` return type, // then we often can skip doing some work. fn has_ret(ty: &FnRetTy) -> bool { @@ -147,7 +157,7 @@ mod llvm_enzyme { }; (sig.clone(), false) } - Annotatable::AssocItem(assoc_item, _) => { + Annotatable::AssocItem(assoc_item, Impl { of_trait: false }) => { let sig = match &assoc_item.kind { ast::AssocItemKind::Fn(box ast::Fn { sig, .. }) => sig, _ => { @@ -224,20 +234,8 @@ mod llvm_enzyme { .filter(|a| **a == DiffActivity::Active || **a == DiffActivity::ActiveOnly) .count() as u32; let (d_sig, new_args, idents, errored) = gen_enzyme_decl(ecx, &sig, &x, span); - let new_decl_span = d_sig.span; let d_body = gen_enzyme_body( - ecx, - &x, - n_active, - &sig, - &d_sig, - primal, - &new_args, - span, - sig_span, - new_decl_span, - idents, - errored, + ecx, &x, n_active, &sig, &d_sig, primal, &new_args, span, sig_span, idents, errored, ); let d_ident = first_ident(&meta_item_vec[0]); @@ -270,36 +268,39 @@ mod llvm_enzyme { }; let inline_never_attr = P(ast::NormalAttr { item: inline_item, tokens: None }); let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id(); - let attr: ast::Attribute = ast::Attribute { - kind: ast::AttrKind::Normal(rustc_ad_attr.clone()), - id: new_id, - style: ast::AttrStyle::Outer, - span, - }; + let attr = outer_normal_attr(&rustc_ad_attr, new_id, span); let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id(); - let inline_never: ast::Attribute = ast::Attribute { - kind: ast::AttrKind::Normal(inline_never_attr), - id: new_id, - style: ast::AttrStyle::Outer, - span, - }; + let inline_never = outer_normal_attr(&inline_never_attr, new_id, span); + + // We're avoid duplicating the attributes `#[rustc_autodiff]` and `#[inline(never)]`. + fn same_attribute(attr: &ast::AttrKind, item: &ast::AttrKind) -> bool { + match (attr, item) { + (ast::AttrKind::Normal(a), ast::AttrKind::Normal(b)) => { + let a = &a.item.path; + let b = &b.item.path; + a.segments.len() == b.segments.len() + && a.segments.iter().zip(b.segments.iter()).all(|(a, b)| a.ident == b.ident) + } + _ => false, + } + } // Don't add it multiple times: let orig_annotatable: Annotatable = match item { Annotatable::Item(ref mut iitem) => { - if !iitem.attrs.iter().any(|a| a.id == attr.id) { + if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &attr.kind)) { iitem.attrs.push(attr); } - if !iitem.attrs.iter().any(|a| a.id == inline_never.id) { + if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind)) { iitem.attrs.push(inline_never.clone()); } Annotatable::Item(iitem.clone()) } - Annotatable::AssocItem(ref mut assoc_item, i @ Impl) => { - if !assoc_item.attrs.iter().any(|a| a.id == attr.id) { + Annotatable::AssocItem(ref mut assoc_item, i @ Impl { of_trait: false }) => { + if !assoc_item.attrs.iter().any(|a| same_attribute(&a.kind, &attr.kind)) { assoc_item.attrs.push(attr); } - if !assoc_item.attrs.iter().any(|a| a.id == inline_never.id) { + if !assoc_item.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind)) { assoc_item.attrs.push(inline_never.clone()); } Annotatable::AssocItem(assoc_item.clone(), i) @@ -314,13 +315,7 @@ mod llvm_enzyme { delim: rustc_ast::token::Delimiter::Parenthesis, tokens: ts, }); - let d_attr: ast::Attribute = ast::Attribute { - kind: ast::AttrKind::Normal(rustc_ad_attr.clone()), - id: new_id, - style: ast::AttrStyle::Outer, - span, - }; - + let d_attr = outer_normal_attr(&rustc_ad_attr, new_id, span); let d_annotatable = if is_impl { let assoc_item: AssocItemKind = ast::AssocItemKind::Fn(asdf); let d_fn = P(ast::AssocItem { @@ -332,7 +327,7 @@ mod llvm_enzyme { kind: assoc_item, tokens: None, }); - Annotatable::AssocItem(d_fn, Impl) + Annotatable::AssocItem(d_fn, Impl { of_trait: false }) } else { let mut d_fn = ecx.item(span, d_ident, thin_vec![d_attr, inline_never], ItemKind::Fn(asdf)); @@ -361,30 +356,27 @@ mod llvm_enzyme { ty } - /// We only want this function to type-check, since we will replace the body - /// later on llvm level. Using `loop {}` does not cover all return types anymore, - /// so instead we build something that should pass. We also add a inline_asm - /// line, as one more barrier for rustc to prevent inlining of this function. - /// FIXME(ZuseZ4): We still have cases of incorrect inlining across modules, see - /// <https://github.com/EnzymeAD/rust/issues/173>, so this isn't sufficient. - /// It also triggers an Enzyme crash if we due to a bug ever try to differentiate - /// this function (which should never happen, since it is only a placeholder). - /// Finally, we also add back_box usages of all input arguments, to prevent rustc - /// from optimizing any arguments away. - fn gen_enzyme_body( + // Will generate a body of the type: + // ``` + // { + // unsafe { + // asm!("NOP"); + // } + // ::core::hint::black_box(primal(args)); + // ::core::hint::black_box((args, ret)); + // <This part remains to be done by following function> + // } + // ``` + fn init_body_helper( ecx: &ExtCtxt<'_>, - x: &AutoDiffAttrs, - n_active: u32, - sig: &ast::FnSig, - d_sig: &ast::FnSig, + span: Span, primal: Ident, new_names: &[String], - span: Span, sig_span: Span, new_decl_span: Span, - idents: Vec<Ident>, + idents: &[Ident], errored: bool, - ) -> P<ast::Block> { + ) -> (P<ast::Block>, P<ast::Expr>, P<ast::Expr>, P<ast::Expr>) { let blackbox_path = ecx.std_path(&[sym::hint, sym::black_box]); let noop = ast::InlineAsm { asm_macro: ast::AsmMacro::Asm, @@ -433,6 +425,51 @@ mod llvm_enzyme { } body.stmts.push(ecx.stmt_semi(black_box_remaining_args)); + (body, primal_call, black_box_primal_call, blackbox_call_expr) + } + + /// We only want this function to type-check, since we will replace the body + /// later on llvm level. Using `loop {}` does not cover all return types anymore, + /// so instead we manually build something that should pass the type checker. + /// We also add a inline_asm line, as one more barrier for rustc to prevent inlining + /// or const propagation. inline_asm will also triggers an Enzyme crash if due to another + /// bug would ever try to accidentially differentiate this placeholder function body. + /// Finally, we also add back_box usages of all input arguments, to prevent rustc + /// from optimizing any arguments away. + fn gen_enzyme_body( + ecx: &ExtCtxt<'_>, + x: &AutoDiffAttrs, + n_active: u32, + sig: &ast::FnSig, + d_sig: &ast::FnSig, + primal: Ident, + new_names: &[String], + span: Span, + sig_span: Span, + idents: Vec<Ident>, + errored: bool, + ) -> P<ast::Block> { + let new_decl_span = d_sig.span; + + // Just adding some default inline-asm and black_box usages to prevent early inlining + // and optimizations which alter the function signature. + // + // The bb_primal_call is the black_box call of the primal function. We keep it around, + // since it has the convenient property of returning the type of the primal function, + // Remember, we only care to match types here. + // No matter which return we pick, we always wrap it into a std::hint::black_box call, + // to prevent rustc from propagating it into the caller. + let (mut body, primal_call, bb_primal_call, bb_call_expr) = init_body_helper( + ecx, + span, + primal, + new_names, + sig_span, + new_decl_span, + &idents, + errored, + ); + if !has_ret(&d_sig.decl.output) { // there is no return type that we have to match, () works fine. return body; @@ -444,7 +481,7 @@ mod llvm_enzyme { if primal_ret && n_active == 0 && x.mode.is_rev() { // We only have the primal ret. - body.stmts.push(ecx.stmt_expr(black_box_primal_call)); + body.stmts.push(ecx.stmt_expr(bb_primal_call)); return body; } @@ -536,11 +573,11 @@ mod llvm_enzyme { return body; } [arg] => { - ret = ecx.expr_call(new_decl_span, blackbox_call_expr, thin_vec![arg.clone()]); + ret = ecx.expr_call(new_decl_span, bb_call_expr, thin_vec![arg.clone()]); } args => { let ret_tuple: P<ast::Expr> = ecx.expr_tuple(span, args.into()); - ret = ecx.expr_call(new_decl_span, blackbox_call_expr, thin_vec![ret_tuple]); + ret = ecx.expr_call(new_decl_span, bb_call_expr, thin_vec![ret_tuple]); } } assert!(has_ret(&d_sig.decl.output)); @@ -553,7 +590,7 @@ mod llvm_enzyme { ecx: &ExtCtxt<'_>, span: Span, primal: Ident, - idents: Vec<Ident>, + idents: &[Ident], ) -> P<ast::Expr> { let has_self = idents.len() > 0 && idents[0].name == kw::SelfLower; if has_self { diff --git a/compiler/rustc_builtin_macros/src/cfg_eval.rs b/compiler/rustc_builtin_macros/src/cfg_eval.rs index 5788966b6e7..b3ba9073118 100644 --- a/compiler/rustc_builtin_macros/src/cfg_eval.rs +++ b/compiler/rustc_builtin_macros/src/cfg_eval.rs @@ -121,18 +121,11 @@ impl CfgEval<'_> { let item = parser.parse_item(ForceCollect::Yes)?.unwrap(); Annotatable::Item(self.flat_map_item(item).pop().unwrap()) } - Annotatable::AssocItem(_, AssocCtxt::Trait) => { + Annotatable::AssocItem(_, ctxt) => { let item = parser.parse_trait_item(ForceCollect::Yes)?.unwrap().unwrap(); Annotatable::AssocItem( - self.flat_map_assoc_item(item, AssocCtxt::Trait).pop().unwrap(), - AssocCtxt::Trait, - ) - } - Annotatable::AssocItem(_, AssocCtxt::Impl) => { - let item = parser.parse_impl_item(ForceCollect::Yes)?.unwrap().unwrap(); - Annotatable::AssocItem( - self.flat_map_assoc_item(item, AssocCtxt::Impl).pop().unwrap(), - AssocCtxt::Impl, + self.flat_map_assoc_item(item, ctxt).pop().unwrap(), + ctxt, ) } Annotatable::ForeignItem(_) => { diff --git a/compiler/rustc_builtin_macros/src/define_opaque.rs b/compiler/rustc_builtin_macros/src/define_opaque.rs index 9777e772cf2..cd02e81f568 100644 --- a/compiler/rustc_builtin_macros/src/define_opaque.rs +++ b/compiler/rustc_builtin_macros/src/define_opaque.rs @@ -11,15 +11,20 @@ pub(crate) fn expand( let define_opaque = match &mut item { Annotatable::Item(p) => match &mut p.kind { ast::ItemKind::Fn(f) => Some(&mut f.define_opaque), + ast::ItemKind::Const(ct) => Some(&mut ct.define_opaque), + ast::ItemKind::Static(si) => Some(&mut si.define_opaque), _ => None, }, Annotatable::AssocItem(i, _assoc_ctxt) => match &mut i.kind { ast::AssocItemKind::Fn(func) => Some(&mut func.define_opaque), + ast::AssocItemKind::Const(ct) => Some(&mut ct.define_opaque), _ => None, }, Annotatable::Stmt(s) => match &mut s.kind { ast::StmtKind::Item(p) => match &mut p.kind { ast::ItemKind::Fn(f) => Some(&mut f.define_opaque), + ast::ItemKind::Const(ct) => Some(&mut ct.define_opaque), + ast::ItemKind::Static(si) => Some(&mut si.define_opaque), _ => None, }, _ => None, @@ -47,7 +52,10 @@ pub(crate) fn expand( .collect(), ); } else { - ecx.dcx().span_err(meta_item.span, "only functions and methods can define opaque types"); + ecx.dcx().span_err( + meta_item.span, + "only functions, statics, and consts can define opaque types", + ); } vec![item] diff --git a/compiler/rustc_builtin_macros/src/derive.rs b/compiler/rustc_builtin_macros/src/derive.rs index 2653a9f48b9..e259f5b3955 100644 --- a/compiler/rustc_builtin_macros/src/derive.rs +++ b/compiler/rustc_builtin_macros/src/derive.rs @@ -103,7 +103,7 @@ impl MultiItemModifier for Expander { fn dummy_annotatable() -> Annotatable { Annotatable::GenericParam(ast::GenericParam { id: ast::DUMMY_NODE_ID, - ident: Ident::empty(), + ident: Ident::dummy(), attrs: Default::default(), bounds: Default::default(), is_placeholder: false, diff --git a/compiler/rustc_builtin_macros/src/test.rs b/compiler/rustc_builtin_macros/src/test.rs index a05fff2dcd1..239f8657284 100644 --- a/compiler/rustc_builtin_macros/src/test.rs +++ b/compiler/rustc_builtin_macros/src/test.rs @@ -285,6 +285,7 @@ pub(crate) fn expand_test_or_bench( defaultness: ast::Defaultness::Final, generics: ast::Generics::default(), ty: cx.ty(sp, ast::TyKind::Path(None, test_path("TestDescAndFn"))), + define_opaque: None, // test::TestDescAndFn { expr: Some( cx.expr_struct( |