Rollup merge of #136419 - EnzymeAD:autodiff-tests, r=onur-ozkan,jieyouxu

adding autodiff tests

I'd like to get started with upstreaming some tests, even though I'm still waiting for an answer on how to best integrate the enzyme pass. Can we therefore temporarily support the -Z llvm-plugins here without too much effort? And in that case, how would that work? I saw you can do remapping, e.g. `rust-src-base`, but I don't think that will give me the path to libEnzyme.so. Do you have another suggestion?

Other than that this test simply checks that the derivative of `x*x` is `2.0 * x`, which in this case is computed as
`%0 = fadd fast double %x.0.val, %x.0.val`
(I'll add a few more tests and move it to an autodiff folder if we can use the -Z flag)

r? ``@jieyouxu``

Locally at least `-Zllvm-plugins=${PWD}/build/x86_64-unknown-linux-gnu/enzyme/build/Enzyme/libEnzyme-19.so` seems to work if I copy the command I get from x.py test and run it manually. However, running x.py test itself fails.

Tracking:

- https://github.com/rust-lang/rust/issues/124509

Zulip discussion: https://rust-lang.zulipchat.com/#narrow/channel/326414-t-infra.2Fbootstrap/topic/Enzyme.20build.20changes

1 files changed, 6 insertions, 18 deletions

diff --git a/compiler/rustc_ast/src/expand/autodiff_attrs.rs b/compiler/rustc_ast/src/expand/autodiff_attrs.rs
index ecc522ec39d..70222f4acab 100644
--- a/compiler/rustc_ast/src/expand/autodiff_attrs.rs
+++ b/compiler/rustc_ast/src/expand/autodiff_attrs.rs
@@ -30,14 +30,6 @@ pub enum DiffMode {
     Forward,
     /// The target function, to be created using reverse mode AD.
     Reverse,
-    /// The target function, to be created using forward mode AD.
-    /// This target function will also be used as a source for higher order derivatives,
-    /// so compute it before all Forward/Reverse targets and optimize it through llvm.
-    ForwardFirst,
-    /// The target function, to be created using reverse mode AD.
-    /// This target function will also be used as a source for higher order derivatives,
-    /// so compute it before all Forward/Reverse targets and optimize it through llvm.
-    ReverseFirst,
 }
 
 /// Dual and Duplicated (and their Only variants) are getting lowered to the same Enzyme Activity.
@@ -92,10 +84,10 @@ pub struct AutoDiffAttrs {
 
 impl DiffMode {
     pub fn is_rev(&self) -> bool {
-        matches!(self, DiffMode::Reverse | DiffMode::ReverseFirst)
+        matches!(self, DiffMode::Reverse)
     }
     pub fn is_fwd(&self) -> bool {
-        matches!(self, DiffMode::Forward | DiffMode::ForwardFirst)
+        matches!(self, DiffMode::Forward)
     }
 }
 
@@ -106,8 +98,6 @@ impl Display for DiffMode {
             DiffMode::Source => write!(f, "Source"),
             DiffMode::Forward => write!(f, "Forward"),
             DiffMode::Reverse => write!(f, "Reverse"),
-            DiffMode::ForwardFirst => write!(f, "ForwardFirst"),
-            DiffMode::ReverseFirst => write!(f, "ReverseFirst"),
         }
     }
 }
@@ -125,12 +115,12 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
     match mode {
         DiffMode::Error => false,
         DiffMode::Source => false,
-        DiffMode::Forward | DiffMode::ForwardFirst => {
+        DiffMode::Forward => {
             activity == DiffActivity::Dual
                 || activity == DiffActivity::DualOnly
                 || activity == DiffActivity::Const
         }
-        DiffMode::Reverse | DiffMode::ReverseFirst => {
+        DiffMode::Reverse => {
             activity == DiffActivity::Const
                 || activity == DiffActivity::Active
                 || activity == DiffActivity::ActiveOnly
@@ -166,10 +156,10 @@ pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
     return match mode {
         DiffMode::Error => false,
         DiffMode::Source => false,
-        DiffMode::Forward | DiffMode::ForwardFirst => {
+        DiffMode::Forward => {
             matches!(activity, Dual | DualOnly | Const)
         }
-        DiffMode::Reverse | DiffMode::ReverseFirst => {
+        DiffMode::Reverse => {
             matches!(activity, Active | ActiveOnly | Duplicated | DuplicatedOnly | Const)
         }
     };
@@ -200,8 +190,6 @@ impl FromStr for DiffMode {
             "Source" => Ok(DiffMode::Source),
             "Forward" => Ok(DiffMode::Forward),
             "Reverse" => Ok(DiffMode::Reverse),
-            "ForwardFirst" => Ok(DiffMode::ForwardFirst),
-            "ReverseFirst" => Ok(DiffMode::ReverseFirst),
             _ => Err(()),
         }
     }