autodiff batching mvp

Author: ZuseZ4

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -77,6 +77,7 @@ 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,
+    pub width: u32,
     pub ret_activity: DiffActivity,
     pub input_activity: Vec<DiffActivity>,
 }
@@ -222,13 +223,15 @@ impl AutoDiffAttrs {
     pub const fn error() -> Self {
         AutoDiffAttrs {
             mode: DiffMode::Error,
+            width: 0,
             ret_activity: DiffActivity::None,
             input_activity: Vec::new(),
         }
     }
     pub fn source() -> Self {
         AutoDiffAttrs {
             mode: DiffMode::Source,
+            width: 0,
             ret_activity: DiffActivity::None,
             input_activity: Vec::new(),
         }

compiler/rustc_builtin_macros/messages.ftl

@@ -76,6 +76,7 @@ builtin_macros_autodiff_mode_activity = {$act} can not be used in {$mode} Mode
 builtin_macros_autodiff_not_build = this rustc version does not support autodiff
 builtin_macros_autodiff_number_activities = expected {$expected} activities, but found {$found}
 builtin_macros_autodiff_ty_activity = {$act} can not be used for this type
+builtin_macros_autodiff_width = autodiff width must fit u32, but is {$width}
 
 builtin_macros_autodiff_unknown_activity = did not recognize Activity: `{$act}`
 builtin_macros_bad_derive_target = `derive` may only be applied to `struct`s, `enum`s and `union`s

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -193,6 +193,14 @@ mod autodiff {
         pub(crate) mode: String,
     }
 
+    #[derive(Diagnostic)]
+    #[diag(builtin_macros_autodiff_width)]
+    pub(crate) struct AutoDiffInvalidWidth {
+        #[primary_span]
+        pub(crate) span: Span,
+        pub(crate) width: u128,
+    }
+
     #[derive(Diagnostic)]
     #[diag(builtin_macros_autodiff)]
     pub(crate) struct AutoDiffInvalidApplication {

compiler/rustc_builtin_macros/src/errors.rs

@@ -1,3 +1,4 @@
+codegen_llvm_autodiff_unused_args = implementation bug, failed to match all args on llvm level
 codegen_llvm_autodiff_without_enable = using the autodiff feature requires -Z autodiff=Enable
 codegen_llvm_autodiff_without_lto = using the autodiff feature requires using fat-lto
 

compiler/rustc_codegen_llvm/messages.ftl

@@ -655,6 +655,10 @@ pub(crate) fn run_pass_manager(
     unsafe {
         write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
     }
+    // This is the final IR, so people should be able to inspect the optimized autodiff output.
+    if config.autodiff.contains(&config::AutoDiff::PrintModAfter) {
+        unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
+    }
 
     if cfg!(llvm_enzyme) && enable_ad {
         let opt_stage = llvm::OptStage::FatLTO;

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -405,7 +405,7 @@ impl<'ll> CodegenCx<'ll, '_> {
             let val_llty = self.val_ty(v);
 
             let g = self.get_static_inner(def_id, val_llty);
-            let llty = llvm::LLVMGlobalGetValueType(g);
+            let llty = self.get_type_of_global(g);
 
             let g = if val_llty == llty {
                 g

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -37,7 +37,7 @@ use crate::debuginfo::metadata::apply_vcall_visibility_metadata;
 use crate::llvm::Metadata;
 use crate::type_::Type;
 use crate::value::Value;
-use crate::{attributes, coverageinfo, debuginfo, llvm, llvm_util};
+use crate::{attributes, common, coverageinfo, debuginfo, llvm, llvm_util};
 
 /// `TyCtxt` (and related cache datastructures) can't be move between threads.
 /// However, there are various cx related functions which we want to be available to the builder and
@@ -642,7 +642,18 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
         llvm::set_section(g, c"llvm.metadata");
     }
 }
-
+impl<'ll> SimpleCx<'ll> {
+    pub(crate) fn get_return_type(&self, ty: &'ll Type) -> &'ll Type {
+        assert!(unsafe { llvm::LLVMRustIsFunctionTy(ty) });
+        unsafe { llvm::LLVMGetReturnType(ty) }
+    }
+    pub(crate) fn get_type_of_global(&self, val: &'ll Value) -> &'ll Type {
+        unsafe { llvm::LLVMGlobalGetValueType(val) }
+    }
+    pub(crate) fn val_ty(&self, v: &'ll Value) -> &'ll Type {
+        common::val_ty(v)
+    }
+}
 impl<'ll> SimpleCx<'ll> {
     pub(crate) fn new(
         llmod: &'ll llvm::Module,
@@ -659,6 +670,11 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         llvm::LLVMMetadataAsValue(self.llcx(), metadata)
     }
 
+    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) }
+    }
+
     pub(crate) fn get_function(&self, name: &str) -> Option<&'ll Value> {
         let name = SmallCStr::new(name);
         unsafe { llvm::LLVMGetNamedFunction((**self).borrow().llmod, name.as_ptr()) }

compiler/rustc_codegen_llvm/src/consts.rs

@@ -94,6 +94,10 @@ impl<G: EmissionGuarantee> Diagnostic<'_, G> for ParseTargetMachineConfig<'_> {
 #[diag(codegen_llvm_autodiff_without_lto)]
 pub(crate) struct AutoDiffWithoutLTO;
 
+#[derive(Diagnostic)]
+#[diag(codegen_llvm_autodiff_unused_args)]
+pub(crate) struct AutoDiffUnusedArgs;
+
 #[derive(Diagnostic)]
 #[diag(codegen_llvm_autodiff_without_enable)]
 pub(crate) struct AutoDiffWithoutEnable;

compiler/rustc_codegen_llvm/src/context.rs

@@ -4,7 +4,7 @@
 use libc::{c_char, c_uint};
 
 use super::MetadataKindId;
-use super::ffi::{BasicBlock, Metadata, Module, Type, Value};
+use super::ffi::{AttributeKind, BasicBlock, Metadata, Module, Type, Value};
 use crate::llvm::Bool;
 
 #[link(name = "llvm-wrapper", kind = "static")]
@@ -17,6 +17,11 @@ unsafe extern "C" {
     pub(crate) fn LLVMRustEraseInstFromParent(V: &Value);
     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;
 }
 
 unsafe extern "C" {

compiler/rustc_codegen_llvm/src/errors.rs

@@ -4,7 +4,7 @@ use rustc_abi::ExternAbi;
 use rustc_ast::expand::autodiff_attrs::{
     AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ret_activity,
 };
-use rustc_ast::{MetaItem, MetaItemInner, attr};
+use rustc_ast::{LitKind, MetaItem, MetaItemInner, attr};
 use rustc_attr_parsing::ReprAttr::ReprAlign;
 use rustc_attr_parsing::{AttributeKind, InlineAttr, InstructionSetAttr, OptimizeAttr};
 use rustc_data_structures::fx::FxHashMap;
@@ -819,8 +819,8 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
         return Some(AutoDiffAttrs::source());
     }
 
-    let [mode, input_activities @ .., ret_activity] = &list[..] else {
-        span_bug!(attr.span(), "rustc_autodiff attribute must contain mode and activities");
+    let [mode, width_meta, input_activities @ .., ret_activity] = &list[..] else {
+        span_bug!(attr.span(), "rustc_autodiff attribute must contain mode, width and activities");
     };
     let mode = if let MetaItemInner::MetaItem(MetaItem { path: p1, .. }) = mode {
         p1.segments.first().unwrap().ident
@@ -837,6 +837,34 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
         }
     };
 
+    let width: u32;
+    match width_meta {
+        MetaItemInner::MetaItem(MetaItem { path: p1, .. }) => {
+            let w = p1.segments.first().unwrap().ident;
+            width = match w.as_str().parse() {
+                Ok(val) => val,
+                Err(_) => {
+                    span_bug!(w.span, "rustc_autodiff width should fit u32");
+                }
+            };
+        }
+        MetaItemInner::Lit(lit) => {
+            if let LitKind::Int(val, _) = lit.kind {
+                width = match val.get().try_into() {
+                    Ok(val) => val,
+                    Err(_) => {
+                        span_bug!(lit.span, "rustc_autodiff width should fit u32");
+                    }
+                };
+            } else {
+                span_bug!(lit.span, "rustc_autodiff width should be an integer");
+            }
+        }
+        _ => {
+            span_bug!(width_meta.span(), "failed to parse rustc_autodiff width");
+        }
+    }
+
     // First read the ret symbol from the attribute
     let ret_symbol = if let MetaItemInner::MetaItem(MetaItem { path: p1, .. }) = ret_activity {
         p1.segments.first().unwrap().ident
@@ -883,7 +911,7 @@ fn autodiff_attrs(tcx: TyCtxt<'_>, id: DefId) -> Option<AutoDiffAttrs> {
         span_bug!(attr.span(), "Invalid return activity {} for {} mode", ret_activity, mode);
     }
 
-    Some(AutoDiffAttrs { mode, ret_activity, input_activity: arg_activities })
+    Some(AutoDiffAttrs { mode, width, ret_activity, input_activity: arg_activities })
 }
 
 pub(crate) fn provide(providers: &mut Providers) {

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -384,6 +384,12 @@ static inline void AddAttributes(T *t, unsigned Index, LLVMAttributeRef *Attrs,
   t->setAttributes(PALNew);
 }
 
+extern "C" bool LLVMRustHasAttributeAtIndex(LLVMValueRef Fn, unsigned Index,
+                                       LLVMRustAttributeKind RustAttr) {
+  Function *F = unwrap<Function>(Fn);
+  return F->hasParamAttribute(Index, fromRust(RustAttr));
+}
+
 extern "C" void LLVMRustAddFunctionAttributes(LLVMValueRef Fn, unsigned Index,
                                               LLVMAttributeRef *Attrs,
                                               size_t AttrsLen) {
@@ -635,6 +641,17 @@ 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();
+}
 
 extern "C" LLVMValueRef
 LLVMRustInlineAsm(LLVMTypeRef Ty, char *AsmString, size_t AsmStringLen,

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -0,0 +1,53 @@
+//@ compile-flags: -Zautodiff=Enable -C opt-level=3  -Clto=fat
+//@ no-prefer-dynamic
+//@ needs-enzyme
+#![feature(autodiff)]
+
+use std::autodiff::autodiff;
+
+#[autodiff(d_square, Reverse, 4, Duplicated, Active)]
+#[no_mangle]
+fn square(x: &f64) -> f64 {
+    x * x
+}
+
+// CHECK:define internal fastcc void @diffe4square([4 x ptr] %"x'"
+// CHECK-NEXT:invertstart:
+// 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_dx1 = 0.0;
+    let mut df_dx2 = 0.0;
+    let mut df_dx3 = 0.0;
+    let mut df_dx4 = 0.0;
+    let [o1, o2, o3, o4] = d_square(&x, &mut df_dx1, &mut df_dx2, &mut df_dx3, &mut df_dx4, 1.0);
+    assert_eq!(output, o1);
+    assert_eq!(output, o2);
+    assert_eq!(output, o3);
+    assert_eq!(output, o4);
+    assert_eq!(6.0, df_dx1);
+    assert_eq!(6.0, df_dx2);
+    assert_eq!(6.0, df_dx3);
+    assert_eq!(6.0, df_dx4);
+}