From c0c6d4280c31a426f48efe15d76d565bbf464136 Mon Sep 17 00:00:00 2001 From: bit-aloo Date: Mon, 21 Apr 2025 13:22:56 +0530 Subject: [PATCH 1/3] support both trait and non-trail associated declaration --- compiler/rustc_builtin_macros/src/autodiff.rs | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) 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); } From a7c7119ed2f98501141d59795a28c2235927e601 Mon Sep 17 00:00:00 2001 From: bit-aloo Date: Mon, 21 Apr 2025 13:26:25 +0530 Subject: [PATCH 2/3] Added test --- tests/pretty/inherent_impl.pp | 41 +++++++++++++++++++++++++++++++++++ tests/pretty/inherent_impl.rs | 24 ++++++++++++++++++++ 2 files changed, 65 insertions(+) create mode 100644 tests/pretty/inherent_impl.pp create mode 100644 tests/pretty/inherent_impl.rs diff --git a/tests/pretty/inherent_impl.pp b/tests/pretty/inherent_impl.pp new file mode 100644 index 0000000000000..97ac766b6b99e --- /dev/null +++ b/tests/pretty/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/inherent_impl.rs b/tests/pretty/inherent_impl.rs new file mode 100644 index 0000000000000..59de93f7e0ffd --- /dev/null +++ b/tests/pretty/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()) + } +} From b8ca0073c8e4d1d9e324db832500beb41c129040 Mon Sep 17 00:00:00 2001 From: bit-aloo Date: Tue, 22 Apr 2025 09:53:08 +0530 Subject: [PATCH 3/3] move autodiff pretty test to a autodiff sub module --- tests/pretty/{ => autodiff}/autodiff_forward.pp | 0 tests/pretty/{ => autodiff}/autodiff_forward.rs | 0 tests/pretty/{ => autodiff}/autodiff_reverse.pp | 0 tests/pretty/{ => autodiff}/autodiff_reverse.rs | 0 tests/pretty/{ => autodiff}/inherent_impl.pp | 0 tests/pretty/{ => autodiff}/inherent_impl.rs | 0 6 files changed, 0 insertions(+), 0 deletions(-) rename tests/pretty/{ => autodiff}/autodiff_forward.pp (100%) rename tests/pretty/{ => autodiff}/autodiff_forward.rs (100%) rename tests/pretty/{ => autodiff}/autodiff_reverse.pp (100%) rename tests/pretty/{ => autodiff}/autodiff_reverse.rs (100%) rename tests/pretty/{ => autodiff}/inherent_impl.pp (100%) rename tests/pretty/{ => autodiff}/inherent_impl.rs (100%) 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/inherent_impl.pp b/tests/pretty/autodiff/inherent_impl.pp similarity index 100% rename from tests/pretty/inherent_impl.pp rename to tests/pretty/autodiff/inherent_impl.pp diff --git a/tests/pretty/inherent_impl.rs b/tests/pretty/autodiff/inherent_impl.rs similarity index 100% rename from tests/pretty/inherent_impl.rs rename to tests/pretty/autodiff/inherent_impl.rs