add autodiff inline

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -28,6 +28,14 @@ pub(crate) fn apply_to_callsite(callsite: &Value, idx: AttributePlace, attrs: &[
     }
 }
 
+pub(crate) fn has_attr(llfn: &Value, idx: u32, attr: AttributeKind) -> bool {
+    llvm::HasAttributeAtIndex(llfn, idx, attr)
+}
+
+pub(crate) fn remove_from_llfn(llfn: &Value, place: AttributePlace, kind: AttributeKind) {
+    llvm::RemoveEnumAttributeAtIndex(llfn, place, kind);
+}
+
 /// Get LLVM attribute for the provided inline heuristic.
 #[inline]
 fn inline_attr<'ll>(cx: &CodegenCx<'ll, '_>, inline: InlineAttr) -> Option<&'ll Attribute> {

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -28,8 +28,9 @@ use crate::back::write::{
 use crate::errors::{
     DynamicLinkingWithLTO, LlvmError, LtoBitcodeFromRlib, LtoDisallowed, LtoDylib, LtoProcMacro,
 };
-use crate::llvm::{self, build_string};
-use crate::{LlvmCodegenBackend, ModuleLlvm};
+use crate::llvm::AttributePlace::Function;
+use crate::llvm::{self, build_string, get_value_name};
+use crate::{LlvmCodegenBackend, ModuleLlvm, SimpleCx, attributes};
 
 /// We keep track of the computed LTO cache keys from the previous
 /// session to determine which CGUs we can reuse.
@@ -664,6 +665,30 @@ pub(crate) fn run_pass_manager(
             unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
         }
 
+        let cx =
+            SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);
+
+        for function in cx.get_functions() {
+            let name = get_value_name(function);
+            let name = std::str::from_utf8(name).unwrap();
+
+            if name.starts_with("__enzyme") {
+                // Ensure `noinline` is present before replacing it.
+                // This is not strictly necessary for correctness, but serves as a sanity check
+                // in case the autodiff pass stops injecting `noinline` in the future.
+                assert!(
+                    !attributes::has_attr(function, 0, llvm::AttributeKind::NoInline),
+                    "Expected __enzyme function to have 'noinline' before adding 'alwaysinline'"
+                );
+
+                // Removing no-inline from function.
+                attributes::remove_from_llfn(function, Function, llvm::AttributeKind::NoInline);
+
+                let attr = llvm::AttributeKind::AlwaysInline.create_attr(cx.llcx);
+                attributes::apply_to_llfn(function, Function, &[attr]);
+            }
+        }
+
         let opt_stage = llvm::OptStage::FatLTO;
         let stage = write::AutodiffStage::PostAD;
         unsafe {

compiler/rustc_codegen_llvm/src/context.rs

@@ -698,6 +698,16 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
             llvm::LLVMMDStringInContext2(self.llcx(), name.as_ptr() as *const c_char, name.len())
         })
     }
+
+    pub(crate) fn get_functions(&self) -> Vec<&'ll Value> {
+        let mut functions = vec![];
+        let mut func = unsafe { llvm::LLVMGetFirstFunction(self.llmod()) };
+        while let Some(f) = func {
+            functions.push(f);
+            func = unsafe { llvm::LLVMGetNextFunction(f) }
+        }
+        functions
+    }
 }
 
 impl<'ll, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1157,6 +1157,11 @@ unsafe extern "C" {
     pub(crate) fn LLVMGetNamedGlobal(M: &Module, Name: *const c_char) -> Option<&Value>;
     pub(crate) fn LLVMGetFirstGlobal(M: &Module) -> Option<&Value>;
     pub(crate) fn LLVMGetNextGlobal(GlobalVar: &Value) -> Option<&Value>;
+
+    pub(crate) fn LLVMGetFirstFunction(M: &Module) -> Option<&Value>;
+    pub(crate) fn LLVMGetNextFunction(Fn: &Value) -> Option<&Value>;
+    pub(crate) fn LLVMRemoveEnumAttributeAtIndex(Fn: &Value, index: c_uint, kind: AttributeKind);
+
     pub(crate) fn LLVMDeleteGlobal(GlobalVar: &Value);
     pub(crate) fn LLVMGetInitializer(GlobalVar: &Value) -> Option<&Value>;
     pub(crate) fn LLVMSetInitializer<'a>(GlobalVar: &'a Value, ConstantVal: &'a Value);

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -41,6 +41,16 @@ pub(crate) fn AddFunctionAttributes<'ll>(
     }
 }
 
+pub(crate) fn HasAttributeAtIndex<'ll>(llfn: &'ll Value, idx: u32, kind: AttributeKind) -> bool {
+    unsafe { LLVMRustHasAttributeAtIndex(llfn, idx, kind) }
+}
+
+pub(crate) fn RemoveEnumAttributeAtIndex(llfn: &Value, place: AttributePlace, kind: AttributeKind) {
+    unsafe {
+        LLVMRemoveEnumAttributeAtIndex(llfn, place.as_uint(), kind);
+    }
+}
+
 pub(crate) fn AddCallSiteAttributes<'ll>(
     callsite: &'ll Value,
     idx: AttributePlace,

compiler/rustc_codegen_llvm/src/type_.rs

@@ -128,6 +128,10 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         (**self).borrow().llcx
     }
 
+    pub(crate) fn llmod(&self) -> &'ll llvm::Module {
+        (**self).borrow().llmod
+    }
+
     pub(crate) fn isize_ty(&self) -> &'ll Type {
         (**self).borrow().isize_ty
     }

tests/codegen/autodiff/inline.rs

@@ -0,0 +1,20 @@
+//@ compile-flags: -Zautodiff=Enable -Zautodiff=NoPostopt -C opt-level=3  -Clto=fat
+//@ no-prefer-dynamic
+//@ needs-enzyme
+
+#![feature(autodiff)]
+
+use std::autodiff::autodiff;
+
+#[autodiff(d_square, Reverse, Duplicated, Active)]
+fn square(x: &f64) -> f64 {
+    x * x
+}
+// CHECK: ; Function Attrs: alwaysinline noinline
+// CHECK-NEXT: declare double @__enzyme_autodiff_ZN6inline8d_square17h021c74e92c259cdeE(...) local_unnamed_addr #8
+fn main() {
+    let x = std::hint::black_box(3.0);
+    let mut dx1 = std::hint::black_box(1.0);
+    let _ = d_square(&x, &mut dx1, 1.0);
+    assert_eq!(dx1, 6.0);
+}