working batching1 tests using &f32 and Vector, failing test batching2 using Simd and Batching

Author: ZuseZ4

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -29,6 +29,8 @@ pub enum DiffMode {
     Forward,
     /// The target function, to be created using reverse mode AD.
     Reverse,
+    /// The target function, to be created using batching.
+    Batch,
 }
 
 /// Dual and Duplicated (and their Only variants) are getting lowered to the same Enzyme Activity.
@@ -69,6 +71,12 @@ pub enum DiffActivity {
     /// length of a slice/vec. This is used for safety checks on slices.
     /// The integer (if given) specifies the size of the slice element in bytes.
     FakeActivitySize(Option<u32>),
+    /// Batching mode A
+    Vector,
+    /// Batching mode B, missing implementation (only available as part of autodiff through dupv)
+    // Leaf,
+    /// Batching mode C, scalar.
+    Scalar,
 }
 
 impl DiffActivity {
@@ -130,6 +138,7 @@ impl Display for DiffMode {
             DiffMode::Source => write!(f, "Source"),
             DiffMode::Forward => write!(f, "Forward"),
             DiffMode::Reverse => write!(f, "Reverse"),
+            DiffMode::Batch => write!(f, "Batch"),
         }
     }
 }
@@ -153,6 +162,13 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
                 || activity == DiffActivity::Active
                 || activity == DiffActivity::ActiveOnly
         }
+        DiffMode::Batch => {
+            // Batching is a special case, since we don't compute derivatives wrt. the return value.
+            // We just compute derivatives wrt. the inputs, so we can ignore the return value.
+            activity == DiffActivity::Const
+                || activity == DiffActivity::Vector
+                || activity == DiffActivity::Scalar
+        }
     }
 }
 
@@ -186,6 +202,11 @@ pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
         DiffMode::Reverse => {
             matches!(activity, Active | ActiveOnly | Duplicated | DuplicatedOnly | Const)
         }
+        DiffMode::Batch => {
+            // Batching is a special case, since we don't compute derivatives wrt. the return value.
+            // We just compute derivatives wrt. the inputs, so we can ignore the return value.
+            matches!(activity, Const | Vector)
+        }
     };
 }
 
@@ -203,6 +224,8 @@ impl Display for DiffActivity {
             DiffActivity::Duplicated => write!(f, "Duplicated"),
             DiffActivity::DuplicatedOnly => write!(f, "DuplicatedOnly"),
             DiffActivity::FakeActivitySize(s) => write!(f, "FakeActivitySize({:?})", s),
+            DiffActivity::Vector => write!(f, "Vector"),
+            DiffActivity::Scalar => write!(f, "Scalar"),
         }
     }
 }
@@ -216,6 +239,7 @@ impl FromStr for DiffMode {
             "Source" => Ok(DiffMode::Source),
             "Forward" => Ok(DiffMode::Forward),
             "Reverse" => Ok(DiffMode::Reverse),
+            "Batch" => Ok(DiffMode::Batch),
             _ => Err(()),
         }
     }
@@ -235,6 +259,8 @@ impl FromStr for DiffActivity {
             "DualvOnly" => Ok(DiffActivity::DualvOnly),
             "Duplicated" => Ok(DiffActivity::Duplicated),
             "DuplicatedOnly" => Ok(DiffActivity::DuplicatedOnly),
+            "Scalar" => Ok(DiffActivity::Scalar),
+            "Vector" => Ok(DiffActivity::Vector),
             _ => Err(()),
         }
     }

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -884,6 +884,7 @@ mod llvm_enzyme {
                 DiffActivity::None | DiffActivity::FakeActivitySize(_) => {
                     panic!("Should not happen");
                 }
+                DiffActivity::Vector | DiffActivity::Scalar => todo!()
             }
             if let PatKind::Ident(_, ident, _) = arg.pat.kind {
                 idents.push(ident.clone());

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -83,6 +83,9 @@ fn match_args_from_caller_to_enzyme<'ll>(
     let enzyme_dupnoneed = cx.create_metadata("enzyme_dupnoneed".to_string()).unwrap();
     let enzyme_dupnoneedv = cx.create_metadata("enzyme_dupnoneedv".to_string()).unwrap();
 
+    let enzyme_scalar = cx.create_metadata("enzyme_scalar".to_string()).unwrap();
+    let enzyme_vector = cx.create_metadata("enzyme_vector".to_string()).unwrap();
+
     while activity_pos < inputs.len() {
         let diff_activity = inputs[activity_pos as usize];
         // Duplicated arguments received a shadow argument, into which enzyme will write the
@@ -99,7 +102,10 @@ fn match_args_from_caller_to_enzyme<'ll>(
             DiffActivity::Duplicated => (enzyme_dup, true),
             DiffActivity::DuplicatedOnly => (enzyme_dupnoneed, true),
             DiffActivity::FakeActivitySize(_) => (enzyme_const, false),
+            DiffActivity::Vector => (enzyme_vector, true),
+            DiffActivity::Scalar => (enzyme_scalar, true),
         };
+        let no_autodiff_only_batching = matches!(diff_activity, DiffActivity::Scalar | DiffActivity::Vector);
         let outer_arg = outer_args[outer_pos];
         args.push(cx.get_metadata_value(activity));
         if matches!(diff_activity, DiffActivity::Dualv) {
@@ -178,8 +184,19 @@ fn match_args_from_caller_to_enzyme<'ll>(
                 outer_pos += 2;
                 activity_pos += 1;
 
+                dbg!(&width);
+                dbg!(&outer_pos);
+                dbg!(&activity_pos);
+                dbg!(&args);
+                let limit = if no_autodiff_only_batching {
+                    // Usually we have one primal arg + width shadow args.
+                    // Here we have `width` primal args, so one less than normal.
+                    width as usize - 1
+                } else {
+                    width as usize
+                };
                 // Now, if width > 1, we need to account for that
-                for _ in 1..width {
+                for _ in 1..limit {
                     let next_outer_arg = outer_args[outer_pos];
                     args.push(next_outer_arg);
                     outer_pos += 1;
@@ -269,6 +286,12 @@ fn compute_enzyme_fn_ty<'ll>(
                 DiffMode::Reverse => {
                     todo!("Handle sret for reverse mode");
                 }
+                DiffMode::Batch => {
+                    let arr_ty = unsafe {
+                        llvm::LLVMArrayType2(inner_ret_ty, attrs.width as u64)
+                    };
+                    ret_ty = arr_ty;
+                }
                 _ => {
                     bug!("unreachable");
                 }
@@ -299,6 +322,7 @@ fn generate_enzyme_call<'ll>(
     let mut ad_name: String = match attrs.mode {
         DiffMode::Forward => "__enzyme_fwddiff",
         DiffMode::Reverse => "__enzyme_autodiff",
+        DiffMode::Batch => "__enzyme_batch",
         _ => panic!("logic bug in autodiff, unrecognized mode"),
     }
     .to_string();

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -855,6 +855,7 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
     let mode = match mode.as_str() {
         "Forward" => DiffMode::Forward,
         "Reverse" => DiffMode::Reverse,
+        "Batch" => DiffMode::Batch,
         _ => {
             span_bug!(mode.span, "rustc_autodiff attribute contains invalid mode");
         }

compiler/rustc_span/src/symbol.rs

@@ -191,6 +191,7 @@ symbols! {
         BTreeEntry,
         BTreeMap,
         BTreeSet,
+        Batching,
         BinaryHeap,
         Borrow,
         BorrowMut,

tests/codegen/autodiff/batching.rs

@@ -0,0 +1,80 @@
+//@ compile-flags: -Zautodiff=Enable -Zautodiff=NoPostopt -C opt-level=3  -Clto=fat
+//@ no-prefer-dynamic
+//@ needs-enzyme
+//
+#![feature(rustc_attrs)]
+#![feature(prelude_import)]
+#![feature(panic_internals)]
+#![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:batching.pp
+
+
+// Test that forward mode ad macros are expanded correctly.
+use std::arch::asm;
+use std::autodiff::autodiff;
+
+// CHECK: ; Function Attrs: mustprogress nofree noinline norecurse nosync nounwind nonlazybind willreturn memory(argmem: read) uwtable
+// CHECK-NEXT: define dso_local noundef float @square(ptr noalias nocapture noundef readonly align 4 dereferenceable(4) %x) unnamed_addr #2 {
+// CHECK-NEXT: start:
+// CHECK-NEXT:   %_2 = load float, ptr %x, align 4, !noundef !4
+// CHECK-NEXT:   %_0 = fmul float %_2, %_2
+// CHECK-NEXT:   ret float %_0
+// CHECK-NEXT: }
+
+// CHECK: ; Function Attrs: alwaysinline nonlazybind uwtable
+// CHECK-NEXT: define dso_local void @d_square2(ptr dead_on_unwind noalias nocapture noundef writable writeonly sret([16 x i8]) align 4 dereferenceable(16) initializes((0, 16)) %_0, ptr noalias nocapture noundef readonly align 4 dereferenceable(4) %x, ptr noalias noundef readonly align 4 dereferenceable(4) %bx_0, ptr noalias noundef readonly align 4 dereferenceable(4) %bx_1, ptr noalias noundef readonly align 4 dereferenceable(4) %bx_2) unnamed_addr #3 personality ptr @rust_eh_personality {
+// CHECK-NEXT: start:
+// CHECK-NEXT:   %0 = insertvalue [4 x ptr] undef, ptr %x, 0
+// CHECK-NEXT:   %1 = insertvalue [4 x ptr] %0, ptr %bx_0, 1
+// CHECK-NEXT:   %2 = insertvalue [4 x ptr] %1, ptr %bx_1, 2
+// CHECK-NEXT:   %3 = insertvalue [4 x ptr] %2, ptr %bx_2, 3
+// CHECK-NEXT:   %4 = call [4 x float] @batch_square([4 x ptr] %3)
+// CHECK-NEXT:   %.elt = extractvalue [4 x float] %4, 0
+// CHECK-NEXT:   store float %.elt, ptr %_0, align 4
+// CHECK-NEXT:   %_0.repack1 = getelementptr inbounds nuw i8, ptr %_0, i64 4
+// CHECK-NEXT:   %.elt2 = extractvalue [4 x float] %4, 1
+// CHECK-NEXT:   store float %.elt2, ptr %_0.repack1, align 4
+// CHECK-NEXT:   %_0.repack3 = getelementptr inbounds nuw i8, ptr %_0, i64 8
+// CHECK-NEXT:   %.elt4 = extractvalue [4 x float] %4, 2
+// CHECK-NEXT:   store float %.elt4, ptr %_0.repack3, align 4
+// CHECK-NEXT:   %_0.repack5 = getelementptr inbounds nuw i8, ptr %_0, i64 12
+// CHECK-NEXT:   %.elt6 = extractvalue [4 x float] %4, 3
+// CHECK-NEXT:   store float %.elt6, ptr %_0.repack5, align 4
+// CHECK-NEXT:   ret void
+// CHECK-NEXT: }
+
+/// Generated from:
+/// ```
+/// #[batching(d_square2, 4, Vector, Vector)]
+/// fn square(x: &f32) -> f32 {
+///   x * x
+/// }
+
+#[no_mangle]
+#[rustc_autodiff]
+#[inline(never)]
+fn square(x: &f32) -> f32 {
+    x * x
+}
+#[rustc_autodiff(Batch, 4, Vector, Vector)]
+#[no_mangle]
+#[inline(never)]
+fn d_square2(x: &f32, bx_0: &f32, bx_1: &f32, bx_2: &f32) -> [f32; 4usize] {
+    unsafe {
+        asm!("NOP", options(nomem));
+    };
+    ::core::hint::black_box(square(x));
+    ::core::hint::black_box((bx_0, bx_1, bx_2));
+    ::core::hint::black_box(<[f32; 4usize]>::default())
+}
+
+
+fn main() {}

tests/codegen/batching/batching1.rs

@@ -0,0 +1,48 @@
+//@ compile-flags: -Zautodiff=Enable -Zautodiff=NoPostopt -C opt-level=3  -Clto=fat
+//@ no-prefer-dynamic
+//@ needs-enzyme
+//
+#![feature(rustc_attrs)]
+#![feature(prelude_import)]
+#![feature(panic_internals)]
+#![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:batching.pp
+
+
+// Test that forward mode ad macros are expanded correctly.
+use std::arch::asm;
+use std::autodiff::autodiff;
+
+// Generated from:
+//// ```
+/// #[batching(d_square2, 4, Batching, Batching)]
+/// fn square(x: f32) -> f32 {
+///    x * x
+/// }
+
+#[no_mangle]
+#[rustc_autodiff]
+#[inline(never)]
+fn square(x: f32) -> f32 {
+    x * x
+}
+#[rustc_autodiff(Batch, 4, Batching, Batching)]
+#[no_mangle]
+#[inline(never)]
+fn d_square2(x: Simd<f32,4>) -> Simd<f32,4> {
+    unsafe {
+        asm!("NOP", options(nomem));
+    };
+    ::core::hint::black_box(x);
+    ::core::hint::black_box(());
+    ::core::hint::black_box(Default::default())
+}
+fn main() {}