Skip to content
This repository was archived by the owner on May 28, 2025. It is now read-only.

Commit 0c72c0d

Browse files
borsbors
committed
Auto merge of rust-lang#133250 - DianQK:embed-bitcode-pgo, r=nikic
The embedded bitcode should always be prepared for LTO/ThinLTO Fixes rust-lang#115344. Fixes rust-lang#117220. There are currently two methods for generating bitcode that used for LTO. One method involves using `-C linker-plugin-lto` to emit object files as bitcode, which is the typical setting used by cargo. The other method is through `-C embed-bitcode=yes`. When using with `-C embed-bitcode=yes -C lto=no`, we run a complete non-LTO LLVM pipeline to obtain bitcode, then the bitcode is used for LTO. We run the Call Graph Profile Pass twice on the same module. This PR is doing something similar to LLVM's `buildFatLTODefaultPipeline`, obtaining the bitcode for embedding after running `buildThinLTOPreLinkDefaultPipeline`. r? nikic
2 parents 002da76 + a897cc0 commit 0c72c0d

File tree

20 files changed

+294
-101
lines changed

20 files changed

+294
-101
lines changed

‎compiler/rustc_codegen_gcc/src/back/lto.rs

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -632,17 +632,16 @@ pub unsafe fn optimize_thin_module(
632632
Arc::new(SyncContext::new(context))
633633
}
634634
};
635-
let module = ModuleCodegen {
636-
module_llvm: GccContext {
635+
let module = ModuleCodegen::new_regular(
636+
thin_module.name().to_string(),
637+
GccContext {
637638
context,
638639
should_combine_object_files,
639640
// TODO(antoyo): use the correct relocation model here.
640641
relocation_model: RelocModel::Pic,
641642
temp_dir: None,
642643
},
643-
name: thin_module.name().to_string(),
644-
kind: ModuleKind::Regular,
645-
};
644+
);
646645
/*{
647646
let target = &*module.module_llvm.tm;
648647
let llmod = module.module_llvm.llmod();

‎compiler/rustc_codegen_gcc/src/base.rs

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,10 @@ use std::sync::Arc;
44
use std::time::Instant;
55

66
use gccjit::{CType, Context, FunctionType, GlobalKind};
7+
use rustc_codegen_ssa::ModuleCodegen;
78
use rustc_codegen_ssa::base::maybe_create_entry_wrapper;
89
use rustc_codegen_ssa::mono_item::MonoItemExt;
910
use rustc_codegen_ssa::traits::DebugInfoCodegenMethods;
10-
use rustc_codegen_ssa::{ModuleCodegen, ModuleKind};
1111
use rustc_middle::dep_graph;
1212
use rustc_middle::mir::mono::Linkage;
1313
#[cfg(feature = "master")]
@@ -237,16 +237,15 @@ pub fn compile_codegen_unit(
237237
}
238238
}
239239

240-
ModuleCodegen {
241-
name: cgu_name.to_string(),
242-
module_llvm: GccContext {
240+
ModuleCodegen::new_regular(
241+
cgu_name.to_string(),
242+
GccContext {
243243
context: Arc::new(SyncContext::new(context)),
244244
relocation_model: tcx.sess.relocation_model(),
245245
should_combine_object_files: false,
246246
temp_dir: None,
247247
},
248-
kind: ModuleKind::Regular,
249-
}
248+
)
250249
}
251250

252251
(module, cost)

‎compiler/rustc_codegen_gcc/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -393,7 +393,7 @@ impl WriteBackendMethods for GccCodegenBackend {
393393
unsafe fn optimize(
394394
_cgcx: &CodegenContext<Self>,
395395
_dcx: DiagCtxtHandle<'_>,
396-
module: &ModuleCodegen<Self::Module>,
396+
module: &mut ModuleCodegen<Self::Module>,
397397
config: &ModuleConfig,
398398
) -> Result<(), FatalError> {
399399
module.module_llvm.context.set_optimization_level(to_gcc_opt_level(config.opt_level));

‎compiler/rustc_codegen_llvm/src/back/lto.rs

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
22
use std::ffi::{CStr, CString};
33
use std::fs::File;
44
use std::path::Path;
5+
use std::ptr::NonNull;
56
use std::sync::Arc;
67
use std::{io, iter, slice};
78

@@ -305,11 +306,8 @@ fn fat_lto(
305306
assert!(!serialized_modules.is_empty(), "must have at least one serialized module");
306307
let (buffer, name) = serialized_modules.remove(0);
307308
info!("no in-memory regular modules to choose from, parsing {:?}", name);
308-
ModuleCodegen {
309-
module_llvm: ModuleLlvm::parse(cgcx, &name, buffer.data(), dcx)?,
310-
name: name.into_string().unwrap(),
311-
kind: ModuleKind::Regular,
312-
}
309+
let llvm_module = ModuleLlvm::parse(cgcx, &name, buffer.data(), dcx)?;
310+
ModuleCodegen::new_regular(name.into_string().unwrap(), llvm_module)
313311
}
314312
};
315313
{
@@ -655,14 +653,14 @@ pub(crate) fn run_pass_manager(
655653
}
656654

657655
unsafe {
658-
write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, stage)?;
656+
write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
659657
}
660658

661659
if cfg!(llvm_enzyme) && enable_ad {
662660
let opt_stage = llvm::OptStage::FatLTO;
663661
let stage = write::AutodiffStage::PostAD;
664662
unsafe {
665-
write::llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, stage)?;
663+
write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
666664
}
667665

668666
// This is the final IR, so people should be able to inspect the optimized autodiff output.
@@ -729,6 +727,11 @@ impl ThinBuffer {
729727
ThinBuffer(buffer)
730728
}
731729
}
730+
731+
pub unsafe fn from_raw_ptr(ptr: *mut llvm::ThinLTOBuffer) -> ThinBuffer {
732+
let mut ptr = NonNull::new(ptr).unwrap();
733+
ThinBuffer(unsafe { ptr.as_mut() })
734+
}
732735
}
733736

734737
impl ThinBufferMethods for ThinBuffer {
@@ -772,11 +775,11 @@ pub(crate) unsafe fn optimize_thin_module(
772775
// crates but for locally codegened modules we may be able to reuse
773776
// that LLVM Context and Module.
774777
let module_llvm = ModuleLlvm::parse(cgcx, module_name, thin_module.data(), dcx)?;
775-
let mut module = ModuleCodegen {
776-
module_llvm,
777-
name: thin_module.name().to_string(),
778-
kind: ModuleKind::Regular,
779-
};
778+
let mut module = ModuleCodegen::new_regular(thin_module.name(), module_llvm);
779+
// Given that the newly created module lacks a thinlto buffer for embedding, we need to re-add it here.
780+
if cgcx.config(ModuleKind::Regular).embed_bitcode() {
781+
module.thin_lto_buffer = Some(thin_module.data().to_vec());
782+
}
780783
{
781784
let target = &*module.module_llvm.tm;
782785
let llmod = module.module_llvm.llmod();

‎compiler/rustc_codegen_llvm/src/back/write.rs

Lines changed: 83 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use std::ffi::{CStr, CString};
22
use std::io::{self, Write};
33
use std::path::{Path, PathBuf};
4+
use std::ptr::null_mut;
45
use std::sync::Arc;
56
use std::{fs, slice, str};
67

@@ -15,7 +16,7 @@ use rustc_codegen_ssa::back::write::{
1516
TargetMachineFactoryFn,
1617
};
1718
use rustc_codegen_ssa::traits::*;
18-
use rustc_codegen_ssa::{CompiledModule, ModuleCodegen};
19+
use rustc_codegen_ssa::{CompiledModule, ModuleCodegen, ModuleKind};
1920
use rustc_data_structures::profiling::SelfProfilerRef;
2021
use rustc_data_structures::small_c_str::SmallCStr;
2122
use rustc_errors::{DiagCtxtHandle, FatalError, Level};
@@ -551,6 +552,7 @@ pub(crate) unsafe fn llvm_optimize(
551552
cgcx: &CodegenContext<LlvmCodegenBackend>,
552553
dcx: DiagCtxtHandle<'_>,
553554
module: &ModuleCodegen<ModuleLlvm>,
555+
thin_lto_buffer: Option<&mut *mut llvm::ThinLTOBuffer>,
554556
config: &ModuleConfig,
555557
opt_level: config::OptLevel,
556558
opt_stage: llvm::OptStage,
@@ -584,7 +586,17 @@ pub(crate) unsafe fn llvm_optimize(
584586
vectorize_loop = config.vectorize_loop;
585587
}
586588
trace!(?unroll_loops, ?vectorize_slp, ?vectorize_loop, ?run_enzyme);
587-
let using_thin_buffers = opt_stage == llvm::OptStage::PreLinkThinLTO || config.bitcode_needed();
589+
if thin_lto_buffer.is_some() {
590+
assert!(
591+
matches!(
592+
opt_stage,
593+
llvm::OptStage::PreLinkNoLTO
594+
| llvm::OptStage::PreLinkFatLTO
595+
| llvm::OptStage::PreLinkThinLTO
596+
),
597+
"the bitcode for LTO can only be obtained at the pre-link stage"
598+
);
599+
}
588600
let pgo_gen_path = get_pgo_gen_path(config);
589601
let pgo_use_path = get_pgo_use_path(config);
590602
let pgo_sample_use_path = get_pgo_sample_use_path(config);
@@ -644,7 +656,9 @@ pub(crate) unsafe fn llvm_optimize(
644656
config.no_prepopulate_passes,
645657
config.verify_llvm_ir,
646658
config.lint_llvm_ir,
647-
using_thin_buffers,
659+
thin_lto_buffer,
660+
config.emit_thin_lto,
661+
config.emit_thin_lto_summary,
648662
config.merge_functions,
649663
unroll_loops,
650664
vectorize_slp,
@@ -675,7 +689,7 @@ pub(crate) unsafe fn llvm_optimize(
675689
pub(crate) unsafe fn optimize(
676690
cgcx: &CodegenContext<LlvmCodegenBackend>,
677691
dcx: DiagCtxtHandle<'_>,
678-
module: &ModuleCodegen<ModuleLlvm>,
692+
module: &mut ModuleCodegen<ModuleLlvm>,
679693
config: &ModuleConfig,
680694
) -> Result<(), FatalError> {
681695
let _timer = cgcx.prof.generic_activity_with_arg("LLVM_module_optimize", &*module.name);
@@ -705,9 +719,53 @@ pub(crate) unsafe fn optimize(
705719
// Otherwise we pretend AD is already done and run the normal opt pipeline (=PostAD).
706720
let consider_ad = cfg!(llvm_enzyme) && config.autodiff.contains(&config::AutoDiff::Enable);
707721
let autodiff_stage = if consider_ad { AutodiffStage::PreAD } else { AutodiffStage::PostAD };
708-
return unsafe {
709-
llvm_optimize(cgcx, dcx, module, config, opt_level, opt_stage, autodiff_stage)
722+
// The embedded bitcode is used to run LTO/ThinLTO.
723+
// The bitcode obtained during the `codegen` phase is no longer suitable for performing LTO.
724+
// It may have undergone LTO due to ThinLocal, so we need to obtain the embedded bitcode at
725+
// this point.
726+
let mut thin_lto_buffer = if (module.kind == ModuleKind::Regular
727+
&& config.emit_obj == EmitObj::ObjectCode(BitcodeSection::Full))
728+
|| config.emit_thin_lto_summary
729+
{
730+
Some(null_mut())
731+
} else {
732+
None
710733
};
734+
unsafe {
735+
llvm_optimize(
736+
cgcx,
737+
dcx,
738+
module,
739+
thin_lto_buffer.as_mut(),
740+
config,
741+
opt_level,
742+
opt_stage,
743+
autodiff_stage,
744+
)
745+
}?;
746+
if let Some(thin_lto_buffer) = thin_lto_buffer {
747+
let thin_lto_buffer = unsafe { ThinBuffer::from_raw_ptr(thin_lto_buffer) };
748+
module.thin_lto_buffer = Some(thin_lto_buffer.data().to_vec());
749+
let bc_summary_out =
750+
cgcx.output_filenames.temp_path(OutputType::ThinLinkBitcode, module_name);
751+
if config.emit_thin_lto_summary
752+
&& let Some(thin_link_bitcode_filename) = bc_summary_out.file_name()
753+
{
754+
let summary_data = thin_lto_buffer.thin_link_data();
755+
cgcx.prof.artifact_size(
756+
"llvm_bitcode_summary",
757+
thin_link_bitcode_filename.to_string_lossy(),
758+
summary_data.len() as u64,
759+
);
760+
let _timer = cgcx.prof.generic_activity_with_arg(
761+
"LLVM_module_codegen_emit_bitcode_summary",
762+
&*module.name,
763+
);
764+
if let Err(err) = fs::write(&bc_summary_out, summary_data) {
765+
dcx.emit_err(WriteBytecode { path: &bc_summary_out, err });
766+
}
767+
}
768+
}
711769
}
712770
Ok(())
713771
}
@@ -760,59 +818,41 @@ pub(crate) unsafe fn codegen(
760818
// otherwise requested.
761819

762820
let bc_out = cgcx.output_filenames.temp_path(OutputType::Bitcode, module_name);
763-
let bc_summary_out =
764-
cgcx.output_filenames.temp_path(OutputType::ThinLinkBitcode, module_name);
765821
let obj_out = cgcx.output_filenames.temp_path(OutputType::Object, module_name);
766822

767823
if config.bitcode_needed() {
768-
let _timer = cgcx
769-
.prof
770-
.generic_activity_with_arg("LLVM_module_codegen_make_bitcode", &*module.name);
771-
let thin = ThinBuffer::new(llmod, config.emit_thin_lto, config.emit_thin_lto_summary);
772-
let data = thin.data();
773-
774-
if let Some(bitcode_filename) = bc_out.file_name() {
775-
cgcx.prof.artifact_size(
776-
"llvm_bitcode",
777-
bitcode_filename.to_string_lossy(),
778-
data.len() as u64,
779-
);
780-
}
781-
782-
if config.emit_thin_lto_summary
783-
&& let Some(thin_link_bitcode_filename) = bc_summary_out.file_name()
784-
{
785-
let summary_data = thin.thin_link_data();
786-
cgcx.prof.artifact_size(
787-
"llvm_bitcode_summary",
788-
thin_link_bitcode_filename.to_string_lossy(),
789-
summary_data.len() as u64,
790-
);
791-
792-
let _timer = cgcx.prof.generic_activity_with_arg(
793-
"LLVM_module_codegen_emit_bitcode_summary",
794-
&*module.name,
795-
);
796-
if let Err(err) = fs::write(&bc_summary_out, summary_data) {
797-
dcx.emit_err(WriteBytecode { path: &bc_summary_out, err });
798-
}
799-
}
800-
801824
if config.emit_bc || config.emit_obj == EmitObj::Bitcode {
825+
let thin = {
826+
let _timer = cgcx.prof.generic_activity_with_arg(
827+
"LLVM_module_codegen_make_bitcode",
828+
&*module.name,
829+
);
830+
ThinBuffer::new(llmod, config.emit_thin_lto, false)
831+
};
832+
let data = thin.data();
802833
let _timer = cgcx
803834
.prof
804835
.generic_activity_with_arg("LLVM_module_codegen_emit_bitcode", &*module.name);
836+
if let Some(bitcode_filename) = bc_out.file_name() {
837+
cgcx.prof.artifact_size(
838+
"llvm_bitcode",
839+
bitcode_filename.to_string_lossy(),
840+
data.len() as u64,
841+
);
842+
}
805843
if let Err(err) = fs::write(&bc_out, data) {
806844
dcx.emit_err(WriteBytecode { path: &bc_out, err });
807845
}
808846
}
809847

810-
if config.emit_obj == EmitObj::ObjectCode(BitcodeSection::Full) {
848+
if config.embed_bitcode() && module.kind == ModuleKind::Regular {
811849
let _timer = cgcx
812850
.prof
813851
.generic_activity_with_arg("LLVM_module_codegen_embed_bitcode", &*module.name);
852+
let thin_bc =
853+
module.thin_lto_buffer.as_deref().expect("cannot find embedded bitcode");
814854
unsafe {
815-
embed_bitcode(cgcx, llcx, llmod, &config.bc_cmdline, data);
855+
embed_bitcode(cgcx, llcx, llmod, &config.bc_cmdline, &thin_bc);
816856
}
817857
}
818858
}

‎compiler/rustc_codegen_llvm/src/base.rs

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,10 @@
1313
1414
use std::time::Instant;
1515

16+
use rustc_codegen_ssa::ModuleCodegen;
1617
use rustc_codegen_ssa::base::maybe_create_entry_wrapper;
1718
use rustc_codegen_ssa::mono_item::MonoItemExt;
1819
use rustc_codegen_ssa::traits::*;
19-
use rustc_codegen_ssa::{ModuleCodegen, ModuleKind};
2020
use rustc_data_structures::small_c_str::SmallCStr;
2121
use rustc_middle::dep_graph;
2222
use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
@@ -133,11 +133,7 @@ pub(crate) fn compile_codegen_unit(
133133
}
134134
}
135135

136-
ModuleCodegen {
137-
name: cgu_name.to_string(),
138-
module_llvm: llvm_module,
139-
kind: ModuleKind::Regular,
140-
}
136+
ModuleCodegen::new_regular(cgu_name.to_string(), llvm_module)
141137
}
142138

143139
(module, cost)

‎compiler/rustc_codegen_llvm/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ impl WriteBackendMethods for LlvmCodegenBackend {
194194
unsafe fn optimize(
195195
cgcx: &CodegenContext<Self>,
196196
dcx: DiagCtxtHandle<'_>,
197-
module: &ModuleCodegen<Self::Module>,
197+
module: &mut ModuleCodegen<Self::Module>,
198198
config: &ModuleConfig,
199199
) -> Result<(), FatalError> {
200200
unsafe { back::write::optimize(cgcx, dcx, module, config) }

‎compiler/rustc_codegen_llvm/src/llvm/ffi.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2425,7 +2425,9 @@ unsafe extern "C" {
24252425
NoPrepopulatePasses: bool,
24262426
VerifyIR: bool,
24272427
LintIR: bool,
2428-
UseThinLTOBuffers: bool,
2428+
ThinLTOBuffer: Option<&mut *mut ThinLTOBuffer>,
2429+
EmitThinLTO: bool,
2430+
EmitThinLTOSummary: bool,
24292431
MergeFunctions: bool,
24302432
UnrollLoops: bool,
24312433
SLPVectorize: bool,

0 commit comments

Comments
 (0)