rust-lang/rust
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Pass in autodiff items when starting the coordinator thread

Browse Source 
As opposed to sending a message to the coordinator thread.
This commit is contained in:
bjorn3
2025-07-03 15:39:19 +00:00
parent 797084dc92
commit 332f1b1bdd
2 changed files with 6 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
compiler/rustc_codegen_ssa/src/back/write.rs
View File

@@ -470,6 +470,7 @@ pub(crate) fn start_async_codegen<B: ExtraBackendMethods>(
backend: B,
tcx: TyCtxt<'_>,
target_cpu: String,
autodiff_items: &[AutoDiffItem],
) -> OngoingCodegen<B> {
let (coordinator_send, coordinator_receive) = channel();
@@ -488,6 +489,7 @@ pub(crate) fn start_async_codegen<B: ExtraBackendMethods>(
backend.clone(),
tcx,
&crate_info,
autodiff_items,
shared_emitter,
codegen_worker_send,
coordinator_receive,
@@ -1044,9 +1046,6 @@ pub(crate) enum Message<B: WriteBackendMethods> {
/// Sent from a backend worker thread.
WorkItem { result: Result<WorkItemResult<B>, Option<WorkerFatalError>>, worker_id: usize },
/// A vector containing all the AutoDiff tasks that we have to pass to Enzyme.
AddAutoDiffItems(Vec<AutoDiffItem>),
/// The frontend has finished generating something (backend IR or a
/// post-LTO artifact) for a codegen unit, and it should be passed to the
/// backend. Sent from the main thread.
@@ -1113,6 +1112,7 @@ fn start_executing_work<B: ExtraBackendMethods>(
backend: B,
tcx: TyCtxt<'_>,
crate_info: &CrateInfo,
autodiff_items: &[AutoDiffItem],
shared_emitter: SharedEmitter,
codegen_worker_send: Sender<CguMessage>,
coordinator_receive: Receiver<Box<dyn Any + Send>>,
@@ -1122,6 +1122,7 @@ fn start_executing_work<B: ExtraBackendMethods>(
) -> thread::JoinHandle<Result<CompiledModules, ()>> {
let coordinator_send = tx_to_llvm_workers;
let sess = tcx.sess;
let autodiff_items = autodiff_items.to_vec();
let mut each_linked_rlib_for_lto = Vec::new();
drop(link::each_linked_rlib(crate_info, None, &mut |cnum, path| {
@@ -1375,7 +1376,6 @@ fn start_executing_work<B: ExtraBackendMethods>(
// This is where we collect codegen units that have gone all the way
// through codegen and LLVM.
let mut autodiff_items = Vec::new();
let mut compiled_modules = vec![];
let mut compiled_allocator_module = None;
let mut needs_link = Vec::new();
@@ -1645,10 +1645,6 @@ fn start_executing_work<B: ExtraBackendMethods>(
main_thread_state = MainThreadState::Idle;
}
Message::AddAutoDiffItems(mut items) => {
autodiff_items.append(&mut items);
}
Message::CodegenComplete => {
if codegen_state != Aborted {
codegen_state = Completed;
@@ -2117,10 +2113,6 @@ impl<B: ExtraBackendMethods> OngoingCodegen<B> {
drop(self.coordinator.sender.send(Box::new(Message::CodegenComplete::<B>)));
}
pub(crate) fn submit_autodiff_items(&self, items: Vec<AutoDiffItem>) {
drop(self.coordinator.sender.send(Box::new(Message::<B>::AddAutoDiffItems(items))));
}
pub(crate) fn check_for_errors(&self, sess: &Session) {
self.shared_emitter_main.check(sess, false);
}