add batch arg lowering to enzyme, update test

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -208,16 +208,21 @@ fn generate_enzyme_call<'ll>(
                     // FIXME(ZuseZ4): We will upstream a safety check later which asserts that
                     // int2 >= int1, which means the shadow vector is large enough to store the gradient.
                     assert!(llvm::LLVMRustGetTypeKind(next_outer_ty) == llvm::TypeKind::Integer);
-                    let next_outer_arg2 = outer_args[outer_pos + 2];
-                    let next_outer_ty2 = cx.val_ty(next_outer_arg2);
-                    assert!(llvm::LLVMRustGetTypeKind(next_outer_ty2) == llvm::TypeKind::Pointer);
-                    let next_outer_arg3 = outer_args[outer_pos + 3];
-                    let next_outer_ty3 = cx.val_ty(next_outer_arg3);
-                    assert!(llvm::LLVMRustGetTypeKind(next_outer_ty3) == llvm::TypeKind::Integer);
-                    args.push(next_outer_arg2);
+
+                    for _ in 0..attrs.width {
+                        let next_outer_arg2 = outer_args[outer_pos + 2];
+                        let next_outer_ty2 = cx.val_ty(next_outer_arg2);
+                        assert!(llvm::LLVMRustGetTypeKind(next_outer_ty2) == llvm::TypeKind::Pointer);
+                        let next_outer_arg3 = outer_args[outer_pos + 3];
+                        let next_outer_ty3 = cx.val_ty(next_outer_arg3);
+                        assert!(llvm::LLVMRustGetTypeKind(next_outer_ty3) == llvm::TypeKind::Integer);
+                        args.push(next_outer_arg2);
+                    }
+
+
                     args.push(cx.get_metadata_value(enzyme_const));
                     args.push(next_outer_arg);
-                    outer_pos += 4;
+                    outer_pos += 2 + 2 * attrs.width as usize;
                     activity_pos += 2;
                 } else {
                     // A duplicated pointer will have the following two outer_fn arguments:
@@ -235,6 +240,14 @@ fn generate_enzyme_call<'ll>(
                     args.push(next_outer_arg);
                     outer_pos += 2;
                     activity_pos += 1;
+
+                    // Now, if width > 1, we need to account for that
+                    for _ in 1..attrs.width {
+                        let next_outer_arg = outer_args[outer_pos];
+                        args.push(next_outer_arg);
+                        outer_pos += 1;
+                    }
+
                 }
             } else {
                 // We do not differentiate with resprect to this argument.

tests/codegen/autodiffv.rs

@@ -11,23 +11,40 @@ fn square(x: &f64) -> f64 {
     x * x
 }
 
-// CHECK:define internal fastcc double @diffesquare(double %x.0.val, ptr nocapture align 8 %"x'"
+// CHECK:define internal fastcc void @diffe4square([4 x ptr] %"x'"
 // CHECK-NEXT:invertstart:
-// CHECK-NEXT:  %_0 = fmul double %x.0.val, %x.0.val
-// CHECK-NEXT:  %0 = fadd fast double %x.0.val, %x.0.val
-// CHECK-NEXT:  %1 = load double, ptr %"x'", align 8
-// CHECK-NEXT:  %2 = fadd fast double %1, %0
-// CHECK-NEXT:  store double %2, ptr %"x'", align 8
-// CHECK-NEXT:  ret double %_0
+// CHECK-NEXT:  %0 = extractvalue [4 x ptr] %"x'", 0
+// CHECK-NEXT:  %1 = load double, ptr %0, align 8, !alias.scope !15950, !noalias !15953
+// CHECK-NEXT:  %2 = fadd fast double %1, 6.000000e+00
+// CHECK-NEXT:  store double %2, ptr %0, align 8, !alias.scope !15950, !noalias !15953
+// CHECK-NEXT:  %3 = extractvalue [4 x ptr] %"x'", 1
+// CHECK-NEXT:  %4 = load double, ptr %3, align 8, !alias.scope !15958, !noalias !15959
+// CHECK-NEXT:  %5 = fadd fast double %4, 6.000000e+00
+// CHECK-NEXT:  store double %5, ptr %3, align 8, !alias.scope !15958, !noalias !15959
+// CHECK-NEXT:  %6 = extractvalue [4 x ptr] %"x'", 2
+// CHECK-NEXT:  %7 = load double, ptr %6, align 8, !alias.scope !15960, !noalias !15961
+// CHECK-NEXT:  %8 = fadd fast double %7, 6.000000e+00
+// CHECK-NEXT:  store double %8, ptr %6, align 8, !alias.scope !15960, !noalias !15961
+// CHECK-NEXT:  %9 = extractvalue [4 x ptr] %"x'", 3
+// CHECK-NEXT:  %10 = load double, ptr %9, align 8, !alias.scope !15962, !noalias !15963
+// CHECK-NEXT:  %11 = fadd fast double %10, 6.000000e+00
+// CHECK-NEXT:  store double %11, ptr %9, align 8, !alias.scope !15962, !noalias !15963
+// CHECK-NEXT:  ret void
 // CHECK-NEXT:}
 
 fn main() {
     let x = 3.0;
     let output = square(&x);
     assert_eq!(9.0, output);
 
-    let mut df_dx = 0.0;
-    let output_ = d_square(&x, &mut df_dx, 1.0);
+    let mut df_dx1 = 0.0;
+    let mut df_dx2 = 0.0;
+    let mut df_dx3 = 0.0;
+    let mut df_dx4 = 0.0;
+    let output_ = d_square(&x, &mut df_dx1, &mut df_dx2, &mut df_dx3,  &mut df_dx4, 1.0);
     assert_eq!(output, output_);
-    assert_eq!(6.0, df_dx);
+    assert_eq!(6.0, df_dx1);
+    assert_eq!(6.0, df_dx2);
+    assert_eq!(6.0, df_dx3);
+    assert_eq!(6.0, df_dx4);
 }