add non-working enzyme_buffer mode, with incorrect logic

Author: ZuseZ4

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -73,9 +73,9 @@ pub enum DiffActivity {
     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.
+    /// Batching mode B, equivalent to *v modes above
+    Buffer,
+    /// "Batching" mode C, scalar. Not batched.
     Scalar,
 }
 
@@ -167,6 +167,7 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
             // We just compute derivatives wrt. the inputs, so we can ignore the return value.
             activity == DiffActivity::Const
                 || activity == DiffActivity::Vector
+                || activity == DiffActivity::Buffer
                 || activity == DiffActivity::Scalar
         }
     }
@@ -205,7 +206,7 @@ pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
         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)
+            matches!(activity, Const | Vector | Buffer)
         }
     };
 }
@@ -225,6 +226,7 @@ impl Display for DiffActivity {
             DiffActivity::DuplicatedOnly => write!(f, "DuplicatedOnly"),
             DiffActivity::FakeActivitySize(s) => write!(f, "FakeActivitySize({:?})", s),
             DiffActivity::Vector => write!(f, "Vector"),
+            DiffActivity::Buffer => write!(f, "Buffer"),
             DiffActivity::Scalar => write!(f, "Scalar"),
         }
     }
@@ -261,6 +263,7 @@ impl FromStr for DiffActivity {
             "DuplicatedOnly" => Ok(DiffActivity::DuplicatedOnly),
             "Scalar" => Ok(DiffActivity::Scalar),
             "Vector" => Ok(DiffActivity::Vector),
+            "Buffer" => Ok(DiffActivity::Buffer),
             _ => Err(()),
         }
     }

compiler/rustc_builtin_macros/src/autodiff.rs

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

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -76,15 +76,18 @@ fn match_args_from_caller_to_enzyme<'ll>(
         outer_pos = 1;
     }
 
+    // Autodiff activities
     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();
     let enzyme_dupnoneedv = cx.create_metadata("enzyme_dupnoneedv".to_string()).unwrap();
 
+    // Batching activities
     let enzyme_scalar = cx.create_metadata("enzyme_scalar".to_string()).unwrap();
     let enzyme_vector = cx.create_metadata("enzyme_vector".to_string()).unwrap();
+    let enzyme_buffer = cx.create_metadata("enzyme_buffer".to_string()).unwrap();
 
     while activity_pos < inputs.len() {
         let diff_activity = inputs[activity_pos as usize];
@@ -103,16 +106,17 @@ fn match_args_from_caller_to_enzyme<'ll>(
             DiffActivity::DuplicatedOnly => (enzyme_dupnoneed, true),
             DiffActivity::FakeActivitySize(_) => (enzyme_const, false),
             DiffActivity::Vector => (enzyme_vector, true),
+            DiffActivity::Buffer => (enzyme_buffer, false),
             DiffActivity::Scalar => (enzyme_scalar, true),
         };
-        let no_autodiff_only_batching = matches!(diff_activity, DiffActivity::Scalar | DiffActivity::Vector);
+        let no_autodiff_only_batching = matches!(diff_activity, DiffActivity::Scalar | DiffActivity::Vector | DiffActivity::Buffer);
         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];
             let elem_bytes_size: u64 = match inputs[activity_pos + 1] {
                 DiffActivity::FakeActivitySize(Some(s)) => s.into(),
-                _ => bug!("incorrect Dualv handling recognized."),
+                _ => bug!("incorrect Dualv/Batching handling recognized."),
             };
             // stride: sizeof(T) * n_elems.
             // n_elems is the next integer.
@@ -127,7 +131,53 @@ fn match_args_from_caller_to_enzyme<'ll>(
             };
             args.push(mul);
         }
+        if matches!(diff_activity, DiffActivity::Buffer) {
+            // There are various cases.
+            // A) We look at a scalar float.
+            // B) We look at a Vector/Array of floats (byVal). Not sure if this is valid.
+            // C) We look at a ptr as part of a slice.
+            // D) We look at a ptr as part of a raw pointer or reference.
+
+            let mut elem_offset = cx.get_const_i64(width.into());
+            let outer_ty = cx.val_ty(outer_arg);
+            dbg!(&outer_ty);
+            let bit_width = if cx.is_float_type(outer_ty) {
+                cx.float_width(outer_ty)
+            } else if cx.is_vec_or_array_type(outer_ty) {
+                let elem_ty = cx.element_type(outer_ty);
+                assert!(cx.is_float_type(elem_ty));
+                let num_vec_elements = cx.vector_length(outer_ty);
+                assert!(num_vec_elements == width as usize);
+                dbg!(&num_vec_elements);
+                cx.float_width(elem_ty)
+            } else if cx.is_ptr_type(outer_ty) {
+                if is_slice(activity_pos, inputs) {
+                    elem_offset = outer_args[outer_pos + 1];
+                    let elem_bytes_size: u64 = match inputs[activity_pos + 1] {
+                        DiffActivity::FakeActivitySize(Some(s)) => s.into(),
+                        _ => bug!("incorrect Dualv/Buffer handling recognized."),
+                    };
+                    elem_bytes_size as usize * 8
+                } else {
+                    // raw pointer or ref, hence `num_elem` = 1
+                    unimplemented!()
+                }
+            } else {
+                bug!("expected float or vector type, found {:?}", outer_ty);
+            };
+            let elem_bytes_size = bit_width as u64 / 8;
+            let mul = unsafe {
+                llvm::LLVMBuildMul(
+                    builder.llbuilder,
+                    cx.get_const_i64(elem_bytes_size),
+                    elem_offset,
+                    UNNAMED,
+                )
+            };
+            args.push(mul);
+        }
         args.push(outer_arg);
+        dbg!(&args);
         if duplicated {
             // We know that duplicated args by construction have a following argument,
             // so this can not be out of bounds.
@@ -136,17 +186,7 @@ fn match_args_from_caller_to_enzyme<'ll>(
             // FIXME(ZuseZ4): We should add support for Vec here too, but it's less urgent since
             // vectors behind references (&Vec<T>) are already supported. Users can not pass a
             // Vec by value for reverse mode, so this would only help forward mode autodiff.
-            let slice = {
-                if activity_pos + 1 >= inputs.len() {
-                    // If there is no arg following our ptr, it also can't be a slice,
-                    // since that would lead to a ptr, int pair.
-                    false
-                } else {
-                    let next_activity = inputs[activity_pos + 1];
-                    // We analyze the MIR types and add this dummy activity if we visit a slice.
-                    matches!(next_activity, DiffActivity::FakeActivitySize(_))
-                }
-            };
+            let slice = is_slice(activity_pos, &inputs);
             if slice {
                 // A duplicated slice will have the following two outer_fn arguments:
                 // (..., ptr1, int1, ptr2, int2, ...). We add the following llvm-ir to our __enzyme call:
@@ -209,6 +249,19 @@ fn match_args_from_caller_to_enzyme<'ll>(
             activity_pos += 1;
         }
     }
+    dbg!("ending");
+}
+
+fn is_slice(activity_pos: usize, inputs: &[DiffActivity]) -> bool {
+     if activity_pos + 1 >= inputs.len() {
+         // If there is no arg following our ptr, it also can't be a slice,
+         // since that would lead to a ptr, int pair.
+         false
+     } else {
+         let next_activity = inputs[activity_pos + 1];
+         // We analyze the MIR types and add this dummy activity if we visit a slice.
+         matches!(next_activity, DiffActivity::FakeActivitySize(_))
+     }
 }
 
 // On LLVM-IR, we can luckily declare __enzyme_ functions without specifying the input
@@ -426,6 +479,7 @@ fn generate_enzyme_call<'ll>(
 
         let call = builder.call(enzyme_ty, ad_fn, &args, None);
 
+        dbg!(&call);
         // This part is a bit iffy. LLVM requires that a call to an inlineable function has some
         // metadata attached to it, but we just created this code oota. Given that the
         // differentiated function already has partly confusing metadata, and given that this
@@ -472,6 +526,7 @@ fn generate_enzyme_call<'ll>(
         } else {
             builder.ret(call);
         }
+        dbg!("Still alive");
 
         // Let's crash in case that we messed something up above and generated invalid IR.
         llvm::LLVMRustVerifyFunction(
@@ -531,6 +586,7 @@ pub(crate) fn differentiate<'ll>(
 
         generate_enzyme_call(&cx, fn_def, fn_target, item.attrs.clone());
     }
+    dbg!("lowered all");
 
     // FIXME(ZuseZ4): support SanitizeHWAddress and prevent illegal/unsupported opts
 

compiler/rustc_codegen_llvm/src/type_.rs

@@ -154,6 +154,23 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
             )
         }
     }
+
+    pub(crate) fn is_float_type(&self, ty: &'ll Type) -> bool {
+        matches!(
+            self.type_kind(ty),
+            TypeKind::Half | TypeKind::Float | TypeKind::Double | TypeKind::X86_FP80
+                | TypeKind::FP128 | TypeKind::PPC_FP128
+        )
+    }
+
+    pub(crate) fn is_vec_or_array_type(&self, ty: &'ll Type) -> bool {
+        matches!(self.type_kind(ty),
+            TypeKind::Array | TypeKind::Vector | TypeKind::ScalableVector)
+    }
+
+    pub(crate) fn is_ptr_type(&self, ty: &'ll Type) -> bool {
+        matches!(self.type_kind(ty), TypeKind::Pointer)
+    }
 }
 
 impl<'ll, CX: Borrow<SCx<'ll>>> BaseTypeCodegenMethods for GenericCx<'ll, CX> {

compiler/rustc_monomorphize/src/partitioning/autodiff.rs

@@ -60,6 +60,9 @@ fn adjust_activity_to_abi<'tcx>(tcx: TyCtxt<'tcx>, fn_ty: Ty<'tcx>, da: &mut Vec
                         | DiffActivity::Duplicated => {
                             DiffActivity::FakeActivitySize(Some(elem_size))
                         }
+                        DiffActivity::Buffer => {
+                            DiffActivity::FakeActivitySize(Some(elem_size))
+                        }
                         DiffActivity::Const => DiffActivity::Const,
                         _ => bug!("unexpected activity for ptr/ref"),
                     };

tests/codegen/batching/batching2.rs

@@ -1,48 +1,62 @@
-//@ compile-flags: -Zautodiff=Enable -Zautodiff=NoPostopt -C opt-level=3  -Clto=fat
+//@ compile-flags: -Zautodiff=Enable,PrintModAfter -Zautodiff=NoPostopt -C opt-level=3  -Clto=fat
 //@ no-prefer-dynamic
 //@ needs-enzyme
 //
 #![feature(rustc_attrs)]
 #![feature(prelude_import)]
 #![feature(panic_internals)]
+#![feature(portable_simd)]
 #![no_std]
 //@ needs-enzyme
 #![feature(autodiff)]
 #[prelude_import]
 use ::std::prelude::rust_2015::*;
 #[macro_use]
 extern crate std;
+use std::simd::Simd;
 //@ pretty-mode:expanded
 //@ pretty-compare-only
 //@ pp-exact:batching.pp
 
+// CHECK: __enzyme
 
 // 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)]
+/// #[batching(d_square2, 4, Buffer, Buffer)]
 /// fn square(x: f32) -> f32 {
 ///    x * x
 /// }
 
 #[no_mangle]
 #[rustc_autodiff]
 #[inline(never)]
-fn square(x: f32) -> f32 {
-    x * x
+fn square(x: &[f32]) -> f32 {
+    x[0] * x[0]
 }
-#[rustc_autodiff(Batch, 4, Batching, Batching)]
+#[rustc_autodiff(Batch, 4, Buffer, Buffer)]
 #[no_mangle]
 #[inline(never)]
-fn d_square2(x: Simd<f32,4>) -> Simd<f32,4> {
+fn d_square2(x: &[f32]) -> Simd<f32,4> {
     unsafe {
         asm!("NOP", options(nomem));
     };
-    ::core::hint::black_box(x);
+    ::core::hint::black_box(square(x));
     ::core::hint::black_box(());
     ::core::hint::black_box(Default::default())
 }
+//#[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(square(x[0]));
+//    ::core::hint::black_box(());
+//    ::core::hint::black_box(Default::default())
+//}
 fn main() {}