working dupv for fwd mode

Author: ZuseZ4

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -50,6 +50,10 @@ 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,
     /// 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 {
         DiffMode::Source => false,
         DiffMode::Forward => {
             activity == DiffActivity::Dual
+                || activity == DiffActivity::Dualv
                 || activity == DiffActivity::DualOnly
                 || activity == DiffActivity::Const
         }
@@ -156,7 +161,7 @@ pub fn valid_ty_for_activity(ty: &P<Ty>, activity: DiffActivity) -> bool {
     if matches!(activity, Const) {
         return true;
     }
-    if matches!(activity, Dual | DualOnly) {
+    if matches!(activity, Dual | DualOnly | Dualv) {
         return true;
     }
     // FIXME(ZuseZ4) We should make this more robust to also
@@ -173,7 +178,7 @@ pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
         DiffMode::Error => false,
         DiffMode::Source => false,
         DiffMode::Forward => {
-            matches!(activity, Dual | DualOnly | Const)
+            matches!(activity, Dual | DualOnly | Dualv | Const)
         }
         DiffMode::Reverse => {
             matches!(activity, Active | ActiveOnly | Duplicated | DuplicatedOnly | Const)
@@ -189,6 +194,7 @@ 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::Duplicated => write!(f, "Duplicated"),
             DiffActivity::DuplicatedOnly => write!(f, "DuplicatedOnly"),
@@ -220,6 +226,7 @@ impl FromStr for DiffActivity {
             "ActiveOnly" => Ok(DiffActivity::ActiveOnly),
             "Const" => Ok(DiffActivity::Const),
             "Dual" => Ok(DiffActivity::Dual),
+            "Dualv" => Ok(DiffActivity::Dualv),
             "DualOnly" => Ok(DiffActivity::DualOnly),
             "Duplicated" => Ok(DiffActivity::Duplicated),
             "DuplicatedOnly" => Ok(DiffActivity::DuplicatedOnly),

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -799,8 +799,13 @@ mod llvm_enzyme {
                         d_inputs.push(shadow_arg.clone());
                     }
                 }
-                DiffActivity::Dual | DiffActivity::DualOnly => {
-                    for i in 0..x.width {
+                DiffActivity::Dual | DiffActivity::DualOnly | DiffActivity::Dualv => {
+                    let iterations = if matches!(activity, DiffActivity::Dualv) {
+                        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 {
                             ident.name
@@ -887,8 +892,8 @@ mod llvm_enzyme {
                 }
             };
 
-            if let DiffActivity::Dual = x.ret_activity {
-                let kind = if x.width == 1 {
+            if matches!(x.ret_activity, DiffActivity::Dual | DiffActivity::Dualv) {
+                let kind = if x.width == 1 || matches!(x.ret_activity, DiffActivity::Dualv) {
                     // Dual can only be used for f32/f64 ret.
                     // In that case we return now a tuple with two floats.
                     TyKind::Tup(thin_vec![ty.clone(), ty.clone()])

compiler/rustc_codegen_llvm/src/builder.rs

@@ -123,7 +123,7 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
 /// Empty string, to be used where LLVM expects an instruction name, indicating
 /// that the instruction is to be left unnamed (i.e. numbered, in textual IR).
 // FIXME(eddyb) pass `&CStr` directly to FFI once it's a thin pointer.
-const UNNAMED: *const c_char = c"".as_ptr();
+pub(crate) const UNNAMED: *const c_char = c"".as_ptr();
 
 impl<'ll, CX: Borrow<SCx<'ll>>> BackendTypes for GenericBuilder<'_, 'll, CX> {
     type Value = <GenericCx<'ll, CX> as BackendTypes>::Value;

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -10,7 +10,7 @@ use rustc_middle::bug;
 use tracing::{debug, trace};
 
 use crate::back::write::llvm_err;
-use crate::builder::SBuilder;
+use crate::builder::{SBuilder, UNNAMED};
 use crate::context::SimpleCx;
 use crate::declare::declare_simple_fn;
 use crate::errors::{AutoDiffWithoutEnable, LlvmError};
@@ -51,6 +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>,
     width: u32,
     args: &mut Vec<&'ll llvm::Value>,
     inputs: &[DiffActivity],
@@ -78,6 +79,7 @@ fn match_args_from_caller_to_enzyme<'ll>(
     let enzyme_const = cx.create_metadata("enzyme_const".to_string()).unwrap();
     let enzyme_out = cx.create_metadata("enzyme_out".to_string()).unwrap();
     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();
 
     while activity_pos < inputs.len() {
@@ -90,13 +92,26 @@ fn match_args_from_caller_to_enzyme<'ll>(
             DiffActivity::Active => (enzyme_out, false),
             DiffActivity::ActiveOnly => (enzyme_out, false),
             DiffActivity::Dual => (enzyme_dup, true),
+            DiffActivity::Dualv => (enzyme_dupv, true),
             DiffActivity::DualOnly => (enzyme_dupnoneed, true),
             DiffActivity::Duplicated => (enzyme_dup, true),
             DiffActivity::DuplicatedOnly => (enzyme_dupnoneed, true),
             DiffActivity::FakeActivitySize => (enzyme_const, false),
         };
         let outer_arg = outer_args[outer_pos];
         args.push(cx.get_metadata_value(activity));
+        if matches!(diff_activity, DiffActivity::Dualv) {
+            let next_outer_arg = outer_args[outer_pos + 1];
+            // stride: sizeof(T) * n_elems.
+            // 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)};
+            args.push(mul);
+        }
         args.push(outer_arg);
         if duplicated {
             // We know that duplicated args by construction have a following argument,
@@ -125,7 +140,13 @@ 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);
 
-                for i in 0..(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)];
                     let next_outer_ty2 = cx.val_ty(next_outer_arg2);
                     assert_eq!(cx.type_kind(next_outer_ty2), TypeKind::Pointer);
@@ -136,7 +157,7 @@ fn match_args_from_caller_to_enzyme<'ll>(
                 }
                 args.push(cx.get_metadata_value(enzyme_const));
                 args.push(next_outer_arg);
-                outer_pos += 2 + 2 * width as usize;
+                outer_pos += 2 + 2 * iterations;
                 activity_pos += 2;
             } else {
                 // A duplicated pointer will have the following two outer_fn arguments:
@@ -360,6 +381,7 @@ fn generate_enzyme_call<'ll>(
         let outer_args: Vec<&llvm::Value> = get_params(outer_fn);
         match_args_from_caller_to_enzyme(
             &cx,
+            &builder,
             attrs.width,
             &mut args,
             &attrs.input_activity,

compiler/rustc_monomorphize/src/partitioning/autodiff.rs

@@ -31,6 +31,7 @@ fn adjust_activity_to_abi<'tcx>(tcx: TyCtxt<'tcx>, fn_ty: Ty<'tcx>, da: &mut Vec
                     let activity = match da[i] {
                         DiffActivity::DualOnly
                         | DiffActivity::Dual
+                        | DiffActivity::Dualv
                         | DiffActivity::DuplicatedOnly
                         | DiffActivity::Duplicated => DiffActivity::FakeActivitySize,
                         DiffActivity::Const => DiffActivity::Const,