working dupvonly for fwd mode

Author: ZuseZ4

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -56,6 +56,10 @@ pub enum DiffActivity {
     /// 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.
@@ -133,6 +137,7 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
             activity == DiffActivity::Dual
                 || activity == DiffActivity::Dualv
                 || activity == DiffActivity::DualOnly
+                || activity == DiffActivity::DualvOnly
                 || activity == DiffActivity::Const
         }
         DiffMode::Reverse => {
@@ -155,7 +160,7 @@ pub fn valid_ty_for_activity(ty: &P<Ty>, activity: DiffActivity) -> bool {
     if matches!(activity, Const) {
         return true;
     }
-    if matches!(activity, Dual | DualOnly | Dualv) {
+    if matches!(activity, Dual | DualOnly | Dualv | DualvOnly) {
         return true;
     }
     // FIXME(ZuseZ4) We should make this more robust to also
@@ -172,7 +177,7 @@ pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
         DiffMode::Error => false,
         DiffMode::Source => false,
         DiffMode::Forward => {
-            matches!(activity, Dual | DualOnly | Dualv | Const)
+            matches!(activity, Dual | DualOnly | Dualv | DualvOnly | Const)
         }
         DiffMode::Reverse => {
             matches!(activity, Active | ActiveOnly | Duplicated | DuplicatedOnly | Const)
@@ -190,6 +195,7 @@ impl Display for DiffActivity {
             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"),
@@ -222,6 +228,7 @@ impl FromStr for DiffActivity {
             "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(()),

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -770,12 +770,18 @@ mod llvm_enzyme {
                         d_inputs.push(shadow_arg.clone());
                     }
                 }
-                DiffActivity::Dual | DiffActivity::DualOnly | DiffActivity::Dualv => {
-                    let iterations = if matches!(activity, DiffActivity::Dualv) {
-                        1
-                    } else {
-                        x.width
-                    };
+                DiffActivity::Dual
+                | DiffActivity::DualOnly
+                | DiffActivity::Dualv
+                | DiffActivity::DualvOnly => {
+                    // the *v variants get lowered to enzyme_dupv and enzyme_dupnoneedv, which cause
+                    // Enzyme to not expect N arguments, but one argument (which is instead larger).
+                    let iterations =
+                        if matches!(activity, DiffActivity::Dualv | DiffActivity::DualvOnly) {
+                            1
+                        } else {
+                            x.width
+                        };
                     for i in 0..iterations {
                         let mut shadow_arg = arg.clone();
                         let old_name = if let PatKind::Ident(_, ident, _) = arg.pat.kind {
@@ -879,7 +885,7 @@ mod llvm_enzyme {
                 let ty = P(rustc_ast::Ty { kind, id: ty.id, span: ty.span, tokens: None });
                 d_decl.output = FnRetTy::Ty(ty);
             }
-            if let DiffActivity::DualOnly = x.ret_activity {
+            if matches!(x.ret_activity, DiffActivity::DualOnly | DiffActivity::DualvOnly) {
                 // No need to change the return type,
                 // we will just return the shadow in place of the primal return.
                 // However, if we have a width > 1, then we don't return -> T, but -> [T; width]

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -51,7 +51,7 @@ fn has_sret(fnc: &Value) -> bool {
 // using iterators and peek()?
 fn match_args_from_caller_to_enzyme<'ll>(
     cx: &SimpleCx<'ll>,
-    builder: &SBuilder<'ll,'ll>,
+    builder: &SBuilder<'ll, 'll>,
     width: u32,
     args: &mut Vec<&'ll llvm::Value>,
     inputs: &[DiffActivity],
@@ -81,6 +81,7 @@ fn match_args_from_caller_to_enzyme<'ll>(
     let enzyme_dup = cx.create_metadata("enzyme_dup".to_string()).unwrap();
     let enzyme_dupv = cx.create_metadata("enzyme_dupv".to_string()).unwrap();
     let enzyme_dupnoneed = cx.create_metadata("enzyme_dupnoneed".to_string()).unwrap();
+    let enzyme_dupnoneedv = cx.create_metadata("enzyme_dupnoneedv".to_string()).unwrap();
 
     while activity_pos < inputs.len() {
         let diff_activity = inputs[activity_pos as usize];
@@ -94,6 +95,7 @@ fn match_args_from_caller_to_enzyme<'ll>(
             DiffActivity::Dual => (enzyme_dup, true),
             DiffActivity::Dualv => (enzyme_dupv, true),
             DiffActivity::DualOnly => (enzyme_dupnoneed, true),
+            DiffActivity::DualvOnly => (enzyme_dupnoneedv, true),
             DiffActivity::Duplicated => (enzyme_dup, true),
             DiffActivity::DuplicatedOnly => (enzyme_dupnoneed, true),
             DiffActivity::FakeActivitySize => (enzyme_const, false),
@@ -106,10 +108,9 @@ fn match_args_from_caller_to_enzyme<'ll>(
             // T=f32 => 4 bytes
             // n_elems is the next integer.
             // Now we multiply `4 * next_outer_arg` to get the stride.
-            //let mul = builder
-            //    .build_mul(cx.get_const_i64(4), next_outer_arg)
-            //    .unwrap();
-            let mul = unsafe {llvm::LLVMBuildMul(builder.llbuilder, cx.get_const_i64(4), next_outer_arg, UNNAMED)};
+            let mul = unsafe {
+                llvm::LLVMBuildMul(builder.llbuilder, cx.get_const_i64(4), next_outer_arg, UNNAMED)
+            };
             args.push(mul);
         }
         args.push(outer_arg);
@@ -140,11 +141,8 @@ fn match_args_from_caller_to_enzyme<'ll>(
                 // int2 >= int1, which means the shadow vector is large enough to store the gradient.
                 assert_eq!(cx.type_kind(next_outer_ty), TypeKind::Integer);
 
-                let iterations = if matches!(diff_activity, DiffActivity::Dualv) {
-                    1
-                } else {
-                    width as usize
-                };
+                let iterations =
+                    if matches!(diff_activity, DiffActivity::Dualv) { 1 } else { width as usize };
 
                 for i in 0..iterations {
                     let next_outer_arg2 = outer_args[outer_pos + 2 * (i + 1)];

compiler/rustc_monomorphize/src/partitioning/autodiff.rs

@@ -40,6 +40,9 @@ fn adjust_activity_to_abi<'tcx>(tcx: TyCtxt<'tcx>, fn_ty: Ty<'tcx>, da: &mut Vec
                     new_activities.push(activity);
                     new_positions.push(i + 1);
                 }
+                // Now we need to figure out the size of each slice element in memory.
+                // Can we actually do that here?
+
                 continue;
             }
         }