@@ -28,8 +28,9 @@ use crate::back::write::{
28
28
use crate :: errors:: {
29
29
DynamicLinkingWithLTO , LlvmError , LtoBitcodeFromRlib , LtoDisallowed , LtoDylib , LtoProcMacro ,
30
30
} ;
31
+ use crate :: llvm:: AttributePlace :: Function ;
31
32
use crate :: llvm:: { self , build_string} ;
32
- use crate :: { LlvmCodegenBackend , ModuleLlvm } ;
33
+ use crate :: { LlvmCodegenBackend , ModuleLlvm , SimpleCx , attributes } ;
33
34
34
35
/// We keep track of the computed LTO cache keys from the previous
35
36
/// session to determine which CGUs we can reuse.
@@ -666,6 +667,31 @@ pub(crate) fn run_pass_manager(
666
667
}
667
668
668
669
if cfg ! ( llvm_enzyme) && enable_ad && !thin {
670
+ let cx =
671
+ SimpleCx :: new ( module. module_llvm . llmod ( ) , & module. module_llvm . llcx , cgcx. pointer_size ) ;
672
+
673
+ for function in cx. get_functions ( ) {
674
+ let enzyme_marker = "enzyme_marker" ;
675
+ if attributes:: has_string_attr ( function, enzyme_marker) {
676
+ // Sanity check: Ensure 'noinline' is present before replacing it.
677
+ assert ! (
678
+ !attributes:: has_attr( function, Function , llvm:: AttributeKind :: NoInline ) ,
679
+ "Expected __enzyme function to have 'noinline' before adding 'alwaysinline'"
680
+ ) ;
681
+
682
+ attributes:: remove_from_llfn ( function, Function , llvm:: AttributeKind :: NoInline ) ;
683
+ attributes:: remove_string_attr_from_llfn ( function, enzyme_marker) ;
684
+
685
+ assert ! (
686
+ !attributes:: has_string_attr( function, enzyme_marker) ,
687
+ "Expected function to not have 'enzyme_marker'"
688
+ ) ;
689
+
690
+ let always_inline = llvm:: AttributeKind :: AlwaysInline . create_attr ( cx. llcx ) ;
691
+ attributes:: apply_to_llfn ( function, Function , & [ always_inline] ) ;
692
+ }
693
+ }
694
+
669
695
let opt_stage = llvm:: OptStage :: FatLTO ;
670
696
let stage = write:: AutodiffStage :: PostAD ;
671
697
if !config. autodiff . contains ( & config:: AutoDiff :: NoPostopt ) {
0 commit comments