From 530b56523b91947476ceee19862fa65191548a46 Mon Sep 17 00:00:00 2001
From: bit-aloo <sshourya17@gmail.com>
Date: Sat, 12 Apr 2025 17:23:19 +0530
Subject: [PATCH] Add explicit inlining to methods after AD layer

---
 compiler/rustc_codegen_llvm/src/back/lto.rs | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/compiler/rustc_codegen_llvm/src/back/lto.rs b/compiler/rustc_codegen_llvm/src/back/lto.rs
index a8b49e9552c..3bab72e4190 100644
--- a/compiler/rustc_codegen_llvm/src/back/lto.rs
+++ b/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,19 @@ 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") {
+                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 {