bjoernager/rust
bjoernager/rust
1
Fork 0
Code Activity

upstream rustc_codegen_llvm changes for enzyme/autodiff

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

21
compiler/rustc_codegen_ssa/src/back/lto.rs
View file

@ -1,11 +1,14 @@
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 rustc_middle::ty::TyCtxt;
use super::write::CodegenContext;
use crate::ModuleCodegen;
use crate::back::write::ModuleConfig;
use crate::traits::*;
pub struct ThinModule<B: WriteBackendMethods> {
@ -81,6 +84,24 @@ impl<B: WriteBackendMethods> LtoModuleCodegen<B> {
LtoModuleCodegen::Thin(ref m) => m.cost(),
}
}
/// Run autodiff on Fat LTO module
pub unsafe fn autodiff(
self,
cgcx: &CodegenContext<B>,
tcx: TyCtxt<'_>,
diff_fncs: Vec<AutoDiffItem>,
config: &ModuleConfig,
) -> Result<LtoModuleCodegen<B>, FatalError> {
match &self {
LtoModuleCodegen::Fat(module) => {
B::autodiff(cgcx, tcx, &module, diff_fncs, config)?;
}
_ => panic!("autodiff called with non-fat LTO module"),
}
Ok(self)
}
}
pub enum SerializedModule<M: ModuleBufferMethods> {

9
compiler/rustc_codegen_ssa/src/traits/write.rs
View file

@ -1,5 +1,7 @@
use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
use rustc_errors::{DiagCtxtHandle, FatalError};
use rustc_middle::dep_graph::WorkProduct;
use rustc_middle::ty::TyCtxt;
use crate::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};
use crate::back::write::{CodegenContext, FatLtoInput, ModuleConfig};
@ -61,6 +63,13 @@ pub trait WriteBackendMethods: 'static + Sized + Clone {
want_summary: bool,
) -> (String, Self::ThinBuffer);
fn serialize_module(module: ModuleCodegen<Self::Module>) -> (String, Self::ModuleBuffer);
fn autodiff(
cgcx: &CodegenContext<Self>,
tcx: TyCtxt<'_>,
module: &ModuleCodegen<Self::Module>,
diff_fncs: Vec<AutoDiffItem>,
config: &ModuleConfig,
) -> Result<(), FatalError>;
}
pub trait ThinBufferMethods: Send + Sync {