addressing most feedback

Author: ZuseZ4

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -77,6 +77,16 @@ pub struct AutoDiffAttrs {
     /// e.g. in the [JAX
     /// Documentation](https://jax.readthedocs.io/en/latest/_tutorials/advanced-autodiff.html#how-it-s-made-two-foundational-autodiff-functions).
     pub mode: DiffMode,
+    /// A user-provided, batching width. If not given, we will default to 1 (no batching).
+    /// Calling a differentiated, non-batched function through a loop 100 times is equivalent to:
+    /// - Calling the function 50 times with a batch size of 2
+    /// - Calling the function 25 times with a batch size of 4,
+    /// etc. A batched function takes more (or longer) arguments, and might be able to benefit from
+    /// cache locality, better re-usal of primal values, and other optimizations.
+    /// We will (before LLVM's vectorizer runs) just generate most LLVM-IR instructions `width`
+    /// times, so this massively increases code size. As such, values like 1024 are unlikely to
+    /// work. We should consider limiting this to u8 or u16, but will leave it at u32 for
+    /// experiments for now and focus on documenting the implications of a large width.
     pub width: u32,
     pub ret_activity: DiffActivity,
     pub input_activity: Vec<DiffActivity>,

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -3,6 +3,8 @@ use std::ptr;
 use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivity, DiffMode};
 use rustc_codegen_ssa::ModuleCodegen;
 use rustc_codegen_ssa::back::write::ModuleConfig;
+use rustc_codegen_ssa::common::TypeKind;
+use rustc_codegen_ssa::traits::BaseTypeCodegenMethods;
 use rustc_errors::FatalError;
 use rustc_middle::bug;
 use tracing::{debug, trace};
@@ -18,18 +20,18 @@ use crate::value::Value;
 use crate::{CodegenContext, LlvmCodegenBackend, ModuleLlvm, attributes, llvm};
 
 fn get_params(fnc: &Value) -> Vec<&Value> {
+    let param_num = llvm::LLVMCountParams(fnc) as usize;
+    let mut fnc_args: Vec<&Value> = vec![];
+    fnc_args.reserve(param_num);
     unsafe {
-        let param_num = llvm::LLVMCountParams(fnc) as usize;
-        let mut fnc_args: Vec<&Value> = vec![];
-        fnc_args.reserve(param_num);
         llvm::LLVMGetParams(fnc, fnc_args.as_mut_ptr());
         fnc_args.set_len(param_num);
-        fnc_args
     }
+    fnc_args
 }
 
 fn has_sret(fnc: &Value) -> bool {
-    let num_args = unsafe { llvm::LLVMCountParams(fnc) as usize };
+    let num_args = llvm::LLVMCountParams(fnc) as usize;
     if num_args == 0 {
         false
     } else {
@@ -121,23 +123,15 @@ fn match_args_from_caller_to_enzyme<'ll>(
                 // (..., metadata! enzyme_dup, ptr, ptr, int1, ...).
                 // 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!(unsafe {
-                    llvm::LLVMRustGetTypeKind(next_outer_ty) == llvm::TypeKind::Integer
-                });
+                assert_eq!(cx.type_kind(next_outer_ty), TypeKind::Integer);
 
                 for _ in 0..width {
                     let next_outer_arg2 = outer_args[outer_pos + 2];
                     let next_outer_ty2 = cx.val_ty(next_outer_arg2);
-                    assert!(
-                        unsafe { llvm::LLVMRustGetTypeKind(next_outer_ty2) }
-                            == llvm::TypeKind::Pointer
-                    );
+                    assert_eq!(cx.type_kind(next_outer_ty2), TypeKind::Pointer);
                     let next_outer_arg3 = outer_args[outer_pos + 3];
                     let next_outer_ty3 = cx.val_ty(next_outer_arg3);
-                    assert!(
-                        unsafe { llvm::LLVMRustGetTypeKind(next_outer_ty3) }
-                            == llvm::TypeKind::Integer
-                    );
+                    assert_eq!(cx.type_kind(next_outer_ty3), TypeKind::Integer);
                     args.push(next_outer_arg2);
                 }
                 args.push(cx.get_metadata_value(enzyme_const));
@@ -150,10 +144,7 @@ fn match_args_from_caller_to_enzyme<'ll>(
                 // (..., metadata! enzyme_dup, ptr, ptr, ...).
                 if matches!(diff_activity, DiffActivity::Duplicated | DiffActivity::DuplicatedOnly)
                 {
-                    assert!(
-                        unsafe { llvm::LLVMRustGetTypeKind(next_outer_ty) }
-                            == llvm::TypeKind::Pointer
-                    );
+                    assert_eq!(cx.type_kind(next_outer_ty), TypeKind::Pointer);
                 }
                 // In the case of Dual we don't have assumptions, e.g. f32 would be valid.
                 args.push(next_outer_arg);
@@ -213,8 +204,8 @@ fn compute_enzyme_fn_ty<'ll>(
             todo!("Handle sret for scalar ad");
         } else {
             // First we check if we also have to deal with the primal return.
-            if attrs.mode.is_fwd() {
-                match attrs.ret_activity {
+            match attrs.mode {
+                DiffMode::Forward => match attrs.ret_activity {
                     DiffActivity::Dual => {
                         let arr_ty =
                             unsafe { llvm::LLVMArrayType2(inner_ret_ty, attrs.width as u64 + 1) };
@@ -231,11 +222,13 @@ fn compute_enzyme_fn_ty<'ll>(
                     _ => {
                         bug!("unreachable");
                     }
+                },
+                DiffMode::Reverse => {
+                    todo!("Handle sret for reverse mode");
+                }
+                _ => {
+                    bug!("unreachable");
                 }
-            } else if attrs.mode.is_rev() {
-                todo!("Handle sret for reverse mode");
-            } else {
-                bug!("unreachable");
             }
         }
     }
@@ -395,7 +388,7 @@ fn generate_enzyme_call<'ll>(
                 // now store the result of the enzyme call into the sret pointer.
                 let sret_ptr = outer_args[0];
                 let call_ty = cx.val_ty(call);
-                assert!(llvm::LLVMRustIsArrayTy(call_ty));
+                assert_eq!(cx.type_kind(call_ty), TypeKind::Array);
                 llvm::LLVMBuildStore(&builder.llbuilder, call, sret_ptr);
             }
             builder.ret_void();

compiler/rustc_codegen_llvm/src/context.rs

@@ -8,6 +8,7 @@ use std::str;
 use rustc_abi::{HasDataLayout, Size, TargetDataLayout, VariantIdx};
 use rustc_codegen_ssa::back::versioned_llvm_target;
 use rustc_codegen_ssa::base::{wants_msvc_seh, wants_wasm_eh};
+use rustc_codegen_ssa::common::TypeKind;
 use rustc_codegen_ssa::errors as ssa_errors;
 use rustc_codegen_ssa::traits::*;
 use rustc_data_structures::base_n::{ALPHANUMERIC_ONLY, ToBaseN};
@@ -645,7 +646,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
 }
 impl<'ll> SimpleCx<'ll> {
     pub(crate) fn get_return_type(&self, ty: &'ll Type) -> &'ll Type {
-        assert!(unsafe { llvm::LLVMRustIsFunctionTy(ty) });
+        assert_eq!(self.type_kind(ty), TypeKind::Function);
         unsafe { llvm::LLVMGetReturnType(ty) }
     }
     pub(crate) fn get_type_of_global(&self, val: &'ll Value) -> &'ll Type {
@@ -671,6 +672,8 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         llvm::LLVMMetadataAsValue(self.llcx(), metadata)
     }
 
+    // FIXME(autodiff): We should split `ConstCodegenMethods` to pull the reusable parts
+    // onto a trait that is also implemented for GenericCx.
     pub(crate) fn get_const_i64(&self, n: u64) -> &'ll Value {
         let ty = unsafe { llvm::LLVMInt64TypeInContext(self.llcx()) };
         unsafe { llvm::LLVMConstInt(ty, n, llvm::False) }

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -18,9 +18,6 @@ unsafe extern "C" {
     pub(crate) fn LLVMRustGetTerminator<'a>(B: &BasicBlock) -> &'a Value;
     pub(crate) fn LLVMRustVerifyFunction(V: &Value, action: LLVMRustVerifierFailureAction) -> Bool;
     pub(crate) fn LLVMRustHasAttributeAtIndex(V: &Value, i: c_uint, Kind: AttributeKind) -> bool;
-
-    pub(crate) fn LLVMRustIsFunctionTy(Ty: &Type) -> bool;
-    pub(crate) fn LLVMRustIsArrayTy(Ty: &Type) -> bool;
     pub(crate) fn LLVMRustGetArrayNumElements(Ty: &Type) -> u64;
 }
 

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1172,7 +1172,7 @@ unsafe extern "C" {
 
     // Operations on parameters
     pub(crate) fn LLVMIsAArgument(Val: &Value) -> Option<&Value>;
-    pub(crate) fn LLVMCountParams(Fn: &Value) -> c_uint;
+    pub(crate) safe fn LLVMCountParams(Fn: &Value) -> c_uint;
     pub(crate) fn LLVMGetParam(Fn: &Value, Index: c_uint) -> &Value;
 
     // Operations on basic blocks

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -831,7 +831,7 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
                 Err(_) => {
                     span_bug!(w.span, "rustc_autodiff width should fit u32");
                 }
-            };
+            }
         }
         MetaItemInner::Lit(lit) => {
             if let LitKind::Int(val, _) = lit.kind {
@@ -840,12 +840,12 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
                     Err(_) => {
                         span_bug!(lit.span, "rustc_autodiff width should fit u32");
                     }
-                };
+                }
             } else {
                 span_bug!(lit.span, "rustc_autodiff width should be an integer");
             }
         }
-    }
+    };
 
     // First read the ret symbol from the attribute
     let ret_symbol = if let MetaItemInner::MetaItem(MetaItem { path: p1, .. }) = ret_activity {

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -641,13 +641,6 @@ static InlineAsm::AsmDialect fromRust(LLVMRustAsmDialect Dialect) {
     report_fatal_error("bad AsmDialect.");
   }
 }
-extern "C" bool LLVMRustIsFunctionTy(LLVMTypeRef Ty) {
-  return unwrap(Ty)->isFunctionTy();
-}
-
-extern "C" bool LLVMRustIsArrayTy(LLVMTypeRef Ty) {
-  return unwrap(Ty)->isArrayTy();
-}
 
 extern "C" uint64_t LLVMRustGetArrayNumElements(LLVMTypeRef Ty) {
   return unwrap(Ty)->getArrayNumElements();

tests/codegen/autodiffv.rs

@@ -18,13 +18,13 @@ fn square(x: &f32) -> f32 {
 // CHECK: define internal fastcc [4 x float] @fwddiffe4square(float %x.0.val, [4 x ptr] %"x'")
 // CHECK-NEXT: start:
 // CHECK-NEXT:   %0 = extractvalue [4 x ptr] %"x'", 0
-// CHECK-NEXT:   %"_2'ipl" = load float, ptr %0, align 4, !alias.scope !38, !noalias !39
+// CHECK-NEXT:   %"_2'ipl" = load float, ptr %0, align 4
 // CHECK-NEXT:   %1 = extractvalue [4 x ptr] %"x'", 1
-// CHECK-NEXT:   %"_2'ipl1" = load float, ptr %1, align 4, !alias.scope !40, !noalias !41
+// CHECK-NEXT:   %"_2'ipl1" = load float, ptr %1, align 4
 // CHECK-NEXT:   %2 = extractvalue [4 x ptr] %"x'", 2
-// CHECK-NEXT:   %"_2'ipl2" = load float, ptr %2, align 4, !alias.scope !42, !noalias !43
+// CHECK-NEXT:   %"_2'ipl2" = load float, ptr %2, align 4
 // CHECK-NEXT:   %3 = extractvalue [4 x ptr] %"x'", 3
-// CHECK-NEXT:   %"_2'ipl3" = load float, ptr %3, align 4, !alias.scope !44, !noalias !45
+// CHECK-NEXT:   %"_2'ipl3" = load float, ptr %3, align 4
 // CHECK-NEXT:   %4 = insertelement <4 x float> poison, float %"_2'ipl", i64 0
 // CHECK-NEXT:   %5 = insertelement <4 x float> %4, float %"_2'ipl1", i64 1
 // CHECK-NEXT:   %6 = insertelement <4 x float> %5, float %"_2'ipl2", i64 2
@@ -48,13 +48,13 @@ fn square(x: &f32) -> f32 {
 // CHECK: define internal fastcc { float, [4 x float] } @fwddiffe4square.1(float %x.0.val, [4 x ptr] %"x'")
 // CHECK-NEXT: start:
 // CHECK-NEXT:   %0 = extractvalue [4 x ptr] %"x'", 0
-// CHECK-NEXT:   %"_2'ipl" = load float, ptr %0, align 4, !alias.scope !46, !noalias !47
+// CHECK-NEXT:   %"_2'ipl" = load float, ptr %0, align 4
 // CHECK-NEXT:   %1 = extractvalue [4 x ptr] %"x'", 1
-// CHECK-NEXT:   %"_2'ipl1" = load float, ptr %1, align 4, !alias.scope !48, !noalias !49
+// CHECK-NEXT:   %"_2'ipl1" = load float, ptr %1, align 4
 // CHECK-NEXT:   %2 = extractvalue [4 x ptr] %"x'", 2
-// CHECK-NEXT:   %"_2'ipl2" = load float, ptr %2, align 4, !alias.scope !50, !noalias !51
+// CHECK-NEXT:   %"_2'ipl2" = load float, ptr %2, align 4
 // CHECK-NEXT:   %3 = extractvalue [4 x ptr] %"x'", 3
-// CHECK-NEXT:   %"_2'ipl3" = load float, ptr %3, align 4, !alias.scope !52, !noalias !53
+// CHECK-NEXT:   %"_2'ipl3" = load float, ptr %3, align 4
 // CHECK-NEXT:   %_0 = fmul float %x.0.val, %x.0.val
 // CHECK-NEXT:   %4 = insertelement <4 x float> poison, float %"_2'ipl", i64 0
 // CHECK-NEXT:   %5 = insertelement <4 x float> %4, float %"_2'ipl1", i64 1