@@ -7,6 +7,7 @@ use std::sync::mpsc::{Receiver, Sender, channel};
7
7
use std:: { fs, io, mem, str, thread} ;
8
8
9
9
use rustc_ast:: attr;
10
+ use rustc_ast:: expand:: autodiff_attrs:: AutoDiffItem ;
10
11
use rustc_data_structures:: fx:: { FxHashMap , FxIndexMap } ;
11
12
use rustc_data_structures:: jobserver:: { self , Acquired } ;
12
13
use rustc_data_structures:: memmap:: Mmap ;
@@ -40,7 +41,7 @@ use tracing::debug;
40
41
use super :: link:: { self , ensure_removed} ;
41
42
use super :: lto:: { self , SerializedModule } ;
42
43
use super :: symbol_export:: symbol_name_for_instance_in_crate;
43
- use crate :: errors:: ErrorCreatingRemarkDir ;
44
+ use crate :: errors:: { AutodiffWithoutLto , ErrorCreatingRemarkDir } ;
44
45
use crate :: traits:: * ;
45
46
use crate :: {
46
47
CachedModuleCodegen , CodegenResults , CompiledModule , CrateInfo , ModuleCodegen , ModuleKind ,
@@ -118,6 +119,7 @@ pub struct ModuleConfig {
118
119
pub merge_functions : bool ,
119
120
pub emit_lifetime_markers : bool ,
120
121
pub llvm_plugins : Vec < String > ,
122
+ pub autodiff : Vec < config:: AutoDiff > ,
121
123
}
122
124
123
125
impl ModuleConfig {
@@ -266,6 +268,7 @@ impl ModuleConfig {
266
268
267
269
emit_lifetime_markers : sess. emit_lifetime_markers ( ) ,
268
270
llvm_plugins : if_regular ! ( sess. opts. unstable_opts. llvm_plugins. clone( ) , vec![ ] ) ,
271
+ autodiff : if_regular ! ( sess. opts. unstable_opts. autodiff. clone( ) , vec![ ] ) ,
269
272
}
270
273
}
271
274
@@ -389,6 +392,7 @@ impl<B: WriteBackendMethods> CodegenContext<B> {
389
392
390
393
fn generate_lto_work < B : ExtraBackendMethods > (
391
394
cgcx : & CodegenContext < B > ,
395
+ autodiff : Vec < AutoDiffItem > ,
392
396
needs_fat_lto : Vec < FatLtoInput < B > > ,
393
397
needs_thin_lto : Vec < ( String , B :: ThinBuffer ) > ,
394
398
import_only_modules : Vec < ( SerializedModule < B :: ModuleBuffer > , WorkProduct ) > ,
@@ -399,9 +403,18 @@ fn generate_lto_work<B: ExtraBackendMethods>(
399
403
assert ! ( needs_thin_lto. is_empty( ) ) ;
400
404
let module =
401
405
B :: run_fat_lto ( cgcx, needs_fat_lto, import_only_modules) . unwrap_or_else ( |e| e. raise ( ) ) ;
406
+ if cgcx. lto == Lto :: Fat {
407
+ let _config = cgcx. config ( ModuleKind :: Regular ) ;
408
+ todo ! ( "fat LTO with autodiff is not yet implemented" ) ;
409
+ //module = unsafe { module.autodiff(cgcx, autodiff, config).unwrap() };
410
+ }
402
411
// We are adding a single work item, so the cost doesn't matter.
403
412
vec ! [ ( WorkItem :: LTO ( module) , 0 ) ]
404
413
} else {
414
+ if !autodiff. is_empty ( ) {
415
+ let dcx = cgcx. create_dcx ( ) ;
416
+ dcx. handle ( ) . emit_fatal ( AutodiffWithoutLto { } ) ;
417
+ }
405
418
assert ! ( needs_fat_lto. is_empty( ) ) ;
406
419
let ( lto_modules, copy_jobs) = B :: run_thin_lto ( cgcx, needs_thin_lto, import_only_modules)
407
420
. unwrap_or_else ( |e| e. raise ( ) ) ;
@@ -1021,6 +1034,9 @@ pub(crate) enum Message<B: WriteBackendMethods> {
1021
1034
/// Sent from a backend worker thread.
1022
1035
WorkItem { result : Result < WorkItemResult < B > , Option < WorkerFatalError > > , worker_id : usize } ,
1023
1036
1037
+ /// A vector containing all the AutoDiff tasks that we have to pass to Enzyme.
1038
+ AddAutoDiffItems ( Vec < AutoDiffItem > ) ,
1039
+
1024
1040
/// The frontend has finished generating something (backend IR or a
1025
1041
/// post-LTO artifact) for a codegen unit, and it should be passed to the
1026
1042
/// backend. Sent from the main thread.
@@ -1348,6 +1364,7 @@ fn start_executing_work<B: ExtraBackendMethods>(
1348
1364
1349
1365
// This is where we collect codegen units that have gone all the way
1350
1366
// through codegen and LLVM.
1367
+ let mut autodiff_items = Vec :: new ( ) ;
1351
1368
let mut compiled_modules = vec ! [ ] ;
1352
1369
let mut compiled_allocator_module = None ;
1353
1370
let mut needs_link = Vec :: new ( ) ;
@@ -1459,9 +1476,13 @@ fn start_executing_work<B: ExtraBackendMethods>(
1459
1476
let needs_thin_lto = mem:: take ( & mut needs_thin_lto) ;
1460
1477
let import_only_modules = mem:: take ( & mut lto_import_only_modules) ;
1461
1478
1462
- for ( work, cost) in
1463
- generate_lto_work ( & cgcx, needs_fat_lto, needs_thin_lto, import_only_modules)
1464
- {
1479
+ for ( work, cost) in generate_lto_work (
1480
+ & cgcx,
1481
+ autodiff_items. clone ( ) ,
1482
+ needs_fat_lto,
1483
+ needs_thin_lto,
1484
+ import_only_modules,
1485
+ ) {
1465
1486
let insertion_index = work_items
1466
1487
. binary_search_by_key ( & cost, |& ( _, cost) | cost)
1467
1488
. unwrap_or_else ( |e| e) ;
@@ -1596,6 +1617,10 @@ fn start_executing_work<B: ExtraBackendMethods>(
1596
1617
main_thread_state = MainThreadState :: Idle ;
1597
1618
}
1598
1619
1620
+ Message :: AddAutoDiffItems ( mut items) => {
1621
+ autodiff_items. append ( & mut items) ;
1622
+ }
1623
+
1599
1624
Message :: CodegenComplete => {
1600
1625
if codegen_state != Aborted {
1601
1626
codegen_state = Completed ;
@@ -2070,6 +2095,10 @@ impl<B: ExtraBackendMethods> OngoingCodegen<B> {
2070
2095
drop ( self . coordinator . sender . send ( Box :: new ( Message :: CodegenComplete :: < B > ) ) ) ;
2071
2096
}
2072
2097
2098
+ pub ( crate ) fn submit_autodiff_items ( & self , items : Vec < AutoDiffItem > ) {
2099
+ drop ( self . coordinator . sender . send ( Box :: new ( Message :: < B > :: AddAutoDiffItems ( items) ) ) ) ;
2100
+ }
2101
+
2073
2102
pub ( crate ) fn check_for_errors ( & self , sess : & Session ) {
2074
2103
self . shared_emitter_main . check ( sess, false ) ;
2075
2104
}
0 commit comments