rustc_builtin_macros/
autodiff.rs

1//! This module contains the implementation of the `#[autodiff]` attribute.
2//! Currently our linter isn't smart enough to see that each import is used in one of the two
3//! configs (autodiff enabled or disabled), so we have to add cfg's to each import.
4//! FIXME(ZuseZ4): Remove this once we have a smarter linter.
5
6mod llvm_enzyme {
7    use std::str::FromStr;
8    use std::string::String;
9
10    use rustc_ast::expand::autodiff_attrs::{
11        AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ret_activity,
12        valid_ty_for_activity,
13    };
14    use rustc_ast::ptr::P;
15    use rustc_ast::token::{Lit, LitKind, Token, TokenKind};
16    use rustc_ast::tokenstream::*;
17    use rustc_ast::visit::AssocCtxt::*;
18    use rustc_ast::{
19        self as ast, AssocItemKind, BindingMode, ExprKind, FnRetTy, FnSig, Generics, ItemKind,
20        MetaItemInner, PatKind, QSelf, TyKind, Visibility,
21    };
22    use rustc_expand::base::{Annotatable, ExtCtxt};
23    use rustc_span::{Ident, Span, Symbol, kw, sym};
24    use thin_vec::{ThinVec, thin_vec};
25    use tracing::{debug, trace};
26
27    use crate::errors;
28
29    pub(crate) fn outer_normal_attr(
30        kind: &P<rustc_ast::NormalAttr>,
31        id: rustc_ast::AttrId,
32        span: Span,
33    ) -> rustc_ast::Attribute {
34        let style = rustc_ast::AttrStyle::Outer;
35        let kind = rustc_ast::AttrKind::Normal(kind.clone());
36        rustc_ast::Attribute { kind, id, style, span }
37    }
38
39    // If we have a default `()` return type or explicitley `()` return type,
40    // then we often can skip doing some work.
41    fn has_ret(ty: &FnRetTy) -> bool {
42        match ty {
43            FnRetTy::Ty(ty) => !ty.kind.is_unit(),
44            FnRetTy::Default(_) => false,
45        }
46    }
47    fn first_ident(x: &MetaItemInner) -> rustc_span::Ident {
48        if let Some(l) = x.lit() {
49            match l.kind {
50                ast::LitKind::Int(val, _) => {
51                    // get an Ident from a lit
52                    return rustc_span::Ident::from_str(val.get().to_string().as_str());
53                }
54                _ => {}
55            }
56        }
57
58        let segments = &x.meta_item().unwrap().path.segments;
59        assert!(segments.len() == 1);
60        segments[0].ident
61    }
62
63    fn name(x: &MetaItemInner) -> String {
64        first_ident(x).name.to_string()
65    }
66
67    fn width(x: &MetaItemInner) -> Option<u128> {
68        let lit = x.lit()?;
69        match lit.kind {
70            ast::LitKind::Int(x, _) => Some(x.get()),
71            _ => return None,
72        }
73    }
74
75    // Get information about the function the macro is applied to
76    fn extract_item_info(iitem: &P<ast::Item>) -> Option<(Visibility, FnSig, Ident)> {
77        match &iitem.kind {
78            ItemKind::Fn(box ast::Fn { sig, ident, .. }) => {
79                Some((iitem.vis.clone(), sig.clone(), ident.clone()))
80            }
81            _ => None,
82        }
83    }
84
85    pub(crate) fn from_ast(
86        ecx: &mut ExtCtxt<'_>,
87        meta_item: &ThinVec<MetaItemInner>,
88        has_ret: bool,
89    ) -> AutoDiffAttrs {
90        let dcx = ecx.sess.dcx();
91        let mode = name(&meta_item[1]);
92        let Ok(mode) = DiffMode::from_str(&mode) else {
93            dcx.emit_err(errors::AutoDiffInvalidMode { span: meta_item[1].span(), mode });
94            return AutoDiffAttrs::error();
95        };
96
97        // Now we check, whether the user wants autodiff in batch/vector mode, or scalar mode.
98        // If he doesn't specify an integer (=width), we default to scalar mode, thus width=1.
99        let mut first_activity = 2;
100
101        let width = if let [_, _, x, ..] = &meta_item[..]
102            && let Some(x) = width(x)
103        {
104            first_activity = 3;
105            match x.try_into() {
106                Ok(x) => x,
107                Err(_) => {
108                    dcx.emit_err(errors::AutoDiffInvalidWidth {
109                        span: meta_item[2].span(),
110                        width: x,
111                    });
112                    return AutoDiffAttrs::error();
113                }
114            }
115        } else {
116            1
117        };
118
119        let mut activities: Vec<DiffActivity> = vec![];
120        let mut errors = false;
121        for x in &meta_item[first_activity..] {
122            let activity_str = name(&x);
123            let res = DiffActivity::from_str(&activity_str);
124            match res {
125                Ok(x) => activities.push(x),
126                Err(_) => {
127                    dcx.emit_err(errors::AutoDiffUnknownActivity {
128                        span: x.span(),
129                        act: activity_str,
130                    });
131                    errors = true;
132                }
133            };
134        }
135        if errors {
136            return AutoDiffAttrs::error();
137        }
138
139        // If a return type exist, we need to split the last activity,
140        // otherwise we return None as placeholder.
141        let (ret_activity, input_activity) = if has_ret {
142            let Some((last, rest)) = activities.split_last() else {
143                unreachable!(
144                    "should not be reachable because we counted the number of activities previously"
145                );
146            };
147            (last, rest)
148        } else {
149            (&DiffActivity::None, activities.as_slice())
150        };
151
152        AutoDiffAttrs {
153            mode,
154            width,
155            ret_activity: *ret_activity,
156            input_activity: input_activity.to_vec(),
157        }
158    }
159
160    fn meta_item_inner_to_ts(t: &MetaItemInner, ts: &mut Vec<TokenTree>) {
161        let comma: Token = Token::new(TokenKind::Comma, Span::default());
162        let val = first_ident(t);
163        let t = Token::from_ast_ident(val);
164        ts.push(TokenTree::Token(t, Spacing::Joint));
165        ts.push(TokenTree::Token(comma.clone(), Spacing::Alone));
166    }
167
168    /// We expand the autodiff macro to generate a new placeholder function which passes
169    /// type-checking and can be called by users. The function body of the placeholder function will
170    /// later be replaced on LLVM-IR level, so the design of the body is less important and for now
171    /// should just prevent early inlining and optimizations which alter the function signature.
172    /// The exact signature of the generated function depends on the configuration provided by the
173    /// user, but here is an example:
174    ///
175    /// ```
176    /// #[autodiff(cos_box, Reverse, Duplicated, Active)]
177    /// fn sin(x: &Box<f32>) -> f32 {
178    ///     f32::sin(**x)
179    /// }
180    /// ```
181    /// which becomes expanded to:
182    /// ```
183    /// #[rustc_autodiff]
184    /// #[inline(never)]
185    /// fn sin(x: &Box<f32>) -> f32 {
186    ///     f32::sin(**x)
187    /// }
188    /// #[rustc_autodiff(Reverse, Duplicated, Active)]
189    /// #[inline(never)]
190    /// fn cos_box(x: &Box<f32>, dx: &mut Box<f32>, dret: f32) -> f32 {
191    ///     unsafe {
192    ///         asm!("NOP");
193    ///     };
194    ///     ::core::hint::black_box(sin(x));
195    ///     ::core::hint::black_box((dx, dret));
196    ///     ::core::hint::black_box(sin(x))
197    /// }
198    /// ```
199    /// FIXME(ZuseZ4): Once autodiff is enabled by default, make this a doc comment which is checked
200    /// in CI.
201    pub(crate) fn expand(
202        ecx: &mut ExtCtxt<'_>,
203        expand_span: Span,
204        meta_item: &ast::MetaItem,
205        mut item: Annotatable,
206    ) -> Vec<Annotatable> {
207        if cfg!(not(llvm_enzyme)) {
208            ecx.sess.dcx().emit_err(errors::AutoDiffSupportNotBuild { span: meta_item.span });
209            return vec![item];
210        }
211        let dcx = ecx.sess.dcx();
212
213        // first get information about the annotable item:
214        let Some((vis, sig, primal)) = (match &item {
215            Annotatable::Item(iitem) => extract_item_info(iitem),
216            Annotatable::Stmt(stmt) => match &stmt.kind {
217                ast::StmtKind::Item(iitem) => extract_item_info(iitem),
218                _ => None,
219            },
220            Annotatable::AssocItem(assoc_item, Impl { .. }) => match &assoc_item.kind {
221                ast::AssocItemKind::Fn(box ast::Fn { sig, ident, .. }) => {
222                    Some((assoc_item.vis.clone(), sig.clone(), ident.clone()))
223                }
224                _ => None,
225            },
226            _ => None,
227        }) else {
228            dcx.emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
229            return vec![item];
230        };
231
232        let meta_item_vec: ThinVec<MetaItemInner> = match meta_item.kind {
233            ast::MetaItemKind::List(ref vec) => vec.clone(),
234            _ => {
235                dcx.emit_err(errors::AutoDiffMissingConfig { span: item.span() });
236                return vec![item];
237            }
238        };
239
240        let has_ret = has_ret(&sig.decl.output);
241        let sig_span = ecx.with_call_site_ctxt(sig.span);
242
243        // create TokenStream from vec elemtents:
244        // meta_item doesn't have a .tokens field
245        let mut ts: Vec<TokenTree> = vec![];
246        if meta_item_vec.len() < 2 {
247            // At the bare minimum, we need a fnc name and a mode, even for a dummy function with no
248            // input and output args.
249            dcx.emit_err(errors::AutoDiffMissingConfig { span: item.span() });
250            return vec![item];
251        }
252
253        meta_item_inner_to_ts(&meta_item_vec[1], &mut ts);
254
255        // Now, if the user gave a width (vector aka batch-mode ad), then we copy it.
256        // If it is not given, we default to 1 (scalar mode).
257        let start_position;
258        let kind: LitKind = LitKind::Integer;
259        let symbol;
260        if meta_item_vec.len() >= 3
261            && let Some(width) = width(&meta_item_vec[2])
262        {
263            start_position = 3;
264            symbol = Symbol::intern(&width.to_string());
265        } else {
266            start_position = 2;
267            symbol = sym::integer(1);
268        }
269        let l: Lit = Lit { kind, symbol, suffix: None };
270        let t = Token::new(TokenKind::Literal(l), Span::default());
271        let comma = Token::new(TokenKind::Comma, Span::default());
272        ts.push(TokenTree::Token(t, Spacing::Joint));
273        ts.push(TokenTree::Token(comma.clone(), Spacing::Alone));
274
275        for t in meta_item_vec.clone()[start_position..].iter() {
276            meta_item_inner_to_ts(t, &mut ts);
277        }
278
279        if !has_ret {
280            // We don't want users to provide a return activity if the function doesn't return anything.
281            // For simplicity, we just add a dummy token to the end of the list.
282            let t = Token::new(TokenKind::Ident(sym::None, false.into()), Span::default());
283            ts.push(TokenTree::Token(t, Spacing::Joint));
284            ts.push(TokenTree::Token(comma, Spacing::Alone));
285        }
286        // We remove the last, trailing comma.
287        ts.pop();
288        let ts: TokenStream = TokenStream::from_iter(ts);
289
290        let x: AutoDiffAttrs = from_ast(ecx, &meta_item_vec, has_ret);
291        if !x.is_active() {
292            // We encountered an error, so we return the original item.
293            // This allows us to potentially parse other attributes.
294            return vec![item];
295        }
296        let span = ecx.with_def_site_ctxt(expand_span);
297
298        let n_active: u32 = x
299            .input_activity
300            .iter()
301            .filter(|a| **a == DiffActivity::Active || **a == DiffActivity::ActiveOnly)
302            .count() as u32;
303        let (d_sig, new_args, idents, errored) = gen_enzyme_decl(ecx, &sig, &x, span);
304        let d_body = gen_enzyme_body(
305            ecx, &x, n_active, &sig, &d_sig, primal, &new_args, span, sig_span, idents, errored,
306        );
307
308        // The first element of it is the name of the function to be generated
309        let asdf = Box::new(ast::Fn {
310            defaultness: ast::Defaultness::Final,
311            sig: d_sig,
312            ident: first_ident(&meta_item_vec[0]),
313            generics: Generics::default(),
314            contract: None,
315            body: Some(d_body),
316            define_opaque: None,
317        });
318        let mut rustc_ad_attr =
319            P(ast::NormalAttr::from_ident(Ident::with_dummy_span(sym::rustc_autodiff)));
320
321        let ts2: Vec<TokenTree> = vec![TokenTree::Token(
322            Token::new(TokenKind::Ident(sym::never, false.into()), span),
323            Spacing::Joint,
324        )];
325        let never_arg = ast::DelimArgs {
326            dspan: DelimSpan::from_single(span),
327            delim: ast::token::Delimiter::Parenthesis,
328            tokens: TokenStream::from_iter(ts2),
329        };
330        let inline_item = ast::AttrItem {
331            unsafety: ast::Safety::Default,
332            path: ast::Path::from_ident(Ident::with_dummy_span(sym::inline)),
333            args: ast::AttrArgs::Delimited(never_arg),
334            tokens: None,
335        };
336        let inline_never_attr = P(ast::NormalAttr { item: inline_item, tokens: None });
337        let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id();
338        let attr = outer_normal_attr(&rustc_ad_attr, new_id, span);
339        let new_id = ecx.sess.psess.attr_id_generator.mk_attr_id();
340        let inline_never = outer_normal_attr(&inline_never_attr, new_id, span);
341
342        // We're avoid duplicating the attributes `#[rustc_autodiff]` and `#[inline(never)]`.
343        fn same_attribute(attr: &ast::AttrKind, item: &ast::AttrKind) -> bool {
344            match (attr, item) {
345                (ast::AttrKind::Normal(a), ast::AttrKind::Normal(b)) => {
346                    let a = &a.item.path;
347                    let b = &b.item.path;
348                    a.segments.len() == b.segments.len()
349                        && a.segments.iter().zip(b.segments.iter()).all(|(a, b)| a.ident == b.ident)
350                }
351                _ => false,
352            }
353        }
354
355        // Don't add it multiple times:
356        let orig_annotatable: Annotatable = match item {
357            Annotatable::Item(ref mut iitem) => {
358                if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &attr.kind)) {
359                    iitem.attrs.push(attr);
360                }
361                if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind)) {
362                    iitem.attrs.push(inline_never.clone());
363                }
364                Annotatable::Item(iitem.clone())
365            }
366            Annotatable::AssocItem(ref mut assoc_item, i @ Impl { .. }) => {
367                if !assoc_item.attrs.iter().any(|a| same_attribute(&a.kind, &attr.kind)) {
368                    assoc_item.attrs.push(attr);
369                }
370                if !assoc_item.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind)) {
371                    assoc_item.attrs.push(inline_never.clone());
372                }
373                Annotatable::AssocItem(assoc_item.clone(), i)
374            }
375            Annotatable::Stmt(ref mut stmt) => {
376                match stmt.kind {
377                    ast::StmtKind::Item(ref mut iitem) => {
378                        if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &attr.kind)) {
379                            iitem.attrs.push(attr);
380                        }
381                        if !iitem.attrs.iter().any(|a| same_attribute(&a.kind, &inline_never.kind))
382                        {
383                            iitem.attrs.push(inline_never.clone());
384                        }
385                    }
386                    _ => unreachable!("stmt kind checked previously"),
387                };
388
389                Annotatable::Stmt(stmt.clone())
390            }
391            _ => {
392                unreachable!("annotatable kind checked previously")
393            }
394        };
395        // Now update for d_fn
396        rustc_ad_attr.item.args = rustc_ast::AttrArgs::Delimited(rustc_ast::DelimArgs {
397            dspan: DelimSpan::dummy(),
398            delim: rustc_ast::token::Delimiter::Parenthesis,
399            tokens: ts,
400        });
401
402        let d_attr = outer_normal_attr(&rustc_ad_attr, new_id, span);
403        let d_annotatable = match &item {
404            Annotatable::AssocItem(_, _) => {
405                let assoc_item: AssocItemKind = ast::AssocItemKind::Fn(asdf);
406                let d_fn = P(ast::AssocItem {
407                    attrs: thin_vec![d_attr, inline_never],
408                    id: ast::DUMMY_NODE_ID,
409                    span,
410                    vis,
411                    kind: assoc_item,
412                    tokens: None,
413                });
414                Annotatable::AssocItem(d_fn, Impl { of_trait: false })
415            }
416            Annotatable::Item(_) => {
417                let mut d_fn = ecx.item(span, thin_vec![d_attr, inline_never], ItemKind::Fn(asdf));
418                d_fn.vis = vis;
419
420                Annotatable::Item(d_fn)
421            }
422            Annotatable::Stmt(_) => {
423                let mut d_fn = ecx.item(span, thin_vec![d_attr, inline_never], ItemKind::Fn(asdf));
424                d_fn.vis = vis;
425
426                Annotatable::Stmt(P(ast::Stmt {
427                    id: ast::DUMMY_NODE_ID,
428                    kind: ast::StmtKind::Item(d_fn),
429                    span,
430                }))
431            }
432            _ => {
433                unreachable!("item kind checked previously")
434            }
435        };
436
437        return vec![orig_annotatable, d_annotatable];
438    }
439
440    // shadow arguments (the extra ones which were not in the original (primal) function), in reverse mode must be
441    // mutable references or ptrs, because Enzyme will write into them.
442    fn assure_mut_ref(ty: &ast::Ty) -> ast::Ty {
443        let mut ty = ty.clone();
444        match ty.kind {
445            TyKind::Ptr(ref mut mut_ty) => {
446                mut_ty.mutbl = ast::Mutability::Mut;
447            }
448            TyKind::Ref(_, ref mut mut_ty) => {
449                mut_ty.mutbl = ast::Mutability::Mut;
450            }
451            _ => {
452                panic!("unsupported type: {:?}", ty);
453            }
454        }
455        ty
456    }
457
458    // Will generate a body of the type:
459    // ```
460    // {
461    //   unsafe {
462    //   asm!("NOP");
463    //   }
464    //   ::core::hint::black_box(primal(args));
465    //   ::core::hint::black_box((args, ret));
466    //   <This part remains to be done by following function>
467    // }
468    // ```
469    fn init_body_helper(
470        ecx: &ExtCtxt<'_>,
471        span: Span,
472        primal: Ident,
473        new_names: &[String],
474        sig_span: Span,
475        new_decl_span: Span,
476        idents: &[Ident],
477        errored: bool,
478    ) -> (P<ast::Block>, P<ast::Expr>, P<ast::Expr>, P<ast::Expr>) {
479        let blackbox_path = ecx.std_path(&[sym::hint, sym::black_box]);
480        let noop = ast::InlineAsm {
481            asm_macro: ast::AsmMacro::Asm,
482            template: vec![ast::InlineAsmTemplatePiece::String("NOP".into())],
483            template_strs: Box::new([]),
484            operands: vec![],
485            clobber_abis: vec![],
486            options: ast::InlineAsmOptions::PURE | ast::InlineAsmOptions::NOMEM,
487            line_spans: vec![],
488        };
489        let noop_expr = ecx.expr_asm(span, P(noop));
490        let unsf = ast::BlockCheckMode::Unsafe(ast::UnsafeSource::CompilerGenerated);
491        let unsf_block = ast::Block {
492            stmts: thin_vec![ecx.stmt_semi(noop_expr)],
493            id: ast::DUMMY_NODE_ID,
494            tokens: None,
495            rules: unsf,
496            span,
497        };
498        let unsf_expr = ecx.expr_block(P(unsf_block));
499        let blackbox_call_expr = ecx.expr_path(ecx.path(span, blackbox_path));
500        let primal_call = gen_primal_call(ecx, span, primal, idents);
501        let black_box_primal_call = ecx.expr_call(
502            new_decl_span,
503            blackbox_call_expr.clone(),
504            thin_vec![primal_call.clone()],
505        );
506        let tup_args = new_names
507            .iter()
508            .map(|arg| ecx.expr_path(ecx.path_ident(span, Ident::from_str(arg))))
509            .collect();
510
511        let black_box_remaining_args = ecx.expr_call(
512            sig_span,
513            blackbox_call_expr.clone(),
514            thin_vec![ecx.expr_tuple(sig_span, tup_args)],
515        );
516
517        let mut body = ecx.block(span, ThinVec::new());
518        body.stmts.push(ecx.stmt_semi(unsf_expr));
519
520        // This uses primal args which won't be available if we errored before
521        if !errored {
522            body.stmts.push(ecx.stmt_semi(black_box_primal_call.clone()));
523        }
524        body.stmts.push(ecx.stmt_semi(black_box_remaining_args));
525
526        (body, primal_call, black_box_primal_call, blackbox_call_expr)
527    }
528
529    /// We only want this function to type-check, since we will replace the body
530    /// later on llvm level. Using `loop {}` does not cover all return types anymore,
531    /// so instead we manually build something that should pass the type checker.
532    /// We also add a inline_asm line, as one more barrier for rustc to prevent inlining
533    /// or const propagation. inline_asm will also triggers an Enzyme crash if due to another
534    /// bug would ever try to accidentially differentiate this placeholder function body.
535    /// Finally, we also add back_box usages of all input arguments, to prevent rustc
536    /// from optimizing any arguments away.
537    fn gen_enzyme_body(
538        ecx: &ExtCtxt<'_>,
539        x: &AutoDiffAttrs,
540        n_active: u32,
541        sig: &ast::FnSig,
542        d_sig: &ast::FnSig,
543        primal: Ident,
544        new_names: &[String],
545        span: Span,
546        sig_span: Span,
547        idents: Vec<Ident>,
548        errored: bool,
549    ) -> P<ast::Block> {
550        let new_decl_span = d_sig.span;
551
552        // Just adding some default inline-asm and black_box usages to prevent early inlining
553        // and optimizations which alter the function signature.
554        //
555        // The bb_primal_call is the black_box call of the primal function. We keep it around,
556        // since it has the convenient property of returning the type of the primal function,
557        // Remember, we only care to match types here.
558        // No matter which return we pick, we always wrap it into a std::hint::black_box call,
559        // to prevent rustc from propagating it into the caller.
560        let (mut body, primal_call, bb_primal_call, bb_call_expr) = init_body_helper(
561            ecx,
562            span,
563            primal,
564            new_names,
565            sig_span,
566            new_decl_span,
567            &idents,
568            errored,
569        );
570
571        if !has_ret(&d_sig.decl.output) {
572            // there is no return type that we have to match, () works fine.
573            return body;
574        }
575
576        // Everything from here onwards just tries to fullfil the return type. Fun!
577
578        // having an active-only return means we'll drop the original return type.
579        // So that can be treated identical to not having one in the first place.
580        let primal_ret = has_ret(&sig.decl.output) && !x.has_active_only_ret();
581
582        if primal_ret && n_active == 0 && x.mode.is_rev() {
583            // We only have the primal ret.
584            body.stmts.push(ecx.stmt_expr(bb_primal_call));
585            return body;
586        }
587
588        if !primal_ret && n_active == 1 {
589            // Again no tuple return, so return default float val.
590            let ty = match d_sig.decl.output {
591                FnRetTy::Ty(ref ty) => ty.clone(),
592                FnRetTy::Default(span) => {
593                    panic!("Did not expect Default ret ty: {:?}", span);
594                }
595            };
596            let arg = ty.kind.is_simple_path().unwrap();
597            let tmp = ecx.def_site_path(&[arg, kw::Default]);
598            let default_call_expr = ecx.expr_path(ecx.path(span, tmp));
599            let default_call_expr = ecx.expr_call(new_decl_span, default_call_expr, thin_vec![]);
600            body.stmts.push(ecx.stmt_expr(default_call_expr));
601            return body;
602        }
603
604        let mut exprs: P<ast::Expr> = primal_call;
605        let d_ret_ty = match d_sig.decl.output {
606            FnRetTy::Ty(ref ty) => ty.clone(),
607            FnRetTy::Default(span) => {
608                panic!("Did not expect Default ret ty: {:?}", span);
609            }
610        };
611
612        if x.mode.is_fwd() {
613            // Fwd mode is easy. If the return activity is Const, we support arbitrary types.
614            // Otherwise, we only support a scalar, a pair of scalars, or an array of scalars.
615            // We checked that (on a best-effort base) in the preceding gen_enzyme_decl function.
616            // In all three cases, we can return `std::hint::black_box(<T>::default())`.
617            if x.ret_activity == DiffActivity::Const {
618                // Here we call the primal function, since our dummy function has the same return
619                // type due to the Const return activity.
620                exprs = ecx.expr_call(new_decl_span, bb_call_expr, thin_vec![exprs]);
621            } else {
622                let q = QSelf { ty: d_ret_ty, path_span: span, position: 0 };
623                let y =
624                    ExprKind::Path(Some(P(q)), ecx.path_ident(span, Ident::from_str("default")));
625                let default_call_expr = ecx.expr(span, y);
626                let default_call_expr =
627                    ecx.expr_call(new_decl_span, default_call_expr, thin_vec![]);
628                exprs = ecx.expr_call(new_decl_span, bb_call_expr, thin_vec![default_call_expr]);
629            }
630        } else if x.mode.is_rev() {
631            if x.width == 1 {
632                // We either have `-> ArbitraryType` or `-> (ArbitraryType, repeated_float_scalars)`.
633                match d_ret_ty.kind {
634                    TyKind::Tup(ref args) => {
635                        // We have a tuple return type. We need to create a tuple of the same size
636                        // and fill it with default values.
637                        let mut exprs2 = thin_vec![exprs];
638                        for arg in args.iter().skip(1) {
639                            let arg = arg.kind.is_simple_path().unwrap();
640                            let tmp = ecx.def_site_path(&[arg, kw::Default]);
641                            let default_call_expr = ecx.expr_path(ecx.path(span, tmp));
642                            let default_call_expr =
643                                ecx.expr_call(new_decl_span, default_call_expr, thin_vec![]);
644                            exprs2.push(default_call_expr);
645                        }
646                        exprs = ecx.expr_tuple(new_decl_span, exprs2);
647                    }
648                    _ => {
649                        // Interestingly, even the `-> ArbitraryType` case
650                        // ends up getting matched and handled correctly above,
651                        // so we don't have to handle any other case for now.
652                        panic!("Unsupported return type: {:?}", d_ret_ty);
653                    }
654                }
655            }
656            exprs = ecx.expr_call(new_decl_span, bb_call_expr, thin_vec![exprs]);
657        } else {
658            unreachable!("Unsupported mode: {:?}", x.mode);
659        }
660
661        body.stmts.push(ecx.stmt_expr(exprs));
662
663        body
664    }
665
666    fn gen_primal_call(
667        ecx: &ExtCtxt<'_>,
668        span: Span,
669        primal: Ident,
670        idents: &[Ident],
671    ) -> P<ast::Expr> {
672        let has_self = idents.len() > 0 && idents[0].name == kw::SelfLower;
673        if has_self {
674            let args: ThinVec<_> =
675                idents[1..].iter().map(|arg| ecx.expr_path(ecx.path_ident(span, *arg))).collect();
676            let self_expr = ecx.expr_self(span);
677            ecx.expr_method_call(span, self_expr, primal, args)
678        } else {
679            let args: ThinVec<_> =
680                idents.iter().map(|arg| ecx.expr_path(ecx.path_ident(span, *arg))).collect();
681            let primal_call_expr = ecx.expr_path(ecx.path_ident(span, primal));
682            ecx.expr_call(span, primal_call_expr, args)
683        }
684    }
685
686    // Generate the new function declaration. Const arguments are kept as is. Duplicated arguments must
687    // be pointers or references. Those receive a shadow argument, which is a mutable reference/pointer.
688    // Active arguments must be scalars. Their shadow argument is added to the return type (and will be
689    // zero-initialized by Enzyme).
690    // Each argument of the primal function (and the return type if existing) must be annotated with an
691    // activity.
692    //
693    // Error handling: If the user provides an invalid configuration (incorrect numbers, types, or
694    // both), we emit an error and return the original signature. This allows us to continue parsing.
695    // FIXME(Sa4dUs): make individual activities' span available so errors
696    // can point to only the activity instead of the entire attribute
697    fn gen_enzyme_decl(
698        ecx: &ExtCtxt<'_>,
699        sig: &ast::FnSig,
700        x: &AutoDiffAttrs,
701        span: Span,
702    ) -> (ast::FnSig, Vec<String>, Vec<Ident>, bool) {
703        let dcx = ecx.sess.dcx();
704        let has_ret = has_ret(&sig.decl.output);
705        let sig_args = sig.decl.inputs.len() + if has_ret { 1 } else { 0 };
706        let num_activities = x.input_activity.len() + if x.has_ret_activity() { 1 } else { 0 };
707        if sig_args != num_activities {
708            dcx.emit_err(errors::AutoDiffInvalidNumberActivities {
709                span,
710                expected: sig_args,
711                found: num_activities,
712            });
713            // This is not the right signature, but we can continue parsing.
714            return (sig.clone(), vec![], vec![], true);
715        }
716        assert!(sig.decl.inputs.len() == x.input_activity.len());
717        assert!(has_ret == x.has_ret_activity());
718        let mut d_decl = sig.decl.clone();
719        let mut d_inputs = Vec::new();
720        let mut new_inputs = Vec::new();
721        let mut idents = Vec::new();
722        let mut act_ret = ThinVec::new();
723
724        // We have two loops, a first one just to check the activities and types and possibly report
725        // multiple errors in one compilation session.
726        let mut errors = false;
727        for (arg, activity) in sig.decl.inputs.iter().zip(x.input_activity.iter()) {
728            if !valid_input_activity(x.mode, *activity) {
729                dcx.emit_err(errors::AutoDiffInvalidApplicationModeAct {
730                    span,
731                    mode: x.mode.to_string(),
732                    act: activity.to_string(),
733                });
734                errors = true;
735            }
736            if !valid_ty_for_activity(&arg.ty, *activity) {
737                dcx.emit_err(errors::AutoDiffInvalidTypeForActivity {
738                    span: arg.ty.span,
739                    act: activity.to_string(),
740                });
741                errors = true;
742            }
743        }
744
745        if has_ret && !valid_ret_activity(x.mode, x.ret_activity) {
746            dcx.emit_err(errors::AutoDiffInvalidRetAct {
747                span,
748                mode: x.mode.to_string(),
749                act: x.ret_activity.to_string(),
750            });
751            // We don't set `errors = true` to avoid annoying type errors relative
752            // to the expanded macro type signature
753        }
754
755        if errors {
756            // This is not the right signature, but we can continue parsing.
757            return (sig.clone(), new_inputs, idents, true);
758        }
759
760        let unsafe_activities = x
761            .input_activity
762            .iter()
763            .any(|&act| matches!(act, DiffActivity::DuplicatedOnly | DiffActivity::DualOnly));
764        for (arg, activity) in sig.decl.inputs.iter().zip(x.input_activity.iter()) {
765            d_inputs.push(arg.clone());
766            match activity {
767                DiffActivity::Active => {
768                    act_ret.push(arg.ty.clone());
769                    // if width =/= 1, then push [arg.ty; width] to act_ret
770                }
771                DiffActivity::ActiveOnly => {
772                    // We will add the active scalar to the return type.
773                    // This is handled later.
774                }
775                DiffActivity::Duplicated | DiffActivity::DuplicatedOnly => {
776                    for i in 0..x.width {
777                        let mut shadow_arg = arg.clone();
778                        // We += into the shadow in reverse mode.
779                        shadow_arg.ty = P(assure_mut_ref(&arg.ty));
780                        let old_name = if let PatKind::Ident(_, ident, _) = arg.pat.kind {
781                            ident.name
782                        } else {
783                            debug!("{:#?}", &shadow_arg.pat);
784                            panic!("not an ident?");
785                        };
786                        let name: String = format!("d{}_{}", old_name, i);
787                        new_inputs.push(name.clone());
788                        let ident = Ident::from_str_and_span(&name, shadow_arg.pat.span);
789                        shadow_arg.pat = P(ast::Pat {
790                            id: ast::DUMMY_NODE_ID,
791                            kind: PatKind::Ident(BindingMode::NONE, ident, None),
792                            span: shadow_arg.pat.span,
793                            tokens: shadow_arg.pat.tokens.clone(),
794                        });
795                        d_inputs.push(shadow_arg.clone());
796                    }
797                }
798                DiffActivity::Dual
799                | DiffActivity::DualOnly
800                | DiffActivity::Dualv
801                | DiffActivity::DualvOnly => {
802                    // the *v variants get lowered to enzyme_dupv and enzyme_dupnoneedv, which cause
803                    // Enzyme to not expect N arguments, but one argument (which is instead larger).
804                    let iterations =
805                        if matches!(activity, DiffActivity::Dualv | DiffActivity::DualvOnly) {
806                            1
807                        } else {
808                            x.width
809                        };
810                    for i in 0..iterations {
811                        let mut shadow_arg = arg.clone();
812                        let old_name = if let PatKind::Ident(_, ident, _) = arg.pat.kind {
813                            ident.name
814                        } else {
815                            debug!("{:#?}", &shadow_arg.pat);
816                            panic!("not an ident?");
817                        };
818                        let name: String = format!("b{}_{}", old_name, i);
819                        new_inputs.push(name.clone());
820                        let ident = Ident::from_str_and_span(&name, shadow_arg.pat.span);
821                        shadow_arg.pat = P(ast::Pat {
822                            id: ast::DUMMY_NODE_ID,
823                            kind: PatKind::Ident(BindingMode::NONE, ident, None),
824                            span: shadow_arg.pat.span,
825                            tokens: shadow_arg.pat.tokens.clone(),
826                        });
827                        d_inputs.push(shadow_arg.clone());
828                    }
829                }
830                DiffActivity::Const => {
831                    // Nothing to do here.
832                }
833                DiffActivity::None | DiffActivity::FakeActivitySize(_) => {
834                    panic!("Should not happen");
835                }
836            }
837            if let PatKind::Ident(_, ident, _) = arg.pat.kind {
838                idents.push(ident.clone());
839            } else {
840                panic!("not an ident?");
841            }
842        }
843
844        let active_only_ret = x.ret_activity == DiffActivity::ActiveOnly;
845        if active_only_ret {
846            assert!(x.mode.is_rev());
847        }
848
849        // If we return a scalar in the primal and the scalar is active,
850        // then add it as last arg to the inputs.
851        if x.mode.is_rev() {
852            match x.ret_activity {
853                DiffActivity::Active | DiffActivity::ActiveOnly => {
854                    let ty = match d_decl.output {
855                        FnRetTy::Ty(ref ty) => ty.clone(),
856                        FnRetTy::Default(span) => {
857                            panic!("Did not expect Default ret ty: {:?}", span);
858                        }
859                    };
860                    let name = "dret".to_string();
861                    let ident = Ident::from_str_and_span(&name, ty.span);
862                    let shadow_arg = ast::Param {
863                        attrs: ThinVec::new(),
864                        ty: ty.clone(),
865                        pat: P(ast::Pat {
866                            id: ast::DUMMY_NODE_ID,
867                            kind: PatKind::Ident(BindingMode::NONE, ident, None),
868                            span: ty.span,
869                            tokens: None,
870                        }),
871                        id: ast::DUMMY_NODE_ID,
872                        span: ty.span,
873                        is_placeholder: false,
874                    };
875                    d_inputs.push(shadow_arg);
876                    new_inputs.push(name);
877                }
878                _ => {}
879            }
880        }
881        d_decl.inputs = d_inputs.into();
882
883        if x.mode.is_fwd() {
884            let ty = match d_decl.output {
885                FnRetTy::Ty(ref ty) => ty.clone(),
886                FnRetTy::Default(span) => {
887                    // We want to return std::hint::black_box(()).
888                    let kind = TyKind::Tup(ThinVec::new());
889                    let ty = P(rustc_ast::Ty { kind, id: ast::DUMMY_NODE_ID, span, tokens: None });
890                    d_decl.output = FnRetTy::Ty(ty.clone());
891                    assert!(matches!(x.ret_activity, DiffActivity::None));
892                    // this won't be used below, so any type would be fine.
893                    ty
894                }
895            };
896
897            if matches!(x.ret_activity, DiffActivity::Dual | DiffActivity::Dualv) {
898                let kind = if x.width == 1 || matches!(x.ret_activity, DiffActivity::Dualv) {
899                    // Dual can only be used for f32/f64 ret.
900                    // In that case we return now a tuple with two floats.
901                    TyKind::Tup(thin_vec![ty.clone(), ty.clone()])
902                } else {
903                    // We have to return [T; width+1], +1 for the primal return.
904                    let anon_const = rustc_ast::AnonConst {
905                        id: ast::DUMMY_NODE_ID,
906                        value: ecx.expr_usize(span, 1 + x.width as usize),
907                    };
908                    TyKind::Array(ty.clone(), anon_const)
909                };
910                let ty = P(rustc_ast::Ty { kind, id: ty.id, span: ty.span, tokens: None });
911                d_decl.output = FnRetTy::Ty(ty);
912            }
913            if matches!(x.ret_activity, DiffActivity::DualOnly | DiffActivity::DualvOnly) {
914                // No need to change the return type,
915                // we will just return the shadow in place of the primal return.
916                // However, if we have a width > 1, then we don't return -> T, but -> [T; width]
917                if x.width > 1 {
918                    let anon_const = rustc_ast::AnonConst {
919                        id: ast::DUMMY_NODE_ID,
920                        value: ecx.expr_usize(span, x.width as usize),
921                    };
922                    let kind = TyKind::Array(ty.clone(), anon_const);
923                    let ty = P(rustc_ast::Ty { kind, id: ty.id, span: ty.span, tokens: None });
924                    d_decl.output = FnRetTy::Ty(ty);
925                }
926            }
927        }
928
929        // If we use ActiveOnly, drop the original return value.
930        d_decl.output =
931            if active_only_ret { FnRetTy::Default(span) } else { d_decl.output.clone() };
932
933        trace!("act_ret: {:?}", act_ret);
934
935        // If we have an active input scalar, add it's gradient to the
936        // return type. This might require changing the return type to a
937        // tuple.
938        if act_ret.len() > 0 {
939            let ret_ty = match d_decl.output {
940                FnRetTy::Ty(ref ty) => {
941                    if !active_only_ret {
942                        act_ret.insert(0, ty.clone());
943                    }
944                    let kind = TyKind::Tup(act_ret);
945                    P(rustc_ast::Ty { kind, id: ty.id, span: ty.span, tokens: None })
946                }
947                FnRetTy::Default(span) => {
948                    if act_ret.len() == 1 {
949                        act_ret[0].clone()
950                    } else {
951                        let kind = TyKind::Tup(act_ret.iter().map(|arg| arg.clone()).collect());
952                        P(rustc_ast::Ty { kind, id: ast::DUMMY_NODE_ID, span, tokens: None })
953                    }
954                }
955            };
956            d_decl.output = FnRetTy::Ty(ret_ty);
957        }
958
959        let mut d_header = sig.header.clone();
960        if unsafe_activities {
961            d_header.safety = rustc_ast::Safety::Unsafe(span);
962        }
963        let d_sig = FnSig { header: d_header, decl: d_decl, span };
964        trace!("Generated signature: {:?}", d_sig);
965        (d_sig, new_inputs, idents, false)
966    }
967}
968
969pub(crate) use llvm_enzyme::expand;
rustc_builtin_macros/autodiff.rs

rustc_builtin_macros/
autodiff.rs
