working dupv and dupvonly for fwd mode

1 files changed, 26 insertions, 14 deletions

diff --git a/compiler/rustc_ast/src/expand/autodiff_attrs.rs b/compiler/rustc_ast/src/expand/autodiff_attrs.rs
index 13a7c5a1805..2f918faaf75 100644
--- a/compiler/rustc_ast/src/expand/autodiff_attrs.rs
+++ b/compiler/rustc_ast/src/expand/autodiff_attrs.rs
@@ -50,8 +50,16 @@ pub enum DiffActivity {
     /// with it.
     Dual,
     /// Forward Mode, Compute derivatives for this input/output and *overwrite* the shadow argument
+    /// with it. It expects the shadow argument to be `width` times larger than the original
+    /// input/output.
+    Dualv,
+    /// Forward Mode, Compute derivatives for this input/output and *overwrite* the shadow argument
     /// with it. Drop the code which updates the original input/output for maximum performance.
     DualOnly,
+    /// Forward Mode, Compute derivatives for this input/output and *overwrite* the shadow argument
+    /// with it. Drop the code which updates the original input/output for maximum performance.
+    /// It expects the shadow argument to be `width` times larger than the original input/output.
+    DualvOnly,
     /// Reverse Mode, Compute derivatives for this &T or *T input and *add* it to the shadow argument.
     Duplicated,
     /// Reverse Mode, Compute derivatives for this &T or *T input and *add* it to the shadow argument.
@@ -59,7 +67,15 @@ pub enum DiffActivity {
     DuplicatedOnly,
     /// All Integers must be Const, but these are used to mark the integer which represents the
     /// length of a slice/vec. This is used for safety checks on slices.
-    FakeActivitySize,
+    /// The integer (if given) specifies the size of the slice element in bytes.
+    FakeActivitySize(Option<u32>),
+}
+
+impl DiffActivity {
+    pub fn is_dual_or_const(&self) -> bool {
+        use DiffActivity::*;
+        matches!(self, |Dual| DualOnly | Dualv | DualvOnly | Const)
+    }
 }
 /// We generate one of these structs for each `#[autodiff(...)]` attribute.
 #[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
@@ -131,11 +147,7 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
     match mode {
         DiffMode::Error => false,
         DiffMode::Source => false,
-        DiffMode::Forward => {
-            activity == DiffActivity::Dual
-                || activity == DiffActivity::DualOnly
-                || activity == DiffActivity::Const
-        }
+        DiffMode::Forward => activity.is_dual_or_const(),
         DiffMode::Reverse => {
             activity == DiffActivity::Const
                 || activity == DiffActivity::Active
@@ -153,10 +165,8 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
 pub fn valid_ty_for_activity(ty: &P<Ty>, activity: DiffActivity) -> bool {
     use DiffActivity::*;
     // It's always allowed to mark something as Const, since we won't compute derivatives wrt. it.
-    if matches!(activity, Const) {
-        return true;
-    }
-    if matches!(activity, Dual | DualOnly) {
+    // Dual variants also support all types.
+    if activity.is_dual_or_const() {
         return true;
     }
     // FIXME(ZuseZ4) We should make this more robust to also
@@ -172,9 +182,7 @@ pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
     return match mode {
         DiffMode::Error => false,
         DiffMode::Source => false,
-        DiffMode::Forward => {
-            matches!(activity, Dual | DualOnly | Const)
-        }
+        DiffMode::Forward => activity.is_dual_or_const(),
         DiffMode::Reverse => {
             matches!(activity, Active | ActiveOnly | Duplicated | DuplicatedOnly | Const)
         }
@@ -189,10 +197,12 @@ impl Display for DiffActivity {
             DiffActivity::Active => write!(f, "Active"),
             DiffActivity::ActiveOnly => write!(f, "ActiveOnly"),
             DiffActivity::Dual => write!(f, "Dual"),
+            DiffActivity::Dualv => write!(f, "Dualv"),
             DiffActivity::DualOnly => write!(f, "DualOnly"),
+            DiffActivity::DualvOnly => write!(f, "DualvOnly"),
             DiffActivity::Duplicated => write!(f, "Duplicated"),
             DiffActivity::DuplicatedOnly => write!(f, "DuplicatedOnly"),
-            DiffActivity::FakeActivitySize => write!(f, "FakeActivitySize"),
+            DiffActivity::FakeActivitySize(s) => write!(f, "FakeActivitySize({:?})", s),
         }
     }
 }
@@ -220,7 +230,9 @@ impl FromStr for DiffActivity {
             "ActiveOnly" => Ok(DiffActivity::ActiveOnly),
             "Const" => Ok(DiffActivity::Const),
             "Dual" => Ok(DiffActivity::Dual),
+            "Dualv" => Ok(DiffActivity::Dualv),
             "DualOnly" => Ok(DiffActivity::DualOnly),
+            "DualvOnly" => Ok(DiffActivity::DualvOnly),
             "Duplicated" => Ok(DiffActivity::Duplicated),
             "DuplicatedOnly" => Ok(DiffActivity::DuplicatedOnly),
             _ => Err(()),