rust-lang/rust
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Split autodiff into autodiff_forward and autodiff_reverse

Browse Source 
Pending fix.
```
error: cannot find a built-in macro with name `autodiff_forward`
    --> library\core\src\macros\mod.rs:1542:5
     |
1542 | /     pub macro autodiff_forward($item:item) {
1543 | |         /* compiler built-in */
1544 | |     }
     | |_____^

error: cannot find a built-in macro with name `autodiff_reverse`
    --> library\core\src\macros\mod.rs:1549:5
     |
1549 | /     pub macro autodiff_reverse($item:item) {
1550 | |         /* compiler built-in */
1551 | |     }
     | |_____^

error: could not compile `core` (lib) due to 2 previous errors
```
This commit is contained in:
Marcelo Domínguez
2025-05-06 09:19:33 +02:00
committed by Marcelo Domínguez
parent f8e9e7636a
commit b21c9e7bfb
6 changed files with 48 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
compiler/rustc_builtin_macros/src/autodiff.rs
View File

@@ -88,25 +88,20 @@ mod llvm_enzyme {
has_ret: bool,
) -> AutoDiffAttrs {
let dcx = ecx.sess.dcx();
let mode = name(&meta_item[1]);
let Ok(mode) = DiffMode::from_str(&mode) else {
dcx.emit_err(errors::AutoDiffInvalidMode { span: meta_item[1].span(), mode });
return AutoDiffAttrs::error();
};
// Now we check, whether the user wants autodiff in batch/vector mode, or scalar mode.
// If he doesn't specify an integer (=width), we default to scalar mode, thus width=1.
let mut first_activity = 2;
let mut first_activity = 1;
let width = if let [_, _, x, ..] = &meta_item[..]
let width = if let [_, x, ..] = &meta_item[..]
&& let Some(x) = width(x)
{
first_activity = 3;
first_activity = 2;
match x.try_into() {
Ok(x) => x,
Err(_) => {
dcx.emit_err(errors::AutoDiffInvalidWidth {
span: meta_item[2].span(),
span: meta_item[1].span(),
width: x,
});
return AutoDiffAttrs::error();
@@ -150,7 +145,7 @@ mod llvm_enzyme {
};
AutoDiffAttrs {
mode,
mode: DiffMode::Error,
width,
ret_activity: *ret_activity,
input_activity: input_activity.to_vec(),
@@ -165,6 +160,24 @@ mod llvm_enzyme {
ts.push(TokenTree::Token(comma.clone(), Spacing::Alone));
}
pub(crate) fn expand_forward(
ecx: &mut ExtCtxt<'_>,
expand_span: Span,
meta_item: &ast::MetaItem,
item: Annotatable,
) -> Vec<Annotatable> {
expand_with_mode(ecx, expand_span, meta_item, item, DiffMode::Forward)
}
pub(crate) fn expand_reverse(
ecx: &mut ExtCtxt<'_>,
expand_span: Span,
meta_item: &ast::MetaItem,
item: Annotatable,
) -> Vec<Annotatable> {
expand_with_mode(ecx, expand_span, meta_item, item, DiffMode::Reverse)
}
/// We expand the autodiff macro to generate a new placeholder function which passes
/// type-checking and can be called by users. The function body of the placeholder function will
/// later be replaced on LLVM-IR level, so the design of the body is less important and for now
@@ -198,11 +211,12 @@ mod llvm_enzyme {
/// ```
/// FIXME(ZuseZ4): Once autodiff is enabled by default, make this a doc comment which is checked
/// in CI.
pub(crate) fn expand(
pub(crate) fn expand_with_mode(
ecx: &mut ExtCtxt<'_>,
expand_span: Span,
meta_item: &ast::MetaItem,
mut item: Annotatable,
mode: DiffMode,
) -> Vec<Annotatable> {
if cfg!(not(llvm_enzyme)) {
ecx.sess.dcx().emit_err(errors::AutoDiffSupportNotBuild { span: meta_item.span });
@@ -289,7 +303,8 @@ mod llvm_enzyme {
ts.pop();
let ts: TokenStream = TokenStream::from_iter(ts);
let x: AutoDiffAttrs = from_ast(ecx, &meta_item_vec, has_ret);
let mut x: AutoDiffAttrs = from_ast(ecx, &meta_item_vec, has_ret);
x.mode = mode;
if !x.is_active() {
// We encountered an error, so we return the original item.
// This allows us to potentially parse other attributes.
@@ -1017,4 +1032,4 @@ mod llvm_enzyme {
}
}
pub(crate) use llvm_enzyme::expand;
pub(crate) use llvm_enzyme::{expand_forward, expand_reverse};