rustc_codegen_llvm/back/
lto.rs

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