rustc_codegen_llvm/back/
lto.rs

1use std::collections::BTreeMap;
2use std::ffi::{CStr, CString};
3use std::fs::File;
4use std::path::{Path, PathBuf};
5use std::ptr::NonNull;
6use std::sync::Arc;
7use std::{io, iter, slice};
8
9use object::read::archive::ArchiveFile;
10use object::{Object, ObjectSection};
11use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule, ThinShared};
12use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
13use rustc_codegen_ssa::traits::*;
14use rustc_codegen_ssa::{ModuleCodegen, ModuleKind, looks_like_rust_object_file};
15use rustc_data_structures::fx::FxHashMap;
16use rustc_data_structures::memmap::Mmap;
17use rustc_errors::DiagCtxtHandle;
18use rustc_hir::attrs::SanitizerSet;
19use rustc_middle::bug;
20use rustc_middle::dep_graph::WorkProduct;
21use rustc_session::config::{self, Lto};
22use tracing::{debug, info};
23
24use crate::back::write::{
25    self, CodegenDiagnosticsStage, DiagnosticHandlers, bitcode_section_name, save_temp_bitcode,
26};
27use crate::errors::{LlvmError, LtoBitcodeFromRlib};
28use crate::llvm::{self, build_string};
29use crate::{LlvmCodegenBackend, ModuleLlvm, SimpleCx};
30
31/// We keep track of the computed LTO cache keys from the previous
32/// session to determine which CGUs we can reuse.
33const THIN_LTO_KEYS_INCR_COMP_FILE_NAME: &str = "thin-lto-past-keys.bin";
34
35fn prepare_lto(
36    cgcx: &CodegenContext<LlvmCodegenBackend>,
37    exported_symbols_for_lto: &[String],
38    each_linked_rlib_for_lto: &[PathBuf],
39    dcx: DiagCtxtHandle<'_>,
40) -> (Vec<CString>, Vec<(SerializedModule<ModuleBuffer>, CString)>) {
41    let mut symbols_below_threshold = exported_symbols_for_lto
42        .iter()
43        .map(|symbol| CString::new(symbol.to_owned()).unwrap())
44        .collect::<Vec<CString>>();
45
46    if cgcx.module_config.instrument_coverage || cgcx.module_config.pgo_gen.enabled() {
47        // These are weak symbols that point to the profile version and the
48        // profile name, which need to be treated as exported so LTO doesn't nix
49        // them.
50        const PROFILER_WEAK_SYMBOLS: [&CStr; 2] =
51            [c"__llvm_profile_raw_version", c"__llvm_profile_filename"];
52
53        symbols_below_threshold.extend(PROFILER_WEAK_SYMBOLS.iter().map(|&sym| sym.to_owned()));
54    }
55
56    if cgcx.module_config.sanitizer.contains(SanitizerSet::MEMORY) {
57        let mut msan_weak_symbols = Vec::new();
58
59        // Similar to profiling, preserve weak msan symbol during LTO.
60        if cgcx.module_config.sanitizer_recover.contains(SanitizerSet::MEMORY) {
61            msan_weak_symbols.push(c"__msan_keep_going");
62        }
63
64        if cgcx.module_config.sanitizer_memory_track_origins != 0 {
65            msan_weak_symbols.push(c"__msan_track_origins");
66        }
67
68        symbols_below_threshold.extend(msan_weak_symbols.into_iter().map(|sym| sym.to_owned()));
69    }
70
71    // Preserve LLVM-injected, ASAN-related symbols.
72    // See also https://github.com/rust-lang/rust/issues/113404.
73    symbols_below_threshold.push(c"___asan_globals_registered".to_owned());
74
75    // __llvm_profile_counter_bias is pulled in at link time by an undefined reference to
76    // __llvm_profile_runtime, therefore we won't know until link time if this symbol
77    // should have default visibility.
78    symbols_below_threshold.push(c"__llvm_profile_counter_bias".to_owned());
79
80    // If we're performing LTO for the entire crate graph, then for each of our
81    // upstream dependencies, find the corresponding rlib and load the bitcode
82    // from the archive.
83    //
84    // We save off all the bytecode and LLVM module ids for later processing
85    // with either fat or thin LTO
86    let mut upstream_modules = Vec::new();
87    if cgcx.lto != Lto::ThinLocal {
88        for path in each_linked_rlib_for_lto {
89            let archive_data = unsafe {
90                Mmap::map(std::fs::File::open(&path).expect("couldn't open rlib"))
91                    .expect("couldn't map rlib")
92            };
93            let archive = ArchiveFile::parse(&*archive_data).expect("wanted an rlib");
94            let obj_files = archive
95                .members()
96                .filter_map(|child| {
97                    child.ok().and_then(|c| {
98                        std::str::from_utf8(c.name()).ok().map(|name| (name.trim(), c))
99                    })
100                })
101                .filter(|&(name, _)| looks_like_rust_object_file(name));
102            for (name, child) in obj_files {
103                info!("adding bitcode from {}", name);
104                match get_bitcode_slice_from_object_data(
105                    child.data(&*archive_data).expect("corrupt rlib"),
106                    cgcx,
107                ) {
108                    Ok(data) => {
109                        let module = SerializedModule::FromRlib(data.to_vec());
110                        upstream_modules.push((module, CString::new(name).unwrap()));
111                    }
112                    Err(e) => dcx.emit_fatal(e),
113                }
114            }
115        }
116    }
117
118    (symbols_below_threshold, upstream_modules)
119}
120
121fn get_bitcode_slice_from_object_data<'a>(
122    obj: &'a [u8],
123    cgcx: &CodegenContext<LlvmCodegenBackend>,
124) -> Result<&'a [u8], LtoBitcodeFromRlib> {
125    // We're about to assume the data here is an object file with sections, but if it's raw LLVM IR
126    // that won't work. Fortunately, if that's what we have we can just return the object directly,
127    // so we sniff the relevant magic strings here and return.
128    if obj.starts_with(b"\xDE\xC0\x17\x0B") || obj.starts_with(b"BC\xC0\xDE") {
129        return Ok(obj);
130    }
131    // We drop the "__LLVM," prefix here because on Apple platforms there's a notion of "segment
132    // name" which in the public API for sections gets treated as part of the section name, but
133    // internally in MachOObjectFile.cpp gets treated separately.
134    let section_name = bitcode_section_name(cgcx).to_str().unwrap().trim_start_matches("__LLVM,");
135
136    let obj =
137        object::File::parse(obj).map_err(|err| LtoBitcodeFromRlib { err: err.to_string() })?;
138
139    let section = obj
140        .section_by_name(section_name)
141        .ok_or_else(|| LtoBitcodeFromRlib { err: format!("Can't find section {section_name}") })?;
142
143    section.data().map_err(|err| LtoBitcodeFromRlib { err: err.to_string() })
144}
145
146/// Performs fat LTO by merging all modules into a single one and returning it
147/// for further optimization.
148pub(crate) fn run_fat(
149    cgcx: &CodegenContext<LlvmCodegenBackend>,
150    exported_symbols_for_lto: &[String],
151    each_linked_rlib_for_lto: &[PathBuf],
152    modules: Vec<FatLtoInput<LlvmCodegenBackend>>,
153) -> ModuleCodegen<ModuleLlvm> {
154    let dcx = cgcx.create_dcx();
155    let dcx = dcx.handle();
156    let (symbols_below_threshold, upstream_modules) =
157        prepare_lto(cgcx, exported_symbols_for_lto, each_linked_rlib_for_lto, dcx);
158    let symbols_below_threshold =
159        symbols_below_threshold.iter().map(|c| c.as_ptr()).collect::<Vec<_>>();
160    fat_lto(cgcx, dcx, modules, upstream_modules, &symbols_below_threshold)
161}
162
163/// Performs thin LTO by performing necessary global analysis and returning two
164/// lists, one of the modules that need optimization and another for modules that
165/// can simply be copied over from the incr. comp. cache.
166pub(crate) fn run_thin(
167    cgcx: &CodegenContext<LlvmCodegenBackend>,
168    exported_symbols_for_lto: &[String],
169    each_linked_rlib_for_lto: &[PathBuf],
170    modules: Vec<(String, ThinBuffer)>,
171    cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
172) -> (Vec<ThinModule<LlvmCodegenBackend>>, Vec<WorkProduct>) {
173    let dcx = cgcx.create_dcx();
174    let dcx = dcx.handle();
175    let (symbols_below_threshold, upstream_modules) =
176        prepare_lto(cgcx, exported_symbols_for_lto, each_linked_rlib_for_lto, dcx);
177    let symbols_below_threshold =
178        symbols_below_threshold.iter().map(|c| c.as_ptr()).collect::<Vec<_>>();
179    if cgcx.opts.cg.linker_plugin_lto.enabled() {
180        unreachable!(
181            "We should never reach this case if the LTO step \
182                      is deferred to the linker"
183        );
184    }
185    thin_lto(cgcx, dcx, modules, upstream_modules, cached_modules, &symbols_below_threshold)
186}
187
188pub(crate) fn prepare_thin(module: ModuleCodegen<ModuleLlvm>) -> (String, ThinBuffer) {
189    let name = module.name;
190    let buffer = ThinBuffer::new(module.module_llvm.llmod(), true);
191    (name, buffer)
192}
193
194fn fat_lto(
195    cgcx: &CodegenContext<LlvmCodegenBackend>,
196    dcx: DiagCtxtHandle<'_>,
197    modules: Vec<FatLtoInput<LlvmCodegenBackend>>,
198    mut serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
199    symbols_below_threshold: &[*const libc::c_char],
200) -> ModuleCodegen<ModuleLlvm> {
201    let _timer = cgcx.prof.generic_activity("LLVM_fat_lto_build_monolithic_module");
202    info!("going for a fat lto");
203
204    // Sort out all our lists of incoming modules into two lists.
205    //
206    // * `serialized_modules` (also and argument to this function) contains all
207    //   modules that are serialized in-memory.
208    // * `in_memory` contains modules which are already parsed and in-memory,
209    //   such as from multi-CGU builds.
210    let mut in_memory = Vec::new();
211    for module in modules {
212        match module {
213            FatLtoInput::InMemory(m) => in_memory.push(m),
214            FatLtoInput::Serialized { name, buffer } => {
215                info!("pushing serialized module {:?}", name);
216                serialized_modules.push((buffer, CString::new(name).unwrap()));
217            }
218        }
219    }
220
221    // Find the "costliest" module and merge everything into that codegen unit.
222    // All the other modules will be serialized and reparsed into the new
223    // context, so this hopefully avoids serializing and parsing the largest
224    // codegen unit.
225    //
226    // Additionally use a regular module as the base here to ensure that various
227    // file copy operations in the backend work correctly. The only other kind
228    // of module here should be an allocator one, and if your crate is smaller
229    // than the allocator module then the size doesn't really matter anyway.
230    let costliest_module = in_memory
231        .iter()
232        .enumerate()
233        .filter(|&(_, module)| module.kind == ModuleKind::Regular)
234        .map(|(i, module)| {
235            let cost = unsafe { llvm::LLVMRustModuleCost(module.module_llvm.llmod()) };
236            (cost, i)
237        })
238        .max();
239
240    // If we found a costliest module, we're good to go. Otherwise all our
241    // inputs were serialized which could happen in the case, for example, that
242    // all our inputs were incrementally reread from the cache and we're just
243    // re-executing the LTO passes. If that's the case deserialize the first
244    // module and create a linker with it.
245    let module: ModuleCodegen<ModuleLlvm> = match costliest_module {
246        Some((_cost, i)) => in_memory.remove(i),
247        None => {
248            assert!(!serialized_modules.is_empty(), "must have at least one serialized module");
249            let (buffer, name) = serialized_modules.remove(0);
250            info!("no in-memory regular modules to choose from, parsing {:?}", name);
251            let llvm_module = ModuleLlvm::parse(cgcx, &name, buffer.data(), dcx);
252            ModuleCodegen::new_regular(name.into_string().unwrap(), llvm_module)
253        }
254    };
255    {
256        let (llcx, llmod) = {
257            let llvm = &module.module_llvm;
258            (&llvm.llcx, llvm.llmod())
259        };
260        info!("using {:?} as a base module", module.name);
261
262        // The linking steps below may produce errors and diagnostics within LLVM
263        // which we'd like to handle and print, so set up our diagnostic handlers
264        // (which get unregistered when they go out of scope below).
265        let _handler =
266            DiagnosticHandlers::new(cgcx, dcx, llcx, &module, CodegenDiagnosticsStage::LTO);
267
268        // For all other modules we codegened we'll need to link them into our own
269        // bitcode. All modules were codegened in their own LLVM context, however,
270        // and we want to move everything to the same LLVM context. Currently the
271        // way we know of to do that is to serialize them to a string and them parse
272        // them later. Not great but hey, that's why it's "fat" LTO, right?
273        for module in in_memory {
274            let buffer = ModuleBuffer::new(module.module_llvm.llmod());
275            let llmod_id = CString::new(&module.name[..]).unwrap();
276            serialized_modules.push((SerializedModule::Local(buffer), llmod_id));
277        }
278        // Sort the modules to ensure we produce deterministic results.
279        serialized_modules.sort_by(|module1, module2| module1.1.cmp(&module2.1));
280
281        // For all serialized bitcode files we parse them and link them in as we did
282        // above, this is all mostly handled in C++.
283        let mut linker = Linker::new(llmod);
284        for (bc_decoded, name) in serialized_modules {
285            let _timer = cgcx
286                .prof
287                .generic_activity_with_arg_recorder("LLVM_fat_lto_link_module", |recorder| {
288                    recorder.record_arg(format!("{name:?}"))
289                });
290            info!("linking {:?}", name);
291            let data = bc_decoded.data();
292            linker
293                .add(data)
294                .unwrap_or_else(|()| write::llvm_err(dcx, LlvmError::LoadBitcode { name }));
295        }
296        drop(linker);
297        save_temp_bitcode(cgcx, &module, "lto.input");
298
299        // Internalize everything below threshold to help strip out more modules and such.
300        unsafe {
301            let ptr = symbols_below_threshold.as_ptr();
302            llvm::LLVMRustRunRestrictionPass(
303                llmod,
304                ptr as *const *const libc::c_char,
305                symbols_below_threshold.len() as libc::size_t,
306            );
307        }
308        save_temp_bitcode(cgcx, &module, "lto.after-restriction");
309    }
310
311    module
312}
313
314pub(crate) struct Linker<'a>(&'a mut llvm::Linker<'a>);
315
316impl<'a> Linker<'a> {
317    pub(crate) fn new(llmod: &'a llvm::Module) -> Self {
318        unsafe { Linker(llvm::LLVMRustLinkerNew(llmod)) }
319    }
320
321    pub(crate) fn add(&mut self, bytecode: &[u8]) -> Result<(), ()> {
322        unsafe {
323            if llvm::LLVMRustLinkerAdd(
324                self.0,
325                bytecode.as_ptr() as *const libc::c_char,
326                bytecode.len(),
327            ) {
328                Ok(())
329            } else {
330                Err(())
331            }
332        }
333    }
334}
335
336impl Drop for Linker<'_> {
337    fn drop(&mut self) {
338        unsafe {
339            llvm::LLVMRustLinkerFree(&mut *(self.0 as *mut _));
340        }
341    }
342}
343
344/// Prepare "thin" LTO to get run on these modules.
345///
346/// The general structure of ThinLTO is quite different from the structure of
347/// "fat" LTO above. With "fat" LTO all LLVM modules in question are merged into
348/// one giant LLVM module, and then we run more optimization passes over this
349/// big module after internalizing most symbols. Thin LTO, on the other hand,
350/// avoid this large bottleneck through more targeted optimization.
351///
352/// At a high level Thin LTO looks like:
353///
354///    1. Prepare a "summary" of each LLVM module in question which describes
355///       the values inside, cost of the values, etc.
356///    2. Merge the summaries of all modules in question into one "index"
357///    3. Perform some global analysis on this index
358///    4. For each module, use the index and analysis calculated previously to
359///       perform local transformations on the module, for example inlining
360///       small functions from other modules.
361///    5. Run thin-specific optimization passes over each module, and then code
362///       generate everything at the end.
363///
364/// The summary for each module is intended to be quite cheap, and the global
365/// index is relatively quite cheap to create as well. As a result, the goal of
366/// ThinLTO is to reduce the bottleneck on LTO and enable LTO to be used in more
367/// situations. For example one cheap optimization is that we can parallelize
368/// all codegen modules, easily making use of all the cores on a machine.
369///
370/// With all that in mind, the function here is designed at specifically just
371/// calculating the *index* for ThinLTO. This index will then be shared amongst
372/// all of the `LtoModuleCodegen` units returned below and destroyed once
373/// they all go out of scope.
374fn thin_lto(
375    cgcx: &CodegenContext<LlvmCodegenBackend>,
376    dcx: DiagCtxtHandle<'_>,
377    modules: Vec<(String, ThinBuffer)>,
378    serialized_modules: Vec<(SerializedModule<ModuleBuffer>, CString)>,
379    cached_modules: Vec<(SerializedModule<ModuleBuffer>, WorkProduct)>,
380    symbols_below_threshold: &[*const libc::c_char],
381) -> (Vec<ThinModule<LlvmCodegenBackend>>, Vec<WorkProduct>) {
382    let _timer = cgcx.prof.generic_activity("LLVM_thin_lto_global_analysis");
383    unsafe {
384        info!("going for that thin, thin LTO");
385
386        let green_modules: FxHashMap<_, _> =
387            cached_modules.iter().map(|(_, wp)| (wp.cgu_name.clone(), wp.clone())).collect();
388
389        let full_scope_len = modules.len() + serialized_modules.len() + cached_modules.len();
390        let mut thin_buffers = Vec::with_capacity(modules.len());
391        let mut module_names = Vec::with_capacity(full_scope_len);
392        let mut thin_modules = Vec::with_capacity(full_scope_len);
393
394        for (i, (name, buffer)) in modules.into_iter().enumerate() {
395            info!("local module: {} - {}", i, name);
396            let cname = CString::new(name.as_bytes()).unwrap();
397            thin_modules.push(llvm::ThinLTOModule {
398                identifier: cname.as_ptr(),
399                data: buffer.data().as_ptr(),
400                len: buffer.data().len(),
401            });
402            thin_buffers.push(buffer);
403            module_names.push(cname);
404        }
405
406        // FIXME: All upstream crates are deserialized internally in the
407        //        function below to extract their summary and modules. Note that
408        //        unlike the loop above we *must* decode and/or read something
409        //        here as these are all just serialized files on disk. An
410        //        improvement, however, to make here would be to store the
411        //        module summary separately from the actual module itself. Right
412        //        now this is store in one large bitcode file, and the entire
413        //        file is deflate-compressed. We could try to bypass some of the
414        //        decompression by storing the index uncompressed and only
415        //        lazily decompressing the bytecode if necessary.
416        //
417        //        Note that truly taking advantage of this optimization will
418        //        likely be further down the road. We'd have to implement
419        //        incremental ThinLTO first where we could actually avoid
420        //        looking at upstream modules entirely sometimes (the contents,
421        //        we must always unconditionally look at the index).
422        let mut serialized = Vec::with_capacity(serialized_modules.len() + cached_modules.len());
423
424        let cached_modules =
425            cached_modules.into_iter().map(|(sm, wp)| (sm, CString::new(wp.cgu_name).unwrap()));
426
427        for (module, name) in serialized_modules.into_iter().chain(cached_modules) {
428            info!("upstream or cached module {:?}", name);
429            thin_modules.push(llvm::ThinLTOModule {
430                identifier: name.as_ptr(),
431                data: module.data().as_ptr(),
432                len: module.data().len(),
433            });
434            serialized.push(module);
435            module_names.push(name);
436        }
437
438        // Sanity check
439        assert_eq!(thin_modules.len(), module_names.len());
440
441        // Delegate to the C++ bindings to create some data here. Once this is a
442        // tried-and-true interface we may wish to try to upstream some of this
443        // to LLVM itself, right now we reimplement a lot of what they do
444        // upstream...
445        let data = llvm::LLVMRustCreateThinLTOData(
446            thin_modules.as_ptr(),
447            thin_modules.len(),
448            symbols_below_threshold.as_ptr(),
449            symbols_below_threshold.len(),
450        )
451        .unwrap_or_else(|| write::llvm_err(dcx, LlvmError::PrepareThinLtoContext));
452
453        let data = ThinData(data);
454
455        info!("thin LTO data created");
456
457        let (key_map_path, prev_key_map, curr_key_map) = if let Some(ref incr_comp_session_dir) =
458            cgcx.incr_comp_session_dir
459        {
460            let path = incr_comp_session_dir.join(THIN_LTO_KEYS_INCR_COMP_FILE_NAME);
461            // If the previous file was deleted, or we get an IO error
462            // reading the file, then we'll just use `None` as the
463            // prev_key_map, which will force the code to be recompiled.
464            let prev =
465                if path.exists() { ThinLTOKeysMap::load_from_file(&path).ok() } else { None };
466            let curr = ThinLTOKeysMap::from_thin_lto_modules(&data, &thin_modules, &module_names);
467            (Some(path), prev, curr)
468        } else {
469            // If we don't compile incrementally, we don't need to load the
470            // import data from LLVM.
471            assert!(green_modules.is_empty());
472            let curr = ThinLTOKeysMap::default();
473            (None, None, curr)
474        };
475        info!("thin LTO cache key map loaded");
476        info!("prev_key_map: {:#?}", prev_key_map);
477        info!("curr_key_map: {:#?}", curr_key_map);
478
479        // Throw our data in an `Arc` as we'll be sharing it across threads. We
480        // also put all memory referenced by the C++ data (buffers, ids, etc)
481        // into the arc as well. After this we'll create a thin module
482        // codegen per module in this data.
483        let shared = Arc::new(ThinShared {
484            data,
485            thin_buffers,
486            serialized_modules: serialized,
487            module_names,
488        });
489
490        let mut copy_jobs = vec![];
491        let mut opt_jobs = vec![];
492
493        info!("checking which modules can be-reused and which have to be re-optimized.");
494        for (module_index, module_name) in shared.module_names.iter().enumerate() {
495            let module_name = module_name_to_str(module_name);
496            if let (Some(prev_key_map), true) =
497                (prev_key_map.as_ref(), green_modules.contains_key(module_name))
498            {
499                assert!(cgcx.incr_comp_session_dir.is_some());
500
501                // If a module exists in both the current and the previous session,
502                // and has the same LTO cache key in both sessions, then we can re-use it
503                if prev_key_map.keys.get(module_name) == curr_key_map.keys.get(module_name) {
504                    let work_product = green_modules[module_name].clone();
505                    copy_jobs.push(work_product);
506                    info!(" - {}: re-used", module_name);
507                    assert!(cgcx.incr_comp_session_dir.is_some());
508                    continue;
509                }
510            }
511
512            info!(" - {}: re-compiled", module_name);
513            opt_jobs.push(ThinModule { shared: Arc::clone(&shared), idx: module_index });
514        }
515
516        // Save the current ThinLTO import information for the next compilation
517        // session, overwriting the previous serialized data (if any).
518        if let Some(path) = key_map_path
519            && let Err(err) = curr_key_map.save_to_file(&path)
520        {
521            write::llvm_err(dcx, LlvmError::WriteThinLtoKey { err });
522        }
523
524        (opt_jobs, copy_jobs)
525    }
526}
527
528fn enable_autodiff_settings(ad: &[config::AutoDiff]) {
529    for val in ad {
530        // We intentionally don't use a wildcard, to not forget handling anything new.
531        match val {
532            config::AutoDiff::PrintPerf => {
533                llvm::set_print_perf(true);
534            }
535            config::AutoDiff::PrintAA => {
536                llvm::set_print_activity(true);
537            }
538            config::AutoDiff::PrintTA => {
539                llvm::set_print_type(true);
540            }
541            config::AutoDiff::PrintTAFn(fun) => {
542                llvm::set_print_type(true); // Enable general type printing
543                llvm::set_print_type_fun(&fun); // Set specific function to analyze
544            }
545            config::AutoDiff::Inline => {
546                llvm::set_inline(true);
547            }
548            config::AutoDiff::LooseTypes => {
549                llvm::set_loose_types(true);
550            }
551            config::AutoDiff::PrintSteps => {
552                llvm::set_print(true);
553            }
554            // We handle this in the PassWrapper.cpp
555            config::AutoDiff::PrintPasses => {}
556            // We handle this in the PassWrapper.cpp
557            config::AutoDiff::PrintModBefore => {}
558            // We handle this in the PassWrapper.cpp
559            config::AutoDiff::PrintModAfter => {}
560            // We handle this in the PassWrapper.cpp
561            config::AutoDiff::PrintModFinal => {}
562            // This is required and already checked
563            config::AutoDiff::Enable => {}
564            // We handle this below
565            config::AutoDiff::NoPostopt => {}
566            // Disables TypeTree generation
567            config::AutoDiff::NoTT => {}
568        }
569    }
570    // This helps with handling enums for now.
571    llvm::set_strict_aliasing(false);
572    // FIXME(ZuseZ4): Test this, since it was added a long time ago.
573    llvm::set_rust_rules(true);
574}
575
576pub(crate) fn run_pass_manager(
577    cgcx: &CodegenContext<LlvmCodegenBackend>,
578    dcx: DiagCtxtHandle<'_>,
579    module: &mut ModuleCodegen<ModuleLlvm>,
580    thin: bool,
581) {
582    let _timer = cgcx.prof.generic_activity_with_arg("LLVM_lto_optimize", &*module.name);
583    let config = &cgcx.module_config;
584
585    // Now we have one massive module inside of llmod. Time to run the
586    // LTO-specific optimization passes that LLVM provides.
587    //
588    // This code is based off the code found in llvm's LTO code generator:
589    //      llvm/lib/LTO/LTOCodeGenerator.cpp
590    debug!("running the pass manager");
591    let opt_stage = if thin { llvm::OptStage::ThinLTO } else { llvm::OptStage::FatLTO };
592    let opt_level = config.opt_level.unwrap_or(config::OptLevel::No);
593
594    // The PostAD behavior is the same that we would have if no autodiff was used.
595    // It will run the default optimization pipeline. If AD is enabled we select
596    // the DuringAD stage, which will disable vectorization and loop unrolling, and
597    // schedule two autodiff optimization + differentiation passes.
598    // We then run the llvm_optimize function a second time, to optimize the code which we generated
599    // in the enzyme differentiation pass.
600    let enable_ad = config.autodiff.contains(&config::AutoDiff::Enable);
601    let enable_gpu = config.offload.contains(&config::Offload::Enable);
602    let stage = if thin {
603        write::AutodiffStage::PreAD
604    } else {
605        if enable_ad { write::AutodiffStage::DuringAD } else { write::AutodiffStage::PostAD }
606    };
607
608    if enable_ad {
609        enable_autodiff_settings(&config.autodiff);
610    }
611
612    unsafe {
613        write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage);
614    }
615
616    if enable_gpu && !thin {
617        let cx =
618            SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);
619        crate::builder::gpu_offload::handle_gpu_code(cgcx, &cx);
620    }
621
622    if cfg!(feature = "llvm_enzyme") && enable_ad && !thin {
623        let opt_stage = llvm::OptStage::FatLTO;
624        let stage = write::AutodiffStage::PostAD;
625        if !config.autodiff.contains(&config::AutoDiff::NoPostopt) {
626            unsafe {
627                write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage);
628            }
629        }
630
631        // This is the final IR, so people should be able to inspect the optimized autodiff output,
632        // for manual inspection.
633        if config.autodiff.contains(&config::AutoDiff::PrintModFinal) {
634            unsafe { llvm::LLVMDumpModule(module.module_llvm.llmod()) };
635        }
636    }
637
638    debug!("lto done");
639}
640
641pub struct ModuleBuffer(&'static mut llvm::ModuleBuffer);
642
643unsafe impl Send for ModuleBuffer {}
644unsafe impl Sync for ModuleBuffer {}
645
646impl ModuleBuffer {
647    pub(crate) fn new(m: &llvm::Module) -> ModuleBuffer {
648        ModuleBuffer(unsafe { llvm::LLVMRustModuleBufferCreate(m) })
649    }
650}
651
652impl ModuleBufferMethods for ModuleBuffer {
653    fn data(&self) -> &[u8] {
654        unsafe {
655            let ptr = llvm::LLVMRustModuleBufferPtr(self.0);
656            let len = llvm::LLVMRustModuleBufferLen(self.0);
657            slice::from_raw_parts(ptr, len)
658        }
659    }
660}
661
662impl Drop for ModuleBuffer {
663    fn drop(&mut self) {
664        unsafe {
665            llvm::LLVMRustModuleBufferFree(&mut *(self.0 as *mut _));
666        }
667    }
668}
669
670pub struct ThinData(&'static mut llvm::ThinLTOData);
671
672unsafe impl Send for ThinData {}
673unsafe impl Sync for ThinData {}
674
675impl Drop for ThinData {
676    fn drop(&mut self) {
677        unsafe {
678            llvm::LLVMRustFreeThinLTOData(&mut *(self.0 as *mut _));
679        }
680    }
681}
682
683pub struct ThinBuffer(&'static mut llvm::ThinLTOBuffer);
684
685unsafe impl Send for ThinBuffer {}
686unsafe impl Sync for ThinBuffer {}
687
688impl ThinBuffer {
689    pub(crate) fn new(m: &llvm::Module, is_thin: bool) -> ThinBuffer {
690        unsafe {
691            let buffer = llvm::LLVMRustThinLTOBufferCreate(m, is_thin);
692            ThinBuffer(buffer)
693        }
694    }
695
696    pub(crate) unsafe fn from_raw_ptr(ptr: *mut llvm::ThinLTOBuffer) -> ThinBuffer {
697        let mut ptr = NonNull::new(ptr).unwrap();
698        ThinBuffer(unsafe { ptr.as_mut() })
699    }
700
701    pub(crate) fn thin_link_data(&self) -> &[u8] {
702        unsafe {
703            let ptr = llvm::LLVMRustThinLTOBufferThinLinkDataPtr(self.0) as *const _;
704            let len = llvm::LLVMRustThinLTOBufferThinLinkDataLen(self.0);
705            slice::from_raw_parts(ptr, len)
706        }
707    }
708}
709
710impl ThinBufferMethods for ThinBuffer {
711    fn data(&self) -> &[u8] {
712        unsafe {
713            let ptr = llvm::LLVMRustThinLTOBufferPtr(self.0) as *const _;
714            let len = llvm::LLVMRustThinLTOBufferLen(self.0);
715            slice::from_raw_parts(ptr, len)
716        }
717    }
718}
719
720impl Drop for ThinBuffer {
721    fn drop(&mut self) {
722        unsafe {
723            llvm::LLVMRustThinLTOBufferFree(&mut *(self.0 as *mut _));
724        }
725    }
726}
727
728pub(crate) fn optimize_thin_module(
729    thin_module: ThinModule<LlvmCodegenBackend>,
730    cgcx: &CodegenContext<LlvmCodegenBackend>,
731) -> ModuleCodegen<ModuleLlvm> {
732    let dcx = cgcx.create_dcx();
733    let dcx = dcx.handle();
734
735    let module_name = &thin_module.shared.module_names[thin_module.idx];
736
737    // Right now the implementation we've got only works over serialized
738    // modules, so we create a fresh new LLVM context and parse the module
739    // into that context. One day, however, we may do this for upstream
740    // crates but for locally codegened modules we may be able to reuse
741    // that LLVM Context and Module.
742    let module_llvm = ModuleLlvm::parse(cgcx, module_name, thin_module.data(), dcx);
743    let mut module = ModuleCodegen::new_regular(thin_module.name(), module_llvm);
744    // Given that the newly created module lacks a thinlto buffer for embedding, we need to re-add it here.
745    if cgcx.module_config.embed_bitcode() {
746        module.thin_lto_buffer = Some(thin_module.data().to_vec());
747    }
748    {
749        let target = &*module.module_llvm.tm;
750        let llmod = module.module_llvm.llmod();
751        save_temp_bitcode(cgcx, &module, "thin-lto-input");
752
753        // Up next comes the per-module local analyses that we do for Thin LTO.
754        // Each of these functions is basically copied from the LLVM
755        // implementation and then tailored to suit this implementation. Ideally
756        // each of these would be supported by upstream LLVM but that's perhaps
757        // a patch for another day!
758        //
759        // You can find some more comments about these functions in the LLVM
760        // bindings we've got (currently `PassWrapper.cpp`)
761        {
762            let _timer =
763                cgcx.prof.generic_activity_with_arg("LLVM_thin_lto_rename", thin_module.name());
764            unsafe {
765                llvm::LLVMRustPrepareThinLTORename(thin_module.shared.data.0, llmod, target.raw())
766            };
767            save_temp_bitcode(cgcx, &module, "thin-lto-after-rename");
768        }
769
770        {
771            let _timer = cgcx
772                .prof
773                .generic_activity_with_arg("LLVM_thin_lto_resolve_weak", thin_module.name());
774            if unsafe { !llvm::LLVMRustPrepareThinLTOResolveWeak(thin_module.shared.data.0, llmod) }
775            {
776                write::llvm_err(dcx, LlvmError::PrepareThinLtoModule);
777            }
778            save_temp_bitcode(cgcx, &module, "thin-lto-after-resolve");
779        }
780
781        {
782            let _timer = cgcx
783                .prof
784                .generic_activity_with_arg("LLVM_thin_lto_internalize", thin_module.name());
785            if unsafe { !llvm::LLVMRustPrepareThinLTOInternalize(thin_module.shared.data.0, llmod) }
786            {
787                write::llvm_err(dcx, LlvmError::PrepareThinLtoModule);
788            }
789            save_temp_bitcode(cgcx, &module, "thin-lto-after-internalize");
790        }
791
792        {
793            let _timer =
794                cgcx.prof.generic_activity_with_arg("LLVM_thin_lto_import", thin_module.name());
795            if unsafe {
796                !llvm::LLVMRustPrepareThinLTOImport(thin_module.shared.data.0, llmod, target.raw())
797            } {
798                write::llvm_err(dcx, LlvmError::PrepareThinLtoModule);
799            }
800            save_temp_bitcode(cgcx, &module, "thin-lto-after-import");
801        }
802
803        // Alright now that we've done everything related to the ThinLTO
804        // analysis it's time to run some optimizations! Here we use the same
805        // `run_pass_manager` as the "fat" LTO above except that we tell it to
806        // populate a thin-specific pass manager, which presumably LLVM treats a
807        // little differently.
808        {
809            info!("running thin lto passes over {}", module.name);
810            run_pass_manager(cgcx, dcx, &mut module, true);
811            save_temp_bitcode(cgcx, &module, "thin-lto-after-pm");
812        }
813    }
814    module
815}
816
817/// Maps LLVM module identifiers to their corresponding LLVM LTO cache keys
818#[derive(Debug, Default)]
819struct ThinLTOKeysMap {
820    // key = llvm name of importing module, value = LLVM cache key
821    keys: BTreeMap<String, String>,
822}
823
824impl ThinLTOKeysMap {
825    fn save_to_file(&self, path: &Path) -> io::Result<()> {
826        use std::io::Write;
827        let mut writer = File::create_buffered(path)?;
828        // The entries are loaded back into a hash map in `load_from_file()`, so
829        // the order in which we write them to file here does not matter.
830        for (module, key) in &self.keys {
831            writeln!(writer, "{module} {key}")?;
832        }
833        Ok(())
834    }
835
836    fn load_from_file(path: &Path) -> io::Result<Self> {
837        use std::io::BufRead;
838        let mut keys = BTreeMap::default();
839        let file = File::open_buffered(path)?;
840        for line in file.lines() {
841            let line = line?;
842            let mut split = line.split(' ');
843            let module = split.next().unwrap();
844            let key = split.next().unwrap();
845            assert_eq!(split.next(), None, "Expected two space-separated values, found {line:?}");
846            keys.insert(module.to_string(), key.to_string());
847        }
848        Ok(Self { keys })
849    }
850
851    fn from_thin_lto_modules(
852        data: &ThinData,
853        modules: &[llvm::ThinLTOModule],
854        names: &[CString],
855    ) -> Self {
856        let keys = iter::zip(modules, names)
857            .map(|(module, name)| {
858                let key = build_string(|rust_str| unsafe {
859                    llvm::LLVMRustComputeLTOCacheKey(rust_str, module.identifier, data.0);
860                })
861                .expect("Invalid ThinLTO module key");
862                (module_name_to_str(name).to_string(), key)
863            })
864            .collect();
865        Self { keys }
866    }
867}
868
869fn module_name_to_str(c_str: &CStr) -> &str {
870    c_str.to_str().unwrap_or_else(|e| {
871        bug!("Encountered non-utf8 LLVM module name `{}`: {}", c_str.to_string_lossy(), e)
872    })
873}
874
875pub(crate) fn parse_module<'a>(
876    cx: &'a llvm::Context,
877    name: &CStr,
878    data: &[u8],
879    dcx: DiagCtxtHandle<'_>,
880) -> &'a llvm::Module {
881    unsafe {
882        llvm::LLVMRustParseBitcodeForLTO(cx, data.as_ptr(), data.len(), name.as_ptr())
883            .unwrap_or_else(|| write::llvm_err(dcx, LlvmError::ParseBitcode))
884    }
885}
rustc_codegen_llvm/back/lto.rs

rustc_codegen_llvm/back/
lto.rs
