From b28d354eaf062d0469bef417634fc93b5eb2b336 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Mon, 24 Feb 2025 18:44:15 -0500 Subject: [PATCH 1/4] don't ignore autodiff tasks for lib builds --- compiler/rustc_codegen_ssa/src/back/write.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs index d2548deb8c7a4..4c763df6a1fe6 100644 --- a/compiler/rustc_codegen_ssa/src/back/write.rs +++ b/compiler/rustc_codegen_ssa/src/back/write.rs @@ -1459,6 +1459,7 @@ fn start_executing_work( if needs_fat_lto.is_empty() && needs_thin_lto.is_empty() && lto_import_only_modules.is_empty() + && autodiff_items.is_empty() { // Nothing more to do! break; From 0597d21d315953299d5bff7e656ccc3f509da05a Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Mon, 24 Feb 2025 22:11:39 -0500 Subject: [PATCH 2/4] small code quality improvements --- compiler/rustc_codegen_llvm/src/builder/autodiff.rs | 5 ++--- compiler/rustc_codegen_llvm/src/lib.rs | 3 ++- compiler/rustc_codegen_ssa/src/back/write.rs | 3 ++- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs index 2c7899975e3ee..2118380cad462 100644 --- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs +++ b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs @@ -1,7 +1,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_errors::FatalError; use tracing::{debug, trace}; @@ -276,7 +275,7 @@ fn generate_enzyme_call<'ll>( } pub(crate) fn differentiate<'ll>( - module: &'ll ModuleCodegen, + module_llvm: &'ll ModuleLlvm, cgcx: &CodegenContext, diff_items: Vec, _config: &ModuleConfig, @@ -286,7 +285,7 @@ pub(crate) fn differentiate<'ll>( } let diag_handler = cgcx.create_dcx(); - let cx = SimpleCx { llmod: module.module_llvm.llmod(), llcx: module.module_llvm.llcx }; + let cx = SimpleCx { llmod: module_llvm.llmod(), llcx: module_llvm.llcx }; // First of all, did the user try to use autodiff without using the -Zautodiff=Enable flag? if !diff_items.is_empty() diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs index e9e1b644f183e..bfb8675f3e66c 100644 --- a/compiler/rustc_codegen_llvm/src/lib.rs +++ b/compiler/rustc_codegen_llvm/src/lib.rs @@ -241,7 +241,8 @@ impl WriteBackendMethods for LlvmCodegenBackend { let dcx = cgcx.create_dcx(); return Err(dcx.handle().emit_almost_fatal(AutoDiffWithoutLTO)); } - builder::autodiff::differentiate(module, cgcx, diff_fncs, config) + let module_llvm = &module.module_llvm; + builder::autodiff::differentiate(module_llvm, cgcx, diff_fncs, config) } } diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs index 4c763df6a1fe6..91b952728f7a0 100644 --- a/compiler/rustc_codegen_ssa/src/back/write.rs +++ b/compiler/rustc_codegen_ssa/src/back/write.rs @@ -1473,13 +1473,14 @@ fn start_executing_work( assert!(!started_lto); started_lto = true; + let autodiff_items = mem::take(&mut autodiff_items); let needs_fat_lto = mem::take(&mut needs_fat_lto); 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(), + autodiff_items, needs_fat_lto, needs_thin_lto, import_only_modules, From f9f5959fa45cd743b0f51950eadd8caef61a024c Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Mon, 24 Feb 2025 22:14:19 -0500 Subject: [PATCH 3/4] more iffy improvements --- compiler/rustc_codegen_llvm/src/lib.rs | 28 ++++++++++++++++--- compiler/rustc_codegen_ssa/src/back/lto.rs | 5 ++-- compiler/rustc_codegen_ssa/src/back/write.rs | 7 +++-- .../rustc_codegen_ssa/src/traits/write.rs | 6 ++++ 4 files changed, 38 insertions(+), 8 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs index bfb8675f3e66c..ac74e429b10fe 100644 --- a/compiler/rustc_codegen_llvm/src/lib.rs +++ b/compiler/rustc_codegen_llvm/src/lib.rs @@ -237,13 +237,33 @@ impl WriteBackendMethods for LlvmCodegenBackend { diff_fncs: Vec, config: &ModuleConfig, ) -> Result<(), FatalError> { - if cgcx.lto != Lto::Fat { - let dcx = cgcx.create_dcx(); - return Err(dcx.handle().emit_almost_fatal(AutoDiffWithoutLTO)); - } + //if cgcx.lto != Lto::Fat { + // let dcx = cgcx.create_dcx(); + // return Err(dcx.handle().emit_almost_fatal(AutoDiffWithoutLTO)); + //} let module_llvm = &module.module_llvm; builder::autodiff::differentiate(module_llvm, cgcx, diff_fncs, config) } + fn autodiff_thin( + cgcx: &CodegenContext, + thin_module: &ThinModule, + diff_fncs: Vec, + config: &ModuleConfig, + ) -> Result<(), FatalError> { + let dcx = cgcx.create_dcx(); + let dcx = dcx.handle(); + + let module_name = &thin_module.shared.module_names[thin_module.idx]; + + // Right now the implementation we've got only works over serialized + // modules, so we create a fresh new LLVM context and parse the module + // into that context. One day, however, we may do this for upstream + // crates but for locally codegened modules we may be able to reuse + // that LLVM Context and Module. + let module_llvm = ModuleLlvm::parse(cgcx, module_name, thin_module.data(), dcx)?; + + builder::autodiff::differentiate(&module_llvm, cgcx, diff_fncs, config) + } } impl LlvmCodegenBackend { diff --git a/compiler/rustc_codegen_ssa/src/back/lto.rs b/compiler/rustc_codegen_ssa/src/back/lto.rs index 9fd984b6419ee..97fe21ad4d552 100644 --- a/compiler/rustc_codegen_ssa/src/back/lto.rs +++ b/compiler/rustc_codegen_ssa/src/back/lto.rs @@ -94,10 +94,11 @@ impl LtoModuleCodegen { match &self { LtoModuleCodegen::Fat(module) => { B::autodiff(cgcx, &module, diff_fncs, config)?; + }, + LtoModuleCodegen::Thin(module) => { + B::autodiff_thin(cgcx, module, diff_fncs, config)?; } - _ => panic!("autodiff called with non-fat LTO module"), } - Ok(self) } } diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs index 91b952728f7a0..71cc13279ec2b 100644 --- a/compiler/rustc_codegen_ssa/src/back/write.rs +++ b/compiler/rustc_codegen_ssa/src/back/write.rs @@ -412,15 +412,18 @@ fn generate_lto_work( vec![(WorkItem::LTO(module), 0)] } else { if !autodiff.is_empty() { - let dcx = cgcx.create_dcx(); - dcx.handle().emit_fatal(AutodiffWithoutLto {}); + //let dcx = cgcx.create_dcx(); + //dcx.handle().emit_fatal(AutodiffWithoutLto {}); } + let config = cgcx.config(ModuleKind::Regular); 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 mut module = + unsafe { module.autodiff(cgcx, autodiff.clone(), config).unwrap_or_else(|e| e.raise()) }; let cost = module.cost(); (WorkItem::LTO(module), cost) }) diff --git a/compiler/rustc_codegen_ssa/src/traits/write.rs b/compiler/rustc_codegen_ssa/src/traits/write.rs index 97fe614aa10cd..ac916e90ab990 100644 --- a/compiler/rustc_codegen_ssa/src/traits/write.rs +++ b/compiler/rustc_codegen_ssa/src/traits/write.rs @@ -68,6 +68,12 @@ pub trait WriteBackendMethods: 'static + Sized + Clone { diff_fncs: Vec, config: &ModuleConfig, ) -> Result<(), FatalError>; + fn autodiff_thin( + cgcx: &CodegenContext, + thin: &ThinModule, + diff_fncs: Vec, + config: &ModuleConfig, + ) -> Result<(), FatalError>; } pub trait ThinBufferMethods: Send + Sync { From ca85a3e4fcac6cfcb1c3b09ab93dbbc70db021aa Mon Sep 17 00:00:00 2001 From: vayunbiyani <61186953+vayunbiyani@users.noreply.github.com> Date: Sat, 15 Mar 2025 16:48:36 +0530 Subject: [PATCH 4/4] Fix build errors --- compiler/rustc_codegen_llvm/src/errors.rs | 6 +++--- compiler/rustc_codegen_llvm/src/lib.rs | 4 ++-- compiler/rustc_codegen_ssa/src/back/write.rs | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/compiler/rustc_codegen_llvm/src/errors.rs b/compiler/rustc_codegen_llvm/src/errors.rs index 4c5a78ca74fe4..55870e386035a 100644 --- a/compiler/rustc_codegen_llvm/src/errors.rs +++ b/compiler/rustc_codegen_llvm/src/errors.rs @@ -90,9 +90,9 @@ impl Diagnostic<'_, G> for ParseTargetMachineConfig<'_> { } } -#[derive(Diagnostic)] -#[diag(codegen_llvm_autodiff_without_lto)] -pub(crate) struct AutoDiffWithoutLTO; +// #[derive(Diagnostic)] +// #[diag(codegen_llvm_autodiff_without_lto)] +// pub(crate) struct AutoDiffWithoutLTO; #[derive(Diagnostic)] #[diag(codegen_llvm_autodiff_without_enable)] diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs index ac74e429b10fe..fc0ae09c466ac 100644 --- a/compiler/rustc_codegen_llvm/src/lib.rs +++ b/compiler/rustc_codegen_llvm/src/lib.rs @@ -29,7 +29,7 @@ use std::mem::ManuallyDrop; use back::owned_target_machine::OwnedTargetMachine; use back::write::{create_informational_target_machine, create_target_machine}; -use errors::{AutoDiffWithoutLTO, ParseTargetMachineConfig}; +use errors::ParseTargetMachineConfig; pub(crate) use llvm_util::target_features_cfg; use rustc_ast::expand::allocator::AllocatorKind; use rustc_ast::expand::autodiff_attrs::AutoDiffItem; @@ -46,7 +46,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 { diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs index 71cc13279ec2b..8d238b5ca3619 100644 --- a/compiler/rustc_codegen_ssa/src/back/write.rs +++ b/compiler/rustc_codegen_ssa/src/back/write.rs @@ -41,7 +41,7 @@ use tracing::debug; use super::link::{self, ensure_removed}; use super::lto::{self, SerializedModule}; use super::symbol_export::symbol_name_for_instance_in_crate; -use crate::errors::{AutodiffWithoutLto, ErrorCreatingRemarkDir}; +use crate::errors::ErrorCreatingRemarkDir; use crate::traits::*; use crate::{ CachedModuleCodegen, CodegenResults, CompiledModule, CrateInfo, ModuleCodegen, ModuleKind, @@ -422,7 +422,7 @@ fn generate_lto_work( lto_modules .into_iter() .map(|module| { - let mut module = + let module = unsafe { module.autodiff(cgcx, autodiff.clone(), config).unwrap_or_else(|e| e.raise()) }; let cost = module.cost(); (WorkItem::LTO(module), cost)