From 653bb64c7530c86d00b875e1ef14ac75455d2d80 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Thu, 3 Jul 2025 14:24:58 +0000
Subject: [PATCH 1/9] Remove LtoModuleCodegen

Most uses of it either contain a fat or thin lto module. Only
WorkItem::LTO could contain both, but splitting that enum variant
doesn't complicate things much.
---
 compiler/rustc_codegen_gcc/src/back/lto.rs    | 15 +++--
 compiler/rustc_codegen_gcc/src/lib.rs         |  6 +-
 compiler/rustc_codegen_llvm/src/back/lto.rs   | 17 +++---
 compiler/rustc_codegen_llvm/src/lib.rs        |  6 +-
 compiler/rustc_codegen_ssa/src/back/lto.rs    | 60 -------------------
 compiler/rustc_codegen_ssa/src/back/write.rs  | 48 +++++++++++----
 .../rustc_codegen_ssa/src/traits/write.rs     |  6 +-
 7 files changed, 58 insertions(+), 100 deletions(-)

diff --git a/compiler/rustc_codegen_gcc/src/back/lto.rs b/compiler/rustc_codegen_gcc/src/back/lto.rs
index 10fce860b7770..e554dd2500bda 100644
--- a/compiler/rustc_codegen_gcc/src/back/lto.rs
+++ b/compiler/rustc_codegen_gcc/src/back/lto.rs
@@ -24,7 +24,7 @@ use std::sync::Arc;
 
 use gccjit::{Context, OutputKind};
 use object::read::archive::ArchiveFile;
-use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule, ThinShared};
+use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule, ThinShared};
 use rustc_codegen_ssa::back::symbol_export;
 use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
 use rustc_codegen_ssa::traits::*;
@@ -176,7 +176,7 @@ pub(crate) fn run_fat(
     cgcx: &CodegenContext<GccCodegenBackend>,
     modules: Vec<FatLtoInput<GccCodegenBackend>>,
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
-) -> Result<LtoModuleCodegen<GccCodegenBackend>, FatalError> {
+) -> Result<ModuleCodegen<GccContext>, FatalError> {
     let dcx = cgcx.create_dcx();
     let dcx = dcx.handle();
     let lto_data = prepare_lto(cgcx, dcx)?;
@@ -201,7 +201,7 @@ fn fat_lto(
     mut serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
     tmp_path: TempDir,
     //symbols_below_threshold: &[String],
-) -> Result<LtoModuleCodegen<GccCodegenBackend>, FatalError> {
+) -> Result<ModuleCodegen<GccContext>, FatalError> {
     let _timer = cgcx.prof.generic_activity("GCC_fat_lto_build_monolithic_module");
     info!("going for a fat lto");
 
@@ -334,7 +334,7 @@ fn fat_lto(
     // of now.
     module.module_llvm.temp_dir = Some(tmp_path);
 
-    Ok(LtoModuleCodegen::Fat(module))
+    Ok(module)
 }
 
 pub struct ModuleBuffer(PathBuf);
@@ -358,7 +358,7 @@ pub(crate) fn run_thin(
     cgcx: &CodegenContext<GccCodegenBackend>,
     modules: Vec<(String, ThinBuffer)>,
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
-) -> Result<(Vec<LtoModuleCodegen<GccCodegenBackend>>, Vec<WorkProduct>), FatalError> {
+) -> Result<(Vec<ThinModule<GccCodegenBackend>>, Vec<WorkProduct>), FatalError> {
     let dcx = cgcx.create_dcx();
     let dcx = dcx.handle();
     let lto_data = prepare_lto(cgcx, dcx)?;
@@ -427,7 +427,7 @@ fn thin_lto(
     tmp_path: TempDir,
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
     //_symbols_below_threshold: &[String],
-) -> Result<(Vec<LtoModuleCodegen<GccCodegenBackend>>, Vec<WorkProduct>), FatalError> {
+) -> Result<(Vec<ThinModule<GccCodegenBackend>>, Vec<WorkProduct>), FatalError> {
     let _timer = cgcx.prof.generic_activity("LLVM_thin_lto_global_analysis");
     info!("going for that thin, thin LTO");
 
@@ -573,8 +573,7 @@ fn thin_lto(
         }*/
 
         info!(" - {}: re-compiled", module_name);
-        opt_jobs
-            .push(LtoModuleCodegen::Thin(ThinModule { shared: shared.clone(), idx: module_index }));
+        opt_jobs.push(ThinModule { shared: shared.clone(), idx: module_index });
     }
 
     // Save the current ThinLTO import information for the next compilation
diff --git a/compiler/rustc_codegen_gcc/src/lib.rs b/compiler/rustc_codegen_gcc/src/lib.rs
index a912678ef2a10..a151a0ab75534 100644
--- a/compiler/rustc_codegen_gcc/src/lib.rs
+++ b/compiler/rustc_codegen_gcc/src/lib.rs
@@ -97,7 +97,7 @@ use gccjit::{CType, Context, OptimizationLevel};
 use gccjit::{TargetInfo, Version};
 use rustc_ast::expand::allocator::AllocatorKind;
 use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
-use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};
+use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule};
 use rustc_codegen_ssa::back::write::{
     CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryFn,
 };
@@ -361,7 +361,7 @@ impl WriteBackendMethods for GccCodegenBackend {
         cgcx: &CodegenContext<Self>,
         modules: Vec<FatLtoInput<Self>>,
         cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
-    ) -> Result<LtoModuleCodegen<Self>, FatalError> {
+    ) -> Result<ModuleCodegen<Self::Module>, FatalError> {
         back::lto::run_fat(cgcx, modules, cached_modules)
     }
 
@@ -369,7 +369,7 @@ impl WriteBackendMethods for GccCodegenBackend {
         cgcx: &CodegenContext<Self>,
         modules: Vec<(String, Self::ThinBuffer)>,
         cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
-    ) -> Result<(Vec<LtoModuleCodegen<Self>>, Vec<WorkProduct>), FatalError> {
+    ) -> Result<(Vec<ThinModule<Self>>, Vec<WorkProduct>), FatalError> {
         back::lto::run_thin(cgcx, modules, cached_modules)
     }
 
diff --git a/compiler/rustc_codegen_llvm/src/back/lto.rs b/compiler/rustc_codegen_llvm/src/back/lto.rs
index 9c62244f3c9ff..0198b9f0cf0c6 100644
--- a/compiler/rustc_codegen_llvm/src/back/lto.rs
+++ b/compiler/rustc_codegen_llvm/src/back/lto.rs
@@ -7,7 +7,7 @@ use std::sync::Arc;
 use std::{io, iter, slice};
 
 use object::read::archive::ArchiveFile;
-use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule, ThinShared};
+use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule, ThinShared};
 use rustc_codegen_ssa::back::symbol_export;
 use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
 use rustc_codegen_ssa::traits::*;
@@ -201,7 +201,7 @@ pub(crate) fn run_fat(
     cgcx: &CodegenContext<LlvmCodegenBackend>,
     modules: Vec<FatLtoInput<LlvmCodegenBackend>>,
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
-) -> Result<LtoModuleCodegen<LlvmCodegenBackend>, FatalError> {
+) -> Result<ModuleCodegen<ModuleLlvm>, FatalError> {
     let dcx = cgcx.create_dcx();
     let dcx = dcx.handle();
     let (symbols_below_threshold, upstream_modules) = prepare_lto(cgcx, dcx)?;
@@ -217,7 +217,7 @@ pub(crate) fn run_thin(
     cgcx: &CodegenContext<LlvmCodegenBackend>,
     modules: Vec<(String, ThinBuffer)>,
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
-) -> Result<(Vec<LtoModuleCodegen<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
+) -> Result<(Vec<ThinModule<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
     let dcx = cgcx.create_dcx();
     let dcx = dcx.handle();
     let (symbols_below_threshold, upstream_modules) = prepare_lto(cgcx, dcx)?;
@@ -248,7 +248,7 @@ fn fat_lto(
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
     mut serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
     symbols_below_threshold: &[*const libc::c_char],
-) -> Result<LtoModuleCodegen<LlvmCodegenBackend>, FatalError> {
+) -> Result<ModuleCodegen<ModuleLlvm>, FatalError> {
     let _timer = cgcx.prof.generic_activity("LLVM_fat_lto_build_monolithic_module");
     info!("going for a fat lto");
 
@@ -366,7 +366,7 @@ fn fat_lto(
         save_temp_bitcode(cgcx, &module, "lto.after-restriction");
     }
 
-    Ok(LtoModuleCodegen::Fat(module))
+    Ok(module)
 }
 
 pub(crate) struct Linker<'a>(&'a mut llvm::Linker<'a>);
@@ -436,7 +436,7 @@ fn thin_lto(
     serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
     cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
     symbols_below_threshold: &[*const libc::c_char],
-) -> Result<(Vec<LtoModuleCodegen<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
+) -> Result<(Vec<ThinModule<LlvmCodegenBackend>>, Vec<WorkProduct>), FatalError> {
     let _timer = cgcx.prof.generic_activity("LLVM_thin_lto_global_analysis");
     unsafe {
         info!("going for that thin, thin LTO");
@@ -568,10 +568,7 @@ fn thin_lto(
             }
 
             info!(" - {}: re-compiled", module_name);
-            opt_jobs.push(LtoModuleCodegen::Thin(ThinModule {
-                shared: Arc::clone(&shared),
-                idx: module_index,
-            }));
+            opt_jobs.push(ThinModule { shared: Arc::clone(&shared), idx: module_index });
         }
 
         // Save the current ThinLTO import information for the next compilation
diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs
index cdfffbe47bfa5..57b37959e0d01 100644
--- a/compiler/rustc_codegen_llvm/src/lib.rs
+++ b/compiler/rustc_codegen_llvm/src/lib.rs
@@ -30,7 +30,7 @@ use errors::{AutoDiffWithoutLTO, ParseTargetMachineConfig};
 use llvm_util::target_config;
 use rustc_ast::expand::allocator::AllocatorKind;
 use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
-use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};
+use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule};
 use rustc_codegen_ssa::back::write::{
     CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryConfig, TargetMachineFactoryFn,
 };
@@ -178,14 +178,14 @@ impl WriteBackendMethods for LlvmCodegenBackend {
         cgcx: &CodegenContext<Self>,
         modules: Vec<FatLtoInput<Self>>,
         cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
-    ) -> Result<LtoModuleCodegen<Self>, FatalError> {
+    ) -> Result<ModuleCodegen<Self::Module>, FatalError> {
         back::lto::run_fat(cgcx, modules, cached_modules)
     }
     fn run_thin_lto(
         cgcx: &CodegenContext<Self>,
         modules: Vec<(String, Self::ThinBuffer)>,
         cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
-    ) -> Result<(Vec<LtoModuleCodegen<Self>>, Vec<WorkProduct>), FatalError> {
+    ) -> Result<(Vec<ThinModule<Self>>, Vec<WorkProduct>), FatalError> {
         back::lto::run_thin(cgcx, modules, cached_modules)
     }
     fn optimize(
diff --git a/compiler/rustc_codegen_ssa/src/back/lto.rs b/compiler/rustc_codegen_ssa/src/back/lto.rs
index ce6fe8a191b3b..b49b6783bbd9a 100644
--- a/compiler/rustc_codegen_ssa/src/back/lto.rs
+++ b/compiler/rustc_codegen_ssa/src/back/lto.rs
@@ -1,13 +1,8 @@
 use std::ffi::CString;
 use std::sync::Arc;
 
-use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
 use rustc_data_structures::memmap::Mmap;
-use rustc_errors::FatalError;
 
-use super::write::CodegenContext;
-use crate::ModuleCodegen;
-use crate::back::write::ModuleConfig;
 use crate::traits::*;
 
 pub struct ThinModule<B: WriteBackendMethods> {
@@ -42,61 +37,6 @@ pub struct ThinShared<B: WriteBackendMethods> {
     pub module_names: Vec<CString>,
 }
 
-pub enum LtoModuleCodegen<B: WriteBackendMethods> {
-    Fat(ModuleCodegen<B::Module>),
-    Thin(ThinModule<B>),
-}
-
-impl<B: WriteBackendMethods> LtoModuleCodegen<B> {
-    pub fn name(&self) -> &str {
-        match *self {
-            LtoModuleCodegen::Fat(_) => "everything",
-            LtoModuleCodegen::Thin(ref m) => m.name(),
-        }
-    }
-
-    /// Optimize this module within the given codegen context.
-    pub fn optimize(
-        self,
-        cgcx: &CodegenContext<B>,
-    ) -> Result<ModuleCodegen<B::Module>, FatalError> {
-        match self {
-            LtoModuleCodegen::Fat(mut module) => {
-                B::optimize_fat(cgcx, &mut module)?;
-                Ok(module)
-            }
-            LtoModuleCodegen::Thin(thin) => B::optimize_thin(cgcx, thin),
-        }
-    }
-
-    /// A "gauge" of how costly it is to optimize this module, used to sort
-    /// biggest modules first.
-    pub fn cost(&self) -> u64 {
-        match *self {
-            // Only one module with fat LTO, so the cost doesn't matter.
-            LtoModuleCodegen::Fat(_) => 0,
-            LtoModuleCodegen::Thin(ref m) => m.cost(),
-        }
-    }
-
-    /// Run autodiff on Fat LTO module
-    pub fn autodiff(
-        self,
-        cgcx: &CodegenContext<B>,
-        diff_fncs: Vec<AutoDiffItem>,
-        config: &ModuleConfig,
-    ) -> Result<LtoModuleCodegen<B>, FatalError> {
-        match &self {
-            LtoModuleCodegen::Fat(module) => {
-                B::autodiff(cgcx, &module, diff_fncs, config)?;
-            }
-            _ => panic!("autodiff called with non-fat LTO module"),
-        }
-
-        Ok(self)
-    }
-}
-
 pub enum SerializedModule<M: ModuleBufferMethods> {
     Local(M),
     FromRlib(Vec<u8>),
diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
index 8330e4f7af0c7..5395b55092b3c 100644
--- a/compiler/rustc_codegen_ssa/src/back/write.rs
+++ b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -408,14 +408,16 @@ fn generate_lto_work<B: ExtraBackendMethods>(
 
     if !needs_fat_lto.is_empty() {
         assert!(needs_thin_lto.is_empty());
-        let mut module =
+        let module =
             B::run_fat_lto(cgcx, needs_fat_lto, import_only_modules).unwrap_or_else(|e| e.raise());
         if cgcx.lto == Lto::Fat && !autodiff.is_empty() {
             let config = cgcx.config(ModuleKind::Regular);
-            module = module.autodiff(cgcx, autodiff, config).unwrap_or_else(|e| e.raise());
+            if let Err(err) = B::autodiff(cgcx, &module, autodiff, config) {
+                err.raise();
+            }
         }
         // We are adding a single work item, so the cost doesn't matter.
-        vec![(WorkItem::LTO(module), 0)]
+        vec![(WorkItem::FatLto(module), 0)]
     } else {
         if !autodiff.is_empty() {
             let dcx = cgcx.create_dcx();
@@ -428,7 +430,7 @@ fn generate_lto_work<B: ExtraBackendMethods>(
             .into_iter()
             .map(|module| {
                 let cost = module.cost();
-                (WorkItem::LTO(module), cost)
+                (WorkItem::ThinLto(module), cost)
             })
             .chain(copy_jobs.into_iter().map(|wp| {
                 (
@@ -736,15 +738,19 @@ pub(crate) enum WorkItem<B: WriteBackendMethods> {
     /// Copy the post-LTO artifacts from the incremental cache to the output
     /// directory.
     CopyPostLtoArtifacts(CachedModuleCodegen),
-    /// Performs (Thin)LTO on the given module.
-    LTO(lto::LtoModuleCodegen<B>),
+    /// Performs fat LTO on the given module.
+    FatLto(ModuleCodegen<B::Module>),
+    /// Performs thin-LTO on the given module.
+    ThinLto(lto::ThinModule<B>),
 }
 
 impl<B: WriteBackendMethods> WorkItem<B> {
     fn module_kind(&self) -> ModuleKind {
         match *self {
             WorkItem::Optimize(ref m) => m.kind,
-            WorkItem::CopyPostLtoArtifacts(_) | WorkItem::LTO(_) => ModuleKind::Regular,
+            WorkItem::CopyPostLtoArtifacts(_) | WorkItem::FatLto(_) | WorkItem::ThinLto(_) => {
+                ModuleKind::Regular
+            }
         }
     }
 
@@ -792,7 +798,8 @@ impl<B: WriteBackendMethods> WorkItem<B> {
         match self {
             WorkItem::Optimize(m) => desc("opt", "optimize module", &m.name),
             WorkItem::CopyPostLtoArtifacts(m) => desc("cpy", "copy LTO artifacts for", &m.name),
-            WorkItem::LTO(m) => desc("lto", "LTO module", m.name()),
+            WorkItem::FatLto(_) => desc("lto", "fat LTO module", "everything"),
+            WorkItem::ThinLto(m) => desc("lto", "thin-LTO module", m.name()),
         }
     }
 }
@@ -996,12 +1003,21 @@ fn execute_copy_from_cache_work_item<B: ExtraBackendMethods>(
     })
 }
 
-fn execute_lto_work_item<B: ExtraBackendMethods>(
+fn execute_fat_lto_work_item<B: ExtraBackendMethods>(
+    cgcx: &CodegenContext<B>,
+    mut module: ModuleCodegen<B::Module>,
+    module_config: &ModuleConfig,
+) -> Result<WorkItemResult<B>, FatalError> {
+    B::optimize_fat(cgcx, &mut module)?;
+    finish_intra_module_work(cgcx, module, module_config)
+}
+
+fn execute_thin_lto_work_item<B: ExtraBackendMethods>(
     cgcx: &CodegenContext<B>,
-    module: lto::LtoModuleCodegen<B>,
+    module: lto::ThinModule<B>,
     module_config: &ModuleConfig,
 ) -> Result<WorkItemResult<B>, FatalError> {
-    let module = module.optimize(cgcx)?;
+    let module = B::optimize_thin(cgcx, module)?;
     finish_intra_module_work(cgcx, module, module_config)
 }
 
@@ -1842,10 +1858,16 @@ fn spawn_work<'a, B: ExtraBackendMethods>(
                     );
                     Ok(execute_copy_from_cache_work_item(&cgcx, m, module_config))
                 }
-                WorkItem::LTO(m) => {
+                WorkItem::FatLto(m) => {
+                    let _timer = cgcx
+                        .prof
+                        .generic_activity_with_arg("codegen_module_perform_lto", "everything");
+                    execute_fat_lto_work_item(&cgcx, m, module_config)
+                }
+                WorkItem::ThinLto(m) => {
                     let _timer =
                         cgcx.prof.generic_activity_with_arg("codegen_module_perform_lto", m.name());
-                    execute_lto_work_item(&cgcx, m, module_config)
+                    execute_thin_lto_work_item(&cgcx, m, module_config)
                 }
             })
         };
diff --git a/compiler/rustc_codegen_ssa/src/traits/write.rs b/compiler/rustc_codegen_ssa/src/traits/write.rs
index 07a0609fda1a1..78380973ffc80 100644
--- a/compiler/rustc_codegen_ssa/src/traits/write.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/write.rs
@@ -2,7 +2,7 @@ use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
 use rustc_errors::{DiagCtxtHandle, FatalError};
 use rustc_middle::dep_graph::WorkProduct;
 
-use crate::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};
+use crate::back::lto::{SerializedModule, ThinModule};
 use crate::back::write::{CodegenContext, FatLtoInput, ModuleConfig};
 use crate::{CompiledModule, ModuleCodegen};
 
@@ -26,7 +26,7 @@ pub trait WriteBackendMethods: Clone + 'static {
         cgcx: &CodegenContext<Self>,
         modules: Vec<FatLtoInput<Self>>,
         cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
-    ) -> Result<LtoModuleCodegen<Self>, FatalError>;
+    ) -> Result<ModuleCodegen<Self::Module>, FatalError>;
     /// Performs thin LTO by performing necessary global analysis and returning two
     /// lists, one of the modules that need optimization and another for modules that
     /// can simply be copied over from the incr. comp. cache.
@@ -34,7 +34,7 @@ pub trait WriteBackendMethods: Clone + 'static {
         cgcx: &CodegenContext<Self>,
         modules: Vec<(String, Self::ThinBuffer)>,
         cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
-    ) -> Result<(Vec<LtoModuleCodegen<Self>>, Vec<WorkProduct>), FatalError>;
+    ) -> Result<(Vec<ThinModule<Self>>, Vec<WorkProduct>), FatalError>;
     fn print_pass_timings(&self);
     fn print_statistics(&self);
     fn optimize(

From 7fd78df3463c761a693ea79265a6a1faaf2ed51b Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Thu, 3 Jul 2025 14:43:09 +0000
Subject: [PATCH 2/9] Move dcx creation into WriteBackendMethods::codegen

---
 compiler/rustc_codegen_gcc/src/back/write.rs   | 4 +++-
 compiler/rustc_codegen_gcc/src/lib.rs          | 3 +--
 compiler/rustc_codegen_llvm/src/back/write.rs  | 4 +++-
 compiler/rustc_codegen_llvm/src/lib.rs         | 3 +--
 compiler/rustc_codegen_ssa/src/back/write.rs   | 7 ++-----
 compiler/rustc_codegen_ssa/src/traits/write.rs | 1 -
 6 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/compiler/rustc_codegen_gcc/src/back/write.rs b/compiler/rustc_codegen_gcc/src/back/write.rs
index d03d063bdace6..113abe70805b0 100644
--- a/compiler/rustc_codegen_gcc/src/back/write.rs
+++ b/compiler/rustc_codegen_gcc/src/back/write.rs
@@ -16,10 +16,12 @@ use crate::{GccCodegenBackend, GccContext};
 
 pub(crate) fn codegen(
     cgcx: &CodegenContext<GccCodegenBackend>,
-    dcx: DiagCtxtHandle<'_>,
     module: ModuleCodegen<GccContext>,
     config: &ModuleConfig,
 ) -> Result<CompiledModule, FatalError> {
+    let dcx = cgcx.create_dcx();
+    let dcx = dcx.handle();
+
     let _timer = cgcx.prof.generic_activity_with_arg("GCC_module_codegen", &*module.name);
     {
         let context = &module.module_llvm.context;
diff --git a/compiler/rustc_codegen_gcc/src/lib.rs b/compiler/rustc_codegen_gcc/src/lib.rs
index a151a0ab75534..34452bdd2005f 100644
--- a/compiler/rustc_codegen_gcc/src/lib.rs
+++ b/compiler/rustc_codegen_gcc/src/lib.rs
@@ -408,11 +408,10 @@ impl WriteBackendMethods for GccCodegenBackend {
 
     fn codegen(
         cgcx: &CodegenContext<Self>,
-        dcx: DiagCtxtHandle<'_>,
         module: ModuleCodegen<Self::Module>,
         config: &ModuleConfig,
     ) -> Result<CompiledModule, FatalError> {
-        back::write::codegen(cgcx, dcx, module, config)
+        back::write::codegen(cgcx, module, config)
     }
 
     fn prepare_thin(
diff --git a/compiler/rustc_codegen_llvm/src/back/write.rs b/compiler/rustc_codegen_llvm/src/back/write.rs
index bde6a9cf4bc61..1f7a785bbe725 100644
--- a/compiler/rustc_codegen_llvm/src/back/write.rs
+++ b/compiler/rustc_codegen_llvm/src/back/write.rs
@@ -817,10 +817,12 @@ pub(crate) fn link(
 
 pub(crate) fn codegen(
     cgcx: &CodegenContext<LlvmCodegenBackend>,
-    dcx: DiagCtxtHandle<'_>,
     module: ModuleCodegen<ModuleLlvm>,
     config: &ModuleConfig,
 ) -> Result<CompiledModule, FatalError> {
+    let dcx = cgcx.create_dcx();
+    let dcx = dcx.handle();
+
     let _timer = cgcx.prof.generic_activity_with_arg("LLVM_module_codegen", &*module.name);
     {
         let llmod = module.module_llvm.llmod();
diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs
index 57b37959e0d01..4d85683de8c43 100644
--- a/compiler/rustc_codegen_llvm/src/lib.rs
+++ b/compiler/rustc_codegen_llvm/src/lib.rs
@@ -212,11 +212,10 @@ impl WriteBackendMethods for LlvmCodegenBackend {
     }
     fn codegen(
         cgcx: &CodegenContext<Self>,
-        dcx: DiagCtxtHandle<'_>,
         module: ModuleCodegen<Self::Module>,
         config: &ModuleConfig,
     ) -> Result<CompiledModule, FatalError> {
-        back::write::codegen(cgcx, dcx, module, config)
+        back::write::codegen(cgcx, module, config)
     }
     fn prepare_thin(
         module: ModuleCodegen<Self::Module>,
diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
index 5395b55092b3c..29e67be6eb619 100644
--- a/compiler/rustc_codegen_ssa/src/back/write.rs
+++ b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -1026,11 +1026,8 @@ fn finish_intra_module_work<B: ExtraBackendMethods>(
     module: ModuleCodegen<B::Module>,
     module_config: &ModuleConfig,
 ) -> Result<WorkItemResult<B>, FatalError> {
-    let dcx = cgcx.create_dcx();
-    let dcx = dcx.handle();
-
     if !cgcx.opts.unstable_opts.combine_cgu || module.kind == ModuleKind::Allocator {
-        let module = B::codegen(cgcx, dcx, module, module_config)?;
+        let module = B::codegen(cgcx, module, module_config)?;
         Ok(WorkItemResult::Finished(module))
     } else {
         Ok(WorkItemResult::NeedsLink(module))
@@ -1718,7 +1715,7 @@ fn start_executing_work<B: ExtraBackendMethods>(
             let dcx = dcx.handle();
             let module = B::run_link(&cgcx, dcx, needs_link).map_err(|_| ())?;
             let module =
-                B::codegen(&cgcx, dcx, module, cgcx.config(ModuleKind::Regular)).map_err(|_| ())?;
+                B::codegen(&cgcx, module, cgcx.config(ModuleKind::Regular)).map_err(|_| ())?;
             compiled_modules.push(module);
         }
 
diff --git a/compiler/rustc_codegen_ssa/src/traits/write.rs b/compiler/rustc_codegen_ssa/src/traits/write.rs
index 78380973ffc80..c732c4b2ab860 100644
--- a/compiler/rustc_codegen_ssa/src/traits/write.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/write.rs
@@ -53,7 +53,6 @@ pub trait WriteBackendMethods: Clone + 'static {
     ) -> Result<ModuleCodegen<Self::Module>, FatalError>;
     fn codegen(
         cgcx: &CodegenContext<Self>,
-        dcx: DiagCtxtHandle<'_>,
         module: ModuleCodegen<Self::Module>,
         config: &ModuleConfig,
     ) -> Result<CompiledModule, FatalError>;

From ee280a596d3b7c401547be9a6afadf7298ee3a55 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Thu, 3 Jul 2025 14:44:26 +0000
Subject: [PATCH 3/9] Fat LTO always produces a single object file, so
 -Zcombine-cgu has no effect

---
 compiler/rustc_codegen_ssa/src/back/write.rs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
index 29e67be6eb619..762dc8f6c2d6c 100644
--- a/compiler/rustc_codegen_ssa/src/back/write.rs
+++ b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -1009,7 +1009,9 @@ fn execute_fat_lto_work_item<B: ExtraBackendMethods>(
     module_config: &ModuleConfig,
 ) -> Result<WorkItemResult<B>, FatalError> {
     B::optimize_fat(cgcx, &mut module)?;
-    finish_intra_module_work(cgcx, module, module_config)
+
+    let module = B::codegen(cgcx, module, module_config)?;
+    Ok(WorkItemResult::Finished(module))
 }
 
 fn execute_thin_lto_work_item<B: ExtraBackendMethods>(

From 797084dc92c251ff0c986adfdd4a7e261f499b02 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Thu, 3 Jul 2025 14:50:28 +0000
Subject: [PATCH 4/9] Split generate_lto_work between fat and thin-LTO cases

---
 compiler/rustc_codegen_ssa/src/back/write.rs | 112 +++++++++++--------
 1 file changed, 63 insertions(+), 49 deletions(-)

diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
index 762dc8f6c2d6c..21d5d3ced000c 100644
--- a/compiler/rustc_codegen_ssa/src/back/write.rs
+++ b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -397,52 +397,50 @@ impl<B: WriteBackendMethods> CodegenContext<B> {
     }
 }
 
-fn generate_lto_work<B: ExtraBackendMethods>(
+fn generate_fat_lto_work<B: ExtraBackendMethods>(
     cgcx: &CodegenContext<B>,
     autodiff: Vec<AutoDiffItem>,
     needs_fat_lto: Vec<FatLtoInput<B>>,
-    needs_thin_lto: Vec<(String, B::ThinBuffer)>,
     import_only_modules: Vec<(SerializedModule<B::ModuleBuffer>, WorkProduct)>,
-) -> Vec<(WorkItem<B>, u64)> {
+) -> WorkItem<B> {
     let _prof_timer = cgcx.prof.generic_activity("codegen_generate_lto_work");
 
-    if !needs_fat_lto.is_empty() {
-        assert!(needs_thin_lto.is_empty());
-        let module =
-            B::run_fat_lto(cgcx, needs_fat_lto, import_only_modules).unwrap_or_else(|e| e.raise());
-        if cgcx.lto == Lto::Fat && !autodiff.is_empty() {
-            let config = cgcx.config(ModuleKind::Regular);
-            if let Err(err) = B::autodiff(cgcx, &module, autodiff, config) {
-                err.raise();
-            }
+    let module =
+        B::run_fat_lto(cgcx, needs_fat_lto, import_only_modules).unwrap_or_else(|e| e.raise());
+    if cgcx.lto == Lto::Fat && !autodiff.is_empty() {
+        let config = cgcx.config(ModuleKind::Regular);
+        if let Err(err) = B::autodiff(cgcx, &module, autodiff, config) {
+            err.raise();
         }
-        // We are adding a single work item, so the cost doesn't matter.
-        vec![(WorkItem::FatLto(module), 0)]
-    } else {
-        if !autodiff.is_empty() {
-            let dcx = cgcx.create_dcx();
-            dcx.handle().emit_fatal(AutodiffWithoutLto {});
-        }
-        assert!(needs_fat_lto.is_empty());
-        let (lto_modules, copy_jobs) = B::run_thin_lto(cgcx, needs_thin_lto, import_only_modules)
-            .unwrap_or_else(|e| e.raise());
-        lto_modules
-            .into_iter()
-            .map(|module| {
-                let cost = module.cost();
-                (WorkItem::ThinLto(module), cost)
-            })
-            .chain(copy_jobs.into_iter().map(|wp| {
-                (
-                    WorkItem::CopyPostLtoArtifacts(CachedModuleCodegen {
-                        name: wp.cgu_name.clone(),
-                        source: wp,
-                    }),
-                    0, // copying is very cheap
-                )
-            }))
-            .collect()
     }
+    WorkItem::FatLto(module)
+}
+
+fn generate_thin_lto_work<B: ExtraBackendMethods>(
+    cgcx: &CodegenContext<B>,
+    needs_thin_lto: Vec<(String, B::ThinBuffer)>,
+    import_only_modules: Vec<(SerializedModule<B::ModuleBuffer>, WorkProduct)>,
+) -> Vec<(WorkItem<B>, u64)> {
+    let _prof_timer = cgcx.prof.generic_activity("codegen_thin_generate_lto_work");
+
+    let (lto_modules, copy_jobs) =
+        B::run_thin_lto(cgcx, needs_thin_lto, import_only_modules).unwrap_or_else(|e| e.raise());
+    lto_modules
+        .into_iter()
+        .map(|module| {
+            let cost = module.cost();
+            (WorkItem::ThinLto(module), cost)
+        })
+        .chain(copy_jobs.into_iter().map(|wp| {
+            (
+                WorkItem::CopyPostLtoArtifacts(CachedModuleCodegen {
+                    name: wp.cgu_name.clone(),
+                    source: wp,
+                }),
+                0, // copying is very cheap
+            )
+        }))
+        .collect()
 }
 
 struct CompiledModules {
@@ -1489,20 +1487,36 @@ fn start_executing_work<B: ExtraBackendMethods>(
                     let needs_thin_lto = mem::take(&mut needs_thin_lto);
                     let import_only_modules = mem::take(&mut lto_import_only_modules);
 
-                    for (work, cost) in generate_lto_work(
-                        &cgcx,
-                        autodiff_items.clone(),
-                        needs_fat_lto,
-                        needs_thin_lto,
-                        import_only_modules,
-                    ) {
-                        let insertion_index = work_items
-                            .binary_search_by_key(&cost, |&(_, cost)| cost)
-                            .unwrap_or_else(|e| e);
-                        work_items.insert(insertion_index, (work, cost));
+                    if !needs_fat_lto.is_empty() {
+                        assert!(needs_thin_lto.is_empty());
+
+                        let work = generate_fat_lto_work(
+                            &cgcx,
+                            autodiff_items.clone(),
+                            needs_fat_lto,
+                            import_only_modules,
+                        );
+                        work_items.push((work, 0));
                         if cgcx.parallel {
                             helper.request_token();
                         }
+                    } else {
+                        if !autodiff_items.is_empty() {
+                            let dcx = cgcx.create_dcx();
+                            dcx.handle().emit_fatal(AutodiffWithoutLto {});
+                        }
+
+                        for (work, cost) in
+                            generate_thin_lto_work(&cgcx, needs_thin_lto, import_only_modules)
+                        {
+                            let insertion_index = work_items
+                                .binary_search_by_key(&cost, |&(_, cost)| cost)
+                                .unwrap_or_else(|e| e);
+                            work_items.insert(insertion_index, (work, cost));
+                            if cgcx.parallel {
+                                helper.request_token();
+                            }
+                        }
                     }
                 }
 

From 332f1b1bdddb01f2f203f47cd3f30f677ae8d9ea Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Thu, 3 Jul 2025 15:39:19 +0000
Subject: [PATCH 5/9] Pass in autodiff items when starting the coordinator
 thread

As opposed to sending a message to the coordinator thread.
---
 compiler/rustc_codegen_ssa/src/back/write.rs | 16 ++++------------
 compiler/rustc_codegen_ssa/src/base.rs       |  9 ++-------
 2 files changed, 6 insertions(+), 19 deletions(-)

diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
index 21d5d3ced000c..dcee7aec73392 100644
--- a/compiler/rustc_codegen_ssa/src/back/write.rs
+++ b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -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);
     }
diff --git a/compiler/rustc_codegen_ssa/src/base.rs b/compiler/rustc_codegen_ssa/src/base.rs
index 102d4ea2fa6cf..01ce07f4ccfea 100644
--- a/compiler/rustc_codegen_ssa/src/base.rs
+++ b/compiler/rustc_codegen_ssa/src/base.rs
@@ -647,7 +647,7 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
 ) -> OngoingCodegen<B> {
     // Skip crate items and just output metadata in -Z no-codegen mode.
     if tcx.sess.opts.unstable_opts.no_codegen || !tcx.sess.opts.output_types.should_codegen() {
-        let ongoing_codegen = start_async_codegen(backend, tcx, target_cpu);
+        let ongoing_codegen = start_async_codegen(backend, tcx, target_cpu, &[]);
 
         ongoing_codegen.codegen_finished(tcx);
 
@@ -667,7 +667,6 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
     // codegen units.
     let MonoItemPartitions { codegen_units, autodiff_items, .. } =
         tcx.collect_and_partition_mono_items(());
-    let autodiff_fncs = autodiff_items.to_vec();
 
     // Force all codegen_unit queries so they are already either red or green
     // when compile_codegen_unit accesses them. We are not able to re-execute
@@ -680,7 +679,7 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
         }
     }
 
-    let ongoing_codegen = start_async_codegen(backend.clone(), tcx, target_cpu);
+    let ongoing_codegen = start_async_codegen(backend.clone(), tcx, target_cpu, autodiff_items);
 
     // Codegen an allocator shim, if necessary.
     if let Some(kind) = allocator_kind_for_codegen(tcx) {
@@ -710,10 +709,6 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
         );
     }
 
-    if !autodiff_fncs.is_empty() {
-        ongoing_codegen.submit_autodiff_items(autodiff_fncs);
-    }
-
     // For better throughput during parallel processing by LLVM, we used to sort
     // CGUs largest to smallest. This would lead to better thread utilization
     // by, for example, preventing a large CGU from being processed last and

From 779cb00311dc0bbb942dbdaefa24f98bd2d4cbbe Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Thu, 3 Jul 2025 16:05:38 +0000
Subject: [PATCH 6/9] Move run_fat_lto call into execute_fat_lto_work_item

This will allow merging all fat LTO steps together. In addition it
reduces the amount of work done on the coordinator thread without
jobserver token.
---
 compiler/rustc_codegen_ssa/src/back/write.rs | 68 +++++++++++---------
 1 file changed, 36 insertions(+), 32 deletions(-)

diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
index dcee7aec73392..c7d2fcef29f99 100644
--- a/compiler/rustc_codegen_ssa/src/back/write.rs
+++ b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -397,25 +397,6 @@ impl<B: WriteBackendMethods> CodegenContext<B> {
     }
 }
 
-fn generate_fat_lto_work<B: ExtraBackendMethods>(
-    cgcx: &CodegenContext<B>,
-    autodiff: Vec<AutoDiffItem>,
-    needs_fat_lto: Vec<FatLtoInput<B>>,
-    import_only_modules: Vec<(SerializedModule<B::ModuleBuffer>, WorkProduct)>,
-) -> WorkItem<B> {
-    let _prof_timer = cgcx.prof.generic_activity("codegen_generate_lto_work");
-
-    let module =
-        B::run_fat_lto(cgcx, needs_fat_lto, import_only_modules).unwrap_or_else(|e| e.raise());
-    if cgcx.lto == Lto::Fat && !autodiff.is_empty() {
-        let config = cgcx.config(ModuleKind::Regular);
-        if let Err(err) = B::autodiff(cgcx, &module, autodiff, config) {
-            err.raise();
-        }
-    }
-    WorkItem::FatLto(module)
-}
-
 fn generate_thin_lto_work<B: ExtraBackendMethods>(
     cgcx: &CodegenContext<B>,
     needs_thin_lto: Vec<(String, B::ThinBuffer)>,
@@ -739,7 +720,11 @@ pub(crate) enum WorkItem<B: WriteBackendMethods> {
     /// directory.
     CopyPostLtoArtifacts(CachedModuleCodegen),
     /// Performs fat LTO on the given module.
-    FatLto(ModuleCodegen<B::Module>),
+    FatLto {
+        needs_fat_lto: Vec<FatLtoInput<B>>,
+        import_only_modules: Vec<(SerializedModule<B::ModuleBuffer>, WorkProduct)>,
+        autodiff: Vec<AutoDiffItem>,
+    },
     /// Performs thin-LTO on the given module.
     ThinLto(lto::ThinModule<B>),
 }
@@ -748,7 +733,7 @@ impl<B: WriteBackendMethods> WorkItem<B> {
     fn module_kind(&self) -> ModuleKind {
         match *self {
             WorkItem::Optimize(ref m) => m.kind,
-            WorkItem::CopyPostLtoArtifacts(_) | WorkItem::FatLto(_) | WorkItem::ThinLto(_) => {
+            WorkItem::CopyPostLtoArtifacts(_) | WorkItem::FatLto { .. } | WorkItem::ThinLto(_) => {
                 ModuleKind::Regular
             }
         }
@@ -798,7 +783,7 @@ impl<B: WriteBackendMethods> WorkItem<B> {
         match self {
             WorkItem::Optimize(m) => desc("opt", "optimize module", &m.name),
             WorkItem::CopyPostLtoArtifacts(m) => desc("cpy", "copy LTO artifacts for", &m.name),
-            WorkItem::FatLto(_) => desc("lto", "fat LTO module", "everything"),
+            WorkItem::FatLto { .. } => desc("lto", "fat LTO module", "everything"),
             WorkItem::ThinLto(m) => desc("lto", "thin-LTO module", m.name()),
         }
     }
@@ -1005,9 +990,21 @@ fn execute_copy_from_cache_work_item<B: ExtraBackendMethods>(
 
 fn execute_fat_lto_work_item<B: ExtraBackendMethods>(
     cgcx: &CodegenContext<B>,
-    mut module: ModuleCodegen<B::Module>,
+    needs_fat_lto: Vec<FatLtoInput<B>>,
+    import_only_modules: Vec<(SerializedModule<B::ModuleBuffer>, WorkProduct)>,
+    autodiff: Vec<AutoDiffItem>,
     module_config: &ModuleConfig,
 ) -> Result<WorkItemResult<B>, FatalError> {
+    let mut module =
+        B::run_fat_lto(cgcx, needs_fat_lto, import_only_modules).unwrap_or_else(|e| e.raise());
+
+    if cgcx.lto == Lto::Fat && !autodiff.is_empty() {
+        let config = cgcx.config(ModuleKind::Regular);
+        if let Err(err) = B::autodiff(cgcx, &module, autodiff, config) {
+            err.raise();
+        }
+    }
+
     B::optimize_fat(cgcx, &mut module)?;
 
     let module = B::codegen(cgcx, module, module_config)?;
@@ -1490,13 +1487,14 @@ fn start_executing_work<B: ExtraBackendMethods>(
                     if !needs_fat_lto.is_empty() {
                         assert!(needs_thin_lto.is_empty());
 
-                        let work = generate_fat_lto_work(
-                            &cgcx,
-                            autodiff_items.clone(),
-                            needs_fat_lto,
-                            import_only_modules,
-                        );
-                        work_items.push((work, 0));
+                        work_items.push((
+                            WorkItem::FatLto {
+                                needs_fat_lto,
+                                import_only_modules,
+                                autodiff: autodiff_items.clone(),
+                            },
+                            0,
+                        ));
                         if cgcx.parallel {
                             helper.request_token();
                         }
@@ -1867,11 +1865,17 @@ fn spawn_work<'a, B: ExtraBackendMethods>(
                     );
                     Ok(execute_copy_from_cache_work_item(&cgcx, m, module_config))
                 }
-                WorkItem::FatLto(m) => {
+                WorkItem::FatLto { needs_fat_lto, import_only_modules, autodiff } => {
                     let _timer = cgcx
                         .prof
                         .generic_activity_with_arg("codegen_module_perform_lto", "everything");
-                    execute_fat_lto_work_item(&cgcx, m, module_config)
+                    execute_fat_lto_work_item(
+                        &cgcx,
+                        needs_fat_lto,
+                        import_only_modules,
+                        autodiff,
+                        module_config,
+                    )
                 }
                 WorkItem::ThinLto(m) => {
                     let _timer =

From 9a3aa8fdb192b42e665eb9059a0061f865b3d352 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Thu, 3 Jul 2025 16:07:28 +0000
Subject: [PATCH 7/9] Remove unnecessary check for fat LTO

---
 compiler/rustc_codegen_ssa/src/back/write.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
index c7d2fcef29f99..433cb63882f63 100644
--- a/compiler/rustc_codegen_ssa/src/back/write.rs
+++ b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -998,7 +998,7 @@ fn execute_fat_lto_work_item<B: ExtraBackendMethods>(
     let mut module =
         B::run_fat_lto(cgcx, needs_fat_lto, import_only_modules).unwrap_or_else(|e| e.raise());
 
-    if cgcx.lto == Lto::Fat && !autodiff.is_empty() {
+    if !autodiff.is_empty() {
         let config = cgcx.config(ModuleKind::Regular);
         if let Err(err) = B::autodiff(cgcx, &module, autodiff, config) {
             err.raise();

From 8d63c7a1d68cd30383dceb03a76499c658fa7f40 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Thu, 3 Jul 2025 16:09:10 +0000
Subject: [PATCH 8/9] Remove unused config param from
 WriteBackendMethods::autodiff

---
 compiler/rustc_codegen_gcc/src/lib.rs               | 1 -
 compiler/rustc_codegen_llvm/src/builder/autodiff.rs | 2 --
 compiler/rustc_codegen_llvm/src/lib.rs              | 3 +--
 compiler/rustc_codegen_ssa/src/back/write.rs        | 3 +--
 compiler/rustc_codegen_ssa/src/traits/write.rs      | 1 -
 5 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/compiler/rustc_codegen_gcc/src/lib.rs b/compiler/rustc_codegen_gcc/src/lib.rs
index 34452bdd2005f..8e63ebc849406 100644
--- a/compiler/rustc_codegen_gcc/src/lib.rs
+++ b/compiler/rustc_codegen_gcc/src/lib.rs
@@ -437,7 +437,6 @@ impl WriteBackendMethods for GccCodegenBackend {
         _cgcx: &CodegenContext<Self>,
         _module: &ModuleCodegen<Self::Module>,
         _diff_functions: Vec<AutoDiffItem>,
-        _config: &ModuleConfig,
     ) -> Result<(), FatalError> {
         unimplemented!()
     }
diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs
index b07d9a5cfca8c..c27e161ba9948 100644
--- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs
+++ b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs
@@ -2,7 +2,6 @@ use std::ptr;
 
 use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivity, DiffMode};
 use rustc_codegen_ssa::ModuleCodegen;
-use rustc_codegen_ssa::back::write::ModuleConfig;
 use rustc_codegen_ssa::common::TypeKind;
 use rustc_codegen_ssa::traits::BaseTypeCodegenMethods;
 use rustc_errors::FatalError;
@@ -461,7 +460,6 @@ pub(crate) fn differentiate<'ll>(
     module: &'ll ModuleCodegen<ModuleLlvm>,
     cgcx: &CodegenContext<LlvmCodegenBackend>,
     diff_items: Vec<AutoDiffItem>,
-    _config: &ModuleConfig,
 ) -> Result<(), FatalError> {
     for item in &diff_items {
         trace!("{}", item);
diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs
index 4d85683de8c43..7f54416ec0e31 100644
--- a/compiler/rustc_codegen_llvm/src/lib.rs
+++ b/compiler/rustc_codegen_llvm/src/lib.rs
@@ -231,13 +231,12 @@ impl WriteBackendMethods for LlvmCodegenBackend {
         cgcx: &CodegenContext<Self>,
         module: &ModuleCodegen<Self::Module>,
         diff_fncs: Vec<AutoDiffItem>,
-        config: &ModuleConfig,
     ) -> Result<(), FatalError> {
         if cgcx.lto != Lto::Fat {
             let dcx = cgcx.create_dcx();
             return Err(dcx.handle().emit_almost_fatal(AutoDiffWithoutLTO));
         }
-        builder::autodiff::differentiate(module, cgcx, diff_fncs, config)
+        builder::autodiff::differentiate(module, cgcx, diff_fncs)
     }
 }
 
diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
index 433cb63882f63..24f648bdb2e48 100644
--- a/compiler/rustc_codegen_ssa/src/back/write.rs
+++ b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -999,8 +999,7 @@ fn execute_fat_lto_work_item<B: ExtraBackendMethods>(
         B::run_fat_lto(cgcx, needs_fat_lto, import_only_modules).unwrap_or_else(|e| e.raise());
 
     if !autodiff.is_empty() {
-        let config = cgcx.config(ModuleKind::Regular);
-        if let Err(err) = B::autodiff(cgcx, &module, autodiff, config) {
+        if let Err(err) = B::autodiff(cgcx, &module, autodiff) {
             err.raise();
         }
     }
diff --git a/compiler/rustc_codegen_ssa/src/traits/write.rs b/compiler/rustc_codegen_ssa/src/traits/write.rs
index c732c4b2ab860..949fb818c783a 100644
--- a/compiler/rustc_codegen_ssa/src/traits/write.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/write.rs
@@ -65,7 +65,6 @@ pub trait WriteBackendMethods: Clone + 'static {
         cgcx: &CodegenContext<Self>,
         module: &ModuleCodegen<Self::Module>,
         diff_fncs: Vec<AutoDiffItem>,
-        config: &ModuleConfig,
     ) -> Result<(), FatalError>;
 }
 

From 21026cae8d34241b65924198b72c7231ce4f5a3d Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Thu, 3 Jul 2025 16:22:32 +0000
Subject: [PATCH 9/9] Merge run_fat_lto, optimize_fat and autodiff into
 run_and_optimize_fat_lto

---
 compiler/rustc_codegen_gcc/src/lib.rs         | 23 +++--------
 compiler/rustc_codegen_llvm/messages.ftl      |  1 -
 compiler/rustc_codegen_llvm/src/errors.rs     |  4 --
 compiler/rustc_codegen_llvm/src/lib.rs        | 39 +++++++------------
 compiler/rustc_codegen_ssa/src/back/write.rs  | 12 +-----
 .../rustc_codegen_ssa/src/traits/write.rs     | 16 ++------
 6 files changed, 26 insertions(+), 69 deletions(-)

diff --git a/compiler/rustc_codegen_gcc/src/lib.rs b/compiler/rustc_codegen_gcc/src/lib.rs
index 8e63ebc849406..75c36fffec983 100644
--- a/compiler/rustc_codegen_gcc/src/lib.rs
+++ b/compiler/rustc_codegen_gcc/src/lib.rs
@@ -357,11 +357,16 @@ impl WriteBackendMethods for GccCodegenBackend {
     type ThinData = ThinData;
     type ThinBuffer = ThinBuffer;
 
-    fn run_fat_lto(
+    fn run_and_optimize_fat_lto(
         cgcx: &CodegenContext<Self>,
         modules: Vec<FatLtoInput<Self>>,
         cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
+        diff_fncs: Vec<AutoDiffItem>,
     ) -> Result<ModuleCodegen<Self::Module>, FatalError> {
+        if !diff_fncs.is_empty() {
+            unimplemented!();
+        }
+
         back::lto::run_fat(cgcx, modules, cached_modules)
     }
 
@@ -391,14 +396,6 @@ impl WriteBackendMethods for GccCodegenBackend {
         Ok(())
     }
 
-    fn optimize_fat(
-        _cgcx: &CodegenContext<Self>,
-        _module: &mut ModuleCodegen<Self::Module>,
-    ) -> Result<(), FatalError> {
-        // TODO(antoyo)
-        Ok(())
-    }
-
     fn optimize_thin(
         cgcx: &CodegenContext<Self>,
         thin: ThinModule<Self>,
@@ -432,14 +429,6 @@ impl WriteBackendMethods for GccCodegenBackend {
     ) -> Result<ModuleCodegen<Self::Module>, FatalError> {
         back::write::link(cgcx, dcx, modules)
     }
-
-    fn autodiff(
-        _cgcx: &CodegenContext<Self>,
-        _module: &ModuleCodegen<Self::Module>,
-        _diff_functions: Vec<AutoDiffItem>,
-    ) -> Result<(), FatalError> {
-        unimplemented!()
-    }
 }
 
 /// This is the entrypoint for a hot plugged rustc_codegen_gccjit
diff --git a/compiler/rustc_codegen_llvm/messages.ftl b/compiler/rustc_codegen_llvm/messages.ftl
index 3885f18271f1b..f197ea7447326 100644
--- a/compiler/rustc_codegen_llvm/messages.ftl
+++ b/compiler/rustc_codegen_llvm/messages.ftl
@@ -1,5 +1,4 @@
 codegen_llvm_autodiff_without_enable = using the autodiff feature requires -Z autodiff=Enable
-codegen_llvm_autodiff_without_lto = using the autodiff feature requires using fat-lto
 
 codegen_llvm_copy_bitcode = failed to copy bitcode to object file: {$err}
 
diff --git a/compiler/rustc_codegen_llvm/src/errors.rs b/compiler/rustc_codegen_llvm/src/errors.rs
index d50ad8a1a9cb4..31d49e8631955 100644
--- a/compiler/rustc_codegen_llvm/src/errors.rs
+++ b/compiler/rustc_codegen_llvm/src/errors.rs
@@ -37,10 +37,6 @@ impl<G: EmissionGuarantee> Diagnostic<'_, G> for ParseTargetMachineConfig<'_> {
     }
 }
 
-#[derive(Diagnostic)]
-#[diag(codegen_llvm_autodiff_without_lto)]
-pub(crate) struct AutoDiffWithoutLTO;
-
 #[derive(Diagnostic)]
 #[diag(codegen_llvm_autodiff_without_enable)]
 pub(crate) struct AutoDiffWithoutEnable;
diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs
index 7f54416ec0e31..5d551c3af87fb 100644
--- a/compiler/rustc_codegen_llvm/src/lib.rs
+++ b/compiler/rustc_codegen_llvm/src/lib.rs
@@ -26,7 +26,7 @@ use std::mem::ManuallyDrop;
 use back::owned_target_machine::OwnedTargetMachine;
 use back::write::{create_informational_target_machine, create_target_machine};
 use context::SimpleCx;
-use errors::{AutoDiffWithoutLTO, ParseTargetMachineConfig};
+use errors::ParseTargetMachineConfig;
 use llvm_util::target_config;
 use rustc_ast::expand::allocator::AllocatorKind;
 use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
@@ -43,7 +43,7 @@ use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
 use rustc_middle::ty::TyCtxt;
 use rustc_middle::util::Providers;
 use rustc_session::Session;
-use rustc_session::config::{Lto, OptLevel, OutputFilenames, PrintKind, PrintRequest};
+use rustc_session::config::{OptLevel, OutputFilenames, PrintKind, PrintRequest};
 use rustc_span::Symbol;
 
 mod back {
@@ -174,12 +174,23 @@ impl WriteBackendMethods for LlvmCodegenBackend {
     ) -> Result<ModuleCodegen<Self::Module>, FatalError> {
         back::write::link(cgcx, dcx, modules)
     }
-    fn run_fat_lto(
+    fn run_and_optimize_fat_lto(
         cgcx: &CodegenContext<Self>,
         modules: Vec<FatLtoInput<Self>>,
         cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
+        diff_fncs: Vec<AutoDiffItem>,
     ) -> Result<ModuleCodegen<Self::Module>, FatalError> {
-        back::lto::run_fat(cgcx, modules, cached_modules)
+        let mut module = back::lto::run_fat(cgcx, modules, cached_modules)?;
+
+        if !diff_fncs.is_empty() {
+            builder::autodiff::differentiate(&module, cgcx, diff_fncs)?;
+        }
+
+        let dcx = cgcx.create_dcx();
+        let dcx = dcx.handle();
+        back::lto::run_pass_manager(cgcx, dcx, &mut module, false)?;
+
+        Ok(module)
     }
     fn run_thin_lto(
         cgcx: &CodegenContext<Self>,
@@ -196,14 +207,6 @@ impl WriteBackendMethods for LlvmCodegenBackend {
     ) -> Result<(), FatalError> {
         back::write::optimize(cgcx, dcx, module, config)
     }
-    fn optimize_fat(
-        cgcx: &CodegenContext<Self>,
-        module: &mut ModuleCodegen<Self::Module>,
-    ) -> Result<(), FatalError> {
-        let dcx = cgcx.create_dcx();
-        let dcx = dcx.handle();
-        back::lto::run_pass_manager(cgcx, dcx, module, false)
-    }
     fn optimize_thin(
         cgcx: &CodegenContext<Self>,
         thin: ThinModule<Self>,
@@ -226,18 +229,6 @@ impl WriteBackendMethods for LlvmCodegenBackend {
     fn serialize_module(module: ModuleCodegen<Self::Module>) -> (String, Self::ModuleBuffer) {
         (module.name, back::lto::ModuleBuffer::new(module.module_llvm.llmod()))
     }
-    /// Generate autodiff rules
-    fn autodiff(
-        cgcx: &CodegenContext<Self>,
-        module: &ModuleCodegen<Self::Module>,
-        diff_fncs: Vec<AutoDiffItem>,
-    ) -> Result<(), FatalError> {
-        if cgcx.lto != Lto::Fat {
-            let dcx = cgcx.create_dcx();
-            return Err(dcx.handle().emit_almost_fatal(AutoDiffWithoutLTO));
-        }
-        builder::autodiff::differentiate(module, cgcx, diff_fncs)
-    }
 }
 
 impl LlvmCodegenBackend {
diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
index 24f648bdb2e48..26fc593da98a7 100644
--- a/compiler/rustc_codegen_ssa/src/back/write.rs
+++ b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -995,17 +995,7 @@ fn execute_fat_lto_work_item<B: ExtraBackendMethods>(
     autodiff: Vec<AutoDiffItem>,
     module_config: &ModuleConfig,
 ) -> Result<WorkItemResult<B>, FatalError> {
-    let mut module =
-        B::run_fat_lto(cgcx, needs_fat_lto, import_only_modules).unwrap_or_else(|e| e.raise());
-
-    if !autodiff.is_empty() {
-        if let Err(err) = B::autodiff(cgcx, &module, autodiff) {
-            err.raise();
-        }
-    }
-
-    B::optimize_fat(cgcx, &mut module)?;
-
+    let module = B::run_and_optimize_fat_lto(cgcx, needs_fat_lto, import_only_modules, autodiff)?;
     let module = B::codegen(cgcx, module, module_config)?;
     Ok(WorkItemResult::Finished(module))
 }
diff --git a/compiler/rustc_codegen_ssa/src/traits/write.rs b/compiler/rustc_codegen_ssa/src/traits/write.rs
index 949fb818c783a..5e993640472d1 100644
--- a/compiler/rustc_codegen_ssa/src/traits/write.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/write.rs
@@ -20,12 +20,13 @@ pub trait WriteBackendMethods: Clone + 'static {
         dcx: DiagCtxtHandle<'_>,
         modules: Vec<ModuleCodegen<Self::Module>>,
     ) -> Result<ModuleCodegen<Self::Module>, FatalError>;
-    /// Performs fat LTO by merging all modules into a single one and returning it
-    /// for further optimization.
-    fn run_fat_lto(
+    /// Performs fat LTO by merging all modules into a single one, running autodiff
+    /// if necessary and running any further optimizations
+    fn run_and_optimize_fat_lto(
         cgcx: &CodegenContext<Self>,
         modules: Vec<FatLtoInput<Self>>,
         cached_modules: Vec<(SerializedModule<Self::ModuleBuffer>, WorkProduct)>,
+        diff_fncs: Vec<AutoDiffItem>,
     ) -> Result<ModuleCodegen<Self::Module>, FatalError>;
     /// Performs thin LTO by performing necessary global analysis and returning two
     /// lists, one of the modules that need optimization and another for modules that
@@ -43,10 +44,6 @@ pub trait WriteBackendMethods: Clone + 'static {
         module: &mut ModuleCodegen<Self::Module>,
         config: &ModuleConfig,
     ) -> Result<(), FatalError>;
-    fn optimize_fat(
-        cgcx: &CodegenContext<Self>,
-        llmod: &mut ModuleCodegen<Self::Module>,
-    ) -> Result<(), FatalError>;
     fn optimize_thin(
         cgcx: &CodegenContext<Self>,
         thin: ThinModule<Self>,
@@ -61,11 +58,6 @@ pub trait WriteBackendMethods: Clone + 'static {
         want_summary: bool,
     ) -> (String, Self::ThinBuffer);
     fn serialize_module(module: ModuleCodegen<Self::Module>) -> (String, Self::ModuleBuffer);
-    fn autodiff(
-        cgcx: &CodegenContext<Self>,
-        module: &ModuleCodegen<Self::Module>,
-        diff_fncs: Vec<AutoDiffItem>,
-    ) -> Result<(), FatalError>;
 }
 
 pub trait ThinBufferMethods: Send + Sync {