Replace LLVMRustThinLTOBuffer with separate LLVMRustBuffers for bitcode and summary

This commit is contained in:
bjorn3
2026-02-12 16:28:55 +00:00
parent 8b2c10ff82
commit a5372d1dba
4 changed files with 44 additions and 99 deletions
+14 -33
View File
@@ -634,7 +634,9 @@ pub(crate) fn run_pass_manager(
};
unsafe {
write::llvm_optimize(cgcx, prof, dcx, module, None, config, opt_level, opt_stage, stage);
write::llvm_optimize(
cgcx, prof, dcx, module, None, None, config, opt_level, opt_stage, stage,
);
}
if cfg!(feature = "llvm_enzyme") && enable_ad && !thin {
@@ -643,7 +645,7 @@ pub(crate) fn run_pass_manager(
if !config.autodiff.contains(&config::AutoDiff::NoPostopt) {
unsafe {
write::llvm_optimize(
cgcx, prof, dcx, module, None, config, opt_level, opt_stage, stage,
cgcx, prof, dcx, module, None, None, config, opt_level, opt_stage, stage,
);
}
}
@@ -664,6 +666,11 @@ unsafe impl Send for Buffer {}
unsafe impl Sync for Buffer {}
impl Buffer {
pub(crate) unsafe fn from_raw_ptr(ptr: *mut llvm::Buffer) -> Buffer {
let mut ptr = NonNull::new(ptr).unwrap();
Buffer(unsafe { ptr.as_mut() })
}
pub(crate) fn data(&self) -> &[u8] {
unsafe {
let ptr = llvm::LLVMRustBufferPtr(self.0);
@@ -708,48 +715,22 @@ fn drop(&mut self) {
}
}
pub struct ThinBuffer(&'static mut llvm::ThinLTOBuffer);
unsafe impl Send for ThinBuffer {}
unsafe impl Sync for ThinBuffer {}
pub struct ThinBuffer {
data: Buffer,
}
impl ThinBuffer {
pub(crate) fn new(m: &llvm::Module, is_thin: bool) -> ThinBuffer {
unsafe {
let buffer = llvm::LLVMRustThinLTOBufferCreate(m, is_thin);
ThinBuffer(buffer)
}
}
pub(crate) unsafe fn from_raw_ptr(ptr: *mut llvm::ThinLTOBuffer) -> ThinBuffer {
let mut ptr = NonNull::new(ptr).unwrap();
ThinBuffer(unsafe { ptr.as_mut() })
}
pub(crate) fn thin_link_data(&self) -> &[u8] {
unsafe {
let ptr = llvm::LLVMRustThinLTOBufferThinLinkDataPtr(self.0) as *const _;
let len = llvm::LLVMRustThinLTOBufferThinLinkDataLen(self.0);
slice::from_raw_parts(ptr, len)
ThinBuffer { data: Buffer(buffer) }
}
}
}
impl ThinBufferMethods for ThinBuffer {
fn data(&self) -> &[u8] {
unsafe {
let ptr = llvm::LLVMRustThinLTOBufferPtr(self.0) as *const _;
let len = llvm::LLVMRustThinLTOBufferLen(self.0);
slice::from_raw_parts(ptr, len)
}
}
}
impl Drop for ThinBuffer {
fn drop(&mut self) {
unsafe {
llvm::LLVMRustThinLTOBufferFree(&mut *(self.0 as *mut _));
}
self.data.data()
}
}
+13 -6
View File
@@ -563,7 +563,8 @@ pub(crate) unsafe fn llvm_optimize(
prof: &SelfProfilerRef,
dcx: DiagCtxtHandle<'_>,
module: &ModuleCodegen<ModuleLlvm>,
thin_lto_buffer: Option<&mut *mut llvm::ThinLTOBuffer>,
thin_lto_buffer: Option<&mut *mut llvm::Buffer>,
thin_lto_summary_buffer: Option<&mut *mut llvm::Buffer>,
config: &ModuleConfig,
opt_level: config::OptLevel,
opt_stage: llvm::OptStage,
@@ -786,6 +787,7 @@ fn handle_offload<'ll>(cx: &'ll SimpleCx<'_>, old_fn: &llvm::Value) {
config.verify_llvm_ir,
config.lint_llvm_ir,
thin_lto_buffer,
thin_lto_summary_buffer,
config.emit_thin_lto_summary,
merge_functions,
unroll_loops,
@@ -932,13 +934,14 @@ pub(crate) fn optimize(
// The bitcode obtained during the `codegen` phase is no longer suitable for performing LTO.
// It may have undergone LTO due to ThinLocal, so we need to obtain the embedded bitcode at
// this point.
let mut thin_lto_buffer = if (module.kind == ModuleKind::Regular
let (mut thin_lto_buffer, mut thin_lto_summary_buffer) = if (module.kind
== ModuleKind::Regular
&& config.emit_obj == EmitObj::ObjectCode(BitcodeSection::Full))
|| config.emit_thin_lto_summary
{
Some(null_mut())
(Some(null_mut()), Some(null_mut()))
} else {
None
(None, None)
};
unsafe {
llvm_optimize(
@@ -947,6 +950,7 @@ pub(crate) fn optimize(
dcx,
module,
thin_lto_buffer.as_mut(),
thin_lto_summary_buffer.as_mut(),
config,
opt_level,
opt_stage,
@@ -954,7 +958,10 @@ pub(crate) fn optimize(
)
};
if let Some(thin_lto_buffer) = thin_lto_buffer {
let thin_lto_buffer = unsafe { ThinBuffer::from_raw_ptr(thin_lto_buffer) };
let thin_lto_buffer =
unsafe { crate::back::lto::Buffer::from_raw_ptr(thin_lto_buffer) };
let thin_lto_summary_buffer =
unsafe { crate::back::lto::Buffer::from_raw_ptr(thin_lto_summary_buffer.unwrap()) };
module.thin_lto_buffer = Some(thin_lto_buffer.data().to_vec());
let bc_summary_out = cgcx.output_filenames.temp_path_for_cgu(
OutputType::ThinLinkBitcode,
@@ -964,7 +971,7 @@ pub(crate) fn optimize(
if config.emit_thin_lto_summary
&& let Some(thin_link_bitcode_filename) = bc_summary_out.file_name()
{
let summary_data = thin_lto_buffer.thin_link_data();
let summary_data = thin_lto_summary_buffer.data();
prof.artifact_size(
"llvm_bitcode_summary",
thin_link_bitcode_filename.to_string_lossy(),
+3 -13
View File
@@ -536,9 +536,6 @@ pub(crate) enum DiagnosticLevel {
unsafe extern "C" {
// LLVMRustThinLTOData
pub(crate) type ThinLTOData;
// LLVMRustThinLTOBuffer
pub(crate) type ThinLTOBuffer;
}
/// LLVMRustThinLTOModule
@@ -2375,7 +2372,8 @@ pub(crate) fn LLVMRustOptimize<'a>(
NoPrepopulatePasses: bool,
VerifyIR: bool,
LintIR: bool,
ThinLTOBuffer: Option<&mut *mut ThinLTOBuffer>,
ThinLTOBuffer: Option<&mut *mut Buffer>,
ThinLTOSummaryBuffer: Option<&mut *mut Buffer>,
EmitThinLTOSummary: bool,
MergeFunctions: bool,
UnrollLoops: bool,
@@ -2464,15 +2462,7 @@ pub(crate) fn LLVMRustUnpackSMDiagnostic(
pub(crate) fn LLVMRustModuleCost(M: &Module) -> u64;
pub(crate) fn LLVMRustModuleInstructionStats(M: &Module) -> u64;
pub(crate) fn LLVMRustThinLTOBufferCreate(
M: &Module,
is_thin: bool,
) -> &'static mut ThinLTOBuffer;
pub(crate) fn LLVMRustThinLTOBufferFree(M: &'static mut ThinLTOBuffer);
pub(crate) fn LLVMRustThinLTOBufferPtr(M: &ThinLTOBuffer) -> *const c_char;
pub(crate) fn LLVMRustThinLTOBufferLen(M: &ThinLTOBuffer) -> size_t;
pub(crate) fn LLVMRustThinLTOBufferThinLinkDataPtr(M: &ThinLTOBuffer) -> *const c_char;
pub(crate) fn LLVMRustThinLTOBufferThinLinkDataLen(M: &ThinLTOBuffer) -> size_t;
pub(crate) fn LLVMRustThinLTOBufferCreate(M: &Module, is_thin: bool) -> &'static mut Buffer;
pub(crate) fn LLVMRustCreateThinLTOData(
Modules: *const ThinLTOModule,
NumModules: size_t,
@@ -87,19 +87,6 @@ extern "C" void LLVMRustTimeTraceProfilerFinish(const char *FileName) {
timeTraceProfilerCleanup();
}
// This struct and various functions are sort of a hack right now, but the
// problem is that we've got in-memory LLVM modules after we generate and
// optimize all codegen-units for one compilation in rustc. To be compatible
// with the LTO support above we need to serialize the modules plus their
// ThinLTO summary into memory.
//
// This structure is basically an owned version of a serialize module, with
// a ThinLTO summary attached.
struct LLVMRustThinLTOBuffer {
std::string data;
std::string thin_link_data;
};
extern "C" bool LLVMRustHasFeature(LLVMTargetMachineRef TM,
const char *Feature) {
TargetMachine *Target = unwrap(TM);
@@ -566,11 +553,12 @@ extern "C" LLVMRustResult LLVMRustOptimize(
LLVMModuleRef ModuleRef, LLVMTargetMachineRef TMRef,
LLVMRustPassBuilderOptLevel OptLevelRust, LLVMRustOptStage OptStage,
bool IsLinkerPluginLTO, bool NoPrepopulatePasses, bool VerifyIR,
bool LintIR, LLVMRustThinLTOBuffer **ThinLTOBufferRef,
bool EmitThinLTOSummary, bool MergeFunctions, bool UnrollLoops,
bool SLPVectorize, bool LoopVectorize, bool DisableSimplifyLibCalls,
bool EmitLifetimeMarkers, registerEnzymeAndPassPipelineFn EnzymePtr,
bool PrintBeforeEnzyme, bool PrintAfterEnzyme, bool PrintPasses,
bool LintIR, LLVMRustBuffer **ThinLTOBufferRef,
LLVMRustBuffer **ThinLTOSummaryBufferRef, bool EmitThinLTOSummary,
bool MergeFunctions, bool UnrollLoops, bool SLPVectorize,
bool LoopVectorize, bool DisableSimplifyLibCalls, bool EmitLifetimeMarkers,
registerEnzymeAndPassPipelineFn EnzymePtr, bool PrintBeforeEnzyme,
bool PrintAfterEnzyme, bool PrintPasses,
LLVMRustSanitizerOptions *SanitizerOptions, const char *PGOGenPath,
const char *PGOUsePath, bool InstrumentCoverage,
const char *InstrProfileOutput, const char *PGOSampleUsePath,
@@ -809,9 +797,10 @@ extern "C" LLVMRustResult LLVMRustOptimize(
ModulePassManager MPM;
bool NeedThinLTOBufferPasses = true;
auto ThinLTOBuffer = std::make_unique<LLVMRustThinLTOBuffer>();
auto ThinLTOBuffer = std::make_unique<LLVMRustBuffer>();
auto ThinLTOSummaryBuffer = std::make_unique<LLVMRustBuffer>();
raw_string_ostream ThinLTODataOS(ThinLTOBuffer->data);
raw_string_ostream ThinLinkDataOS(ThinLTOBuffer->thin_link_data);
raw_string_ostream ThinLinkDataOS(ThinLTOSummaryBuffer->data);
bool IsLTO = OptStage == LLVMRustOptStage::ThinLTO ||
OptStage == LLVMRustOptStage::FatLTO;
if (!NoPrepopulatePasses) {
@@ -838,6 +827,7 @@ extern "C" LLVMRustResult LLVMRustOptimize(
MPM.addPass(ThinLTOBitcodeWriterPass(
ThinLTODataOS, EmitThinLTOSummary ? &ThinLinkDataOS : nullptr));
*ThinLTOBufferRef = ThinLTOBuffer.release();
*ThinLTOSummaryBufferRef = ThinLTOSummaryBuffer.release();
MPM.addPass(PB.buildModuleOptimizationPipeline(
OptLevel, ThinOrFullLTOPhase::None));
MPM.addPass(
@@ -898,6 +888,7 @@ extern "C" LLVMRustResult LLVMRustOptimize(
MPM.addPass(BitcodeWriterPass(ThinLTODataOS));
}
*ThinLTOBufferRef = ThinLTOBuffer.release();
*ThinLTOSummaryBufferRef = ThinLTOSummaryBuffer.release();
}
// now load "-enzyme" pass:
@@ -1408,9 +1399,9 @@ extern "C" bool LLVMRustPrepareThinLTOImport(const LLVMRustThinLTOData *Data,
return true;
}
extern "C" LLVMRustThinLTOBuffer *LLVMRustThinLTOBufferCreate(LLVMModuleRef M,
bool is_thin) {
auto Ret = std::make_unique<LLVMRustThinLTOBuffer>();
extern "C" LLVMRustBuffer *LLVMRustThinLTOBufferCreate(LLVMModuleRef M,
bool is_thin) {
auto Ret = std::make_unique<LLVMRustBuffer>();
{
auto OS = raw_string_ostream(Ret->data);
{
@@ -1436,30 +1427,6 @@ extern "C" LLVMRustThinLTOBuffer *LLVMRustThinLTOBufferCreate(LLVMModuleRef M,
return Ret.release();
}
extern "C" void LLVMRustThinLTOBufferFree(LLVMRustThinLTOBuffer *Buffer) {
delete Buffer;
}
extern "C" const void *
LLVMRustThinLTOBufferPtr(const LLVMRustThinLTOBuffer *Buffer) {
return Buffer->data.data();
}
extern "C" size_t
LLVMRustThinLTOBufferLen(const LLVMRustThinLTOBuffer *Buffer) {
return Buffer->data.length();
}
extern "C" const void *
LLVMRustThinLTOBufferThinLinkDataPtr(const LLVMRustThinLTOBuffer *Buffer) {
return Buffer->thin_link_data.data();
}
extern "C" size_t
LLVMRustThinLTOBufferThinLinkDataLen(const LLVMRustThinLTOBuffer *Buffer) {
return Buffer->thin_link_data.length();
}
// This is what we used to parse upstream bitcode for actual ThinLTO
// processing. We'll call this once per module optimized through ThinLTO, and
// it'll be called concurrently on many threads.