@@ -13,7 +13,7 @@ use crate::builder::{Builder, PlaceRef, UNNAMED};
1313use crate :: context:: SimpleCx ;
1414use crate :: declare:: declare_simple_fn;
1515use crate :: llvm;
16- use crate :: llvm:: { Metadata , TRUE , Type } ;
16+ use crate :: llvm:: { TRUE , Type } ;
1717use crate :: value:: Value ;
1818
1919pub ( crate ) fn adjust_activity_to_abi < ' tcx > (
@@ -159,32 +159,36 @@ fn match_args_from_caller_to_enzyme<'ll, 'tcx>(
159159 let mut outer_pos: usize = 0 ;
160160 let mut activity_pos = 0 ;
161161
162- let enzyme_const = cx. create_metadata ( b"enzyme_const" ) ;
163- let enzyme_out = cx. create_metadata ( b"enzyme_out" ) ;
164- let enzyme_dup = cx. create_metadata ( b"enzyme_dup" ) ;
165- let enzyme_dupv = cx. create_metadata ( b"enzyme_dupv" ) ;
166- let enzyme_dupnoneed = cx. create_metadata ( b"enzyme_dupnoneed" ) ;
167- let enzyme_dupnoneedv = cx. create_metadata ( b"enzyme_dupnoneedv" ) ;
162+ // We used to use llvm's metadata to instruct enzyme how to differentiate a function.
163+ // In debug mode we would use incremental compilation which caused the metadata to be
164+ // dropped. This is prevented by now using named globals, which are also understood
165+ // by Enzyme.
166+ let global_const = cx. declare_global ( "enzyme_const" , cx. type_ptr ( ) ) ;
167+ let global_out = cx. declare_global ( "enzyme_out" , cx. type_ptr ( ) ) ;
168+ let global_dup = cx. declare_global ( "enzyme_dup" , cx. type_ptr ( ) ) ;
169+ let global_dupv = cx. declare_global ( "enzyme_dupv" , cx. type_ptr ( ) ) ;
170+ let global_dupnoneed = cx. declare_global ( "enzyme_dupnoneed" , cx. type_ptr ( ) ) ;
171+ let global_dupnoneedv = cx. declare_global ( "enzyme_dupnoneedv" , cx. type_ptr ( ) ) ;
168172
169173 while activity_pos < inputs. len ( ) {
170174 let diff_activity = inputs[ activity_pos as usize ] ;
171175 // Duplicated arguments received a shadow argument, into which enzyme will write the
172176 // gradient.
173- let ( activity, duplicated) : ( & Metadata , bool ) = match diff_activity {
177+ let ( activity, duplicated) : ( & llvm :: Value , bool ) = match diff_activity {
174178 DiffActivity :: None => panic ! ( "not a valid input activity" ) ,
175- DiffActivity :: Const => ( enzyme_const , false ) ,
176- DiffActivity :: Active => ( enzyme_out , false ) ,
177- DiffActivity :: ActiveOnly => ( enzyme_out , false ) ,
178- DiffActivity :: Dual => ( enzyme_dup , true ) ,
179- DiffActivity :: Dualv => ( enzyme_dupv , true ) ,
180- DiffActivity :: DualOnly => ( enzyme_dupnoneed , true ) ,
181- DiffActivity :: DualvOnly => ( enzyme_dupnoneedv , true ) ,
182- DiffActivity :: Duplicated => ( enzyme_dup , true ) ,
183- DiffActivity :: DuplicatedOnly => ( enzyme_dupnoneed , true ) ,
184- DiffActivity :: FakeActivitySize ( _) => ( enzyme_const , false ) ,
179+ DiffActivity :: Const => ( global_const , false ) ,
180+ DiffActivity :: Active => ( global_out , false ) ,
181+ DiffActivity :: ActiveOnly => ( global_out , false ) ,
182+ DiffActivity :: Dual => ( global_dup , true ) ,
183+ DiffActivity :: Dualv => ( global_dupv , true ) ,
184+ DiffActivity :: DualOnly => ( global_dupnoneed , true ) ,
185+ DiffActivity :: DualvOnly => ( global_dupnoneedv , true ) ,
186+ DiffActivity :: Duplicated => ( global_dup , true ) ,
187+ DiffActivity :: DuplicatedOnly => ( global_dupnoneed , true ) ,
188+ DiffActivity :: FakeActivitySize ( _) => ( global_const , false ) ,
185189 } ;
186190 let outer_arg = outer_args[ outer_pos] ;
187- args. push ( cx . get_metadata_value ( activity) ) ;
191+ args. push ( activity) ;
188192 if matches ! ( diff_activity, DiffActivity :: Dualv ) {
189193 let next_outer_arg = outer_args[ outer_pos + 1 ] ;
190194 let elem_bytes_size: u64 = match inputs[ activity_pos + 1 ] {
@@ -244,7 +248,7 @@ fn match_args_from_caller_to_enzyme<'ll, 'tcx>(
244248 assert_eq ! ( cx. type_kind( next_outer_ty3) , TypeKind :: Integer ) ;
245249 args. push ( next_outer_arg2) ;
246250 }
247- args. push ( cx . get_metadata_value ( enzyme_const ) ) ;
251+ args. push ( global_const ) ;
248252 args. push ( next_outer_arg) ;
249253 outer_pos += 2 + 2 * iterations;
250254 activity_pos += 2 ;
@@ -353,13 +357,13 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>(
353357 let mut args = Vec :: with_capacity ( num_args as usize + 1 ) ;
354358 args. push ( fn_to_diff) ;
355359
356- let enzyme_primal_ret = cx. create_metadata ( b "enzyme_primal_return") ;
360+ let global_primal_ret = cx. declare_global ( "enzyme_primal_return" , cx . type_ptr ( ) ) ;
357361 if matches ! ( attrs. ret_activity, DiffActivity :: Dual | DiffActivity :: Active ) {
358- args. push ( cx . get_metadata_value ( enzyme_primal_ret ) ) ;
362+ args. push ( global_primal_ret ) ;
359363 }
360364 if attrs. width > 1 {
361- let enzyme_width = cx. create_metadata ( b "enzyme_width") ;
362- args. push ( cx . get_metadata_value ( enzyme_width ) ) ;
365+ let global_width = cx. declare_global ( "enzyme_width" , cx . type_ptr ( ) ) ;
366+ args. push ( global_width ) ;
363367 args. push ( cx. get_const_int ( cx. type_i64 ( ) , attrs. width as u64 ) ) ;
364368 }
365369
0 commit comments