1
Fork 0

upstream rustc_codegen_llvm changes for enzyme/autodiff

This commit is contained in:
Manuel Drehwald 2025-01-01 21:42:45 +01:00
parent 372442fe5f
commit d753cbf779
17 changed files with 610 additions and 28 deletions

View file

@ -28,9 +28,10 @@ use std::mem::ManuallyDrop;
use back::owned_target_machine::OwnedTargetMachine;
use back::write::{create_informational_target_machine, create_target_machine};
use errors::ParseTargetMachineConfig;
use errors::{AutoDiffWithoutLTO, ParseTargetMachineConfig};
pub use llvm_util::target_features_cfg;
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::write::{
CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryConfig, TargetMachineFactoryFn,
@ -44,7 +45,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::{OptLevel, OutputFilenames, PrintKind, PrintRequest};
use rustc_session::config::{Lto, OptLevel, OutputFilenames, PrintKind, PrintRequest};
use rustc_span::Symbol;
mod back {
@ -233,6 +234,20 @@ 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>,
tcx: TyCtxt<'_>,
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, tcx, diff_fncs, config)
}
}
unsafe impl Send for LlvmCodegenBackend {} // Llvm is on a per-thread basis