Post

Split autodiff into autodiff_forward and autodiff_reverse

PR #140697 introduces a new improvement for the autodiff feature: splitting the existing #[autodiff(...)] procedural macro into two specialized macros: #[autodiff_forward(...)] and #[autodiff_reverse(...)].

The original unified #[autodiff(NAME, MODE, INPUT_ACTIVITIES, OUTPUT_ACTIVITIES)] macro posed several issues: 1. Forward and reverse modes require different parameters. 2. While both modes use shadow arguments, their initialization logic differs. 3. Splitting the macros simplifies documentation and makes it a bit easier, as it’s already done at https://enzyme.mit.edu/rust.

Each macro variant now routes through its own adapter, by passing the mode: DiffMode at macro expansion (compiler/rustc_builtin_macros/src/autodiff.rs) to extend and create two different functions: expand_forward and expand_reverse:

1
2
3
4
5
6
7
pub(crate) fn expand_forward(...) -> Vec<Annotatable> {
    expand_with_mode(..., DiffMode::Forward)
}

pub(crate) fn expand_reverse(...) -> Vec<Annotatable> {
    expand_with_mode(..., DiffMode::Reverse)
}

Internally, both expand to the canonical #[rustc_autodiff(...)] form, where the mode is injected manually via:

1
2
3
4
5
6
let mode_symbol = match mode {
    DiffMode::Forward => sym::Forward,
    DiffMode::Reverse => sym::Reverse,
    _ => unreachable!("Unsupported mode"),
};
ts.insert(0, TokenTree::Token(Token::new(TokenKind::Ident(mode_symbol, false.into()), Span::default()), Spacing::Joint));
This post is licensed under CC BY 4.0 by the author.