diff --git a/compiler/rustc_builtin_macros/src/autodiff.rs b/compiler/rustc_builtin_macros/src/autodiff.rs index daebd516499bc..b33856ac23661 100644 --- a/compiler/rustc_builtin_macros/src/autodiff.rs +++ b/compiler/rustc_builtin_macros/src/autodiff.rs @@ -217,14 +217,12 @@ mod llvm_enzyme { ast::StmtKind::Item(iitem) => extract_item_info(iitem), _ => None, }, - Annotatable::AssocItem(assoc_item, Impl { of_trait: false }) => { - match &assoc_item.kind { - ast::AssocItemKind::Fn(box ast::Fn { sig, ident, .. }) => { - Some((assoc_item.vis.clone(), sig.clone(), ident.clone())) - } - _ => None, + Annotatable::AssocItem(assoc_item, Impl { .. }) => match &assoc_item.kind { + ast::AssocItemKind::Fn(box ast::Fn { sig, ident, .. }) => { + Some((assoc_item.vis.clone(), sig.clone(), ident.clone())) } - } + _ => None, + }, _ => None, }) else { dcx.emit_err(errors::AutoDiffInvalidApplication { span: item.span() }); @@ -365,7 +363,7 @@ mod llvm_enzyme { } Annotatable::Item(iitem.clone()) } - Annotatable::AssocItem(ref mut assoc_item, i @ Impl { of_trait: false }) => { + Annotatable::AssocItem(ref mut assoc_item, i @ Impl { .. }) => { if !assoc_item.attrs.iter().any(|a| same_attribute(&a.kind, &attr.kind)) { assoc_item.attrs.push(attr); } diff --git a/tests/pretty/autodiff_forward.pp b/tests/pretty/autodiff/autodiff_forward.pp similarity index 100% rename from tests/pretty/autodiff_forward.pp rename to tests/pretty/autodiff/autodiff_forward.pp diff --git a/tests/pretty/autodiff_forward.rs b/tests/pretty/autodiff/autodiff_forward.rs similarity index 100% rename from tests/pretty/autodiff_forward.rs rename to tests/pretty/autodiff/autodiff_forward.rs diff --git a/tests/pretty/autodiff_reverse.pp b/tests/pretty/autodiff/autodiff_reverse.pp similarity index 100% rename from tests/pretty/autodiff_reverse.pp rename to tests/pretty/autodiff/autodiff_reverse.pp diff --git a/tests/pretty/autodiff_reverse.rs b/tests/pretty/autodiff/autodiff_reverse.rs similarity index 100% rename from tests/pretty/autodiff_reverse.rs rename to tests/pretty/autodiff/autodiff_reverse.rs diff --git a/tests/pretty/autodiff/inherent_impl.pp b/tests/pretty/autodiff/inherent_impl.pp new file mode 100644 index 0000000000000..97ac766b6b99e --- /dev/null +++ b/tests/pretty/autodiff/inherent_impl.pp @@ -0,0 +1,41 @@ +#![feature(prelude_import)] +#![no_std] +//@ needs-enzyme + +#![feature(autodiff)] +#[prelude_import] +use ::std::prelude::rust_2015::*; +#[macro_use] +extern crate std; +//@ pretty-mode:expanded +//@ pretty-compare-only +//@ pp-exact:inherent_impl.pp + +use std::autodiff::autodiff; + +struct Foo { + a: f64, +} + +trait MyTrait { + fn f(&self, x: f64) + -> f64; + fn df(&self, x: f64, seed: f64) + -> (f64, f64); +} + +impl MyTrait for Foo { + #[rustc_autodiff] + #[inline(never)] + fn f(&self, x: f64) -> f64 { + self.a * 0.25 * (x * x - 1.0 - 2.0 * x.ln()) + } + #[rustc_autodiff(Reverse, 1, Const, Active, Active)] + #[inline(never)] + fn df(&self, x: f64, dret: f64) -> (f64, f64) { + unsafe { asm!("NOP", options(pure, nomem)); }; + ::core::hint::black_box(self.f(x)); + ::core::hint::black_box((dret,)); + ::core::hint::black_box((self.f(x), f64::default())) + } +} diff --git a/tests/pretty/autodiff/inherent_impl.rs b/tests/pretty/autodiff/inherent_impl.rs new file mode 100644 index 0000000000000..59de93f7e0ffd --- /dev/null +++ b/tests/pretty/autodiff/inherent_impl.rs @@ -0,0 +1,24 @@ +//@ needs-enzyme + +#![feature(autodiff)] +//@ pretty-mode:expanded +//@ pretty-compare-only +//@ pp-exact:inherent_impl.pp + +use std::autodiff::autodiff; + +struct Foo { + a: f64, +} + +trait MyTrait { + fn f(&self, x: f64) -> f64; + fn df(&self, x: f64, seed: f64) -> (f64, f64); +} + +impl MyTrait for Foo { + #[autodiff(df, Reverse, Const, Active, Active)] + fn f(&self, x: f64) -> f64 { + self.a * 0.25 * (x * x - 1.0 - 2.0 * x.ln()) + } +}