rustc_codegen_spirv/linker/
mod.rs

1#[cfg(test)]
2mod test;
3
4mod dce;
5mod destructure_composites;
6mod duplicates;
7mod entry_interface;
8mod import_export_link;
9mod inline;
10mod ipo;
11mod mem2reg;
12mod param_weakening;
13mod peephole_opts;
14mod simple_passes;
15mod specializer;
16mod spirt_passes;
17mod zombies;
18
19use std::borrow::Cow;
20
21use crate::codegen_cx::{ModuleOutputType, SpirvMetadata};
22use crate::custom_decorations::{CustomDecoration, SrcLocDecoration, ZombieDecoration};
23use crate::custom_insts;
24use either::Either;
25use rspirv::binary::{Assemble, Consumer};
26use rspirv::dr::{Block, Loader, Module, ModuleHeader, Operand};
27use rspirv::spirv::{Op, StorageClass, Word};
28use rustc_data_structures::fx::FxHashMap;
29use rustc_errors::ErrorGuaranteed;
30use rustc_session::Session;
31use rustc_session::config::OutputFilenames;
32use std::collections::BTreeMap;
33use std::ffi::{OsStr, OsString};
34use std::path::PathBuf;
35
36pub type Result<T> = std::result::Result<T, ErrorGuaranteed>;
37
38#[derive(Default)]
39pub struct Options {
40    pub compact_ids: bool,
41    pub dce: bool,
42    pub early_report_zombies: bool,
43    pub infer_storage_classes: bool,
44    pub structurize: bool,
45    pub spirt_passes: Vec<String>,
46
47    pub abort_strategy: Option<String>,
48    pub module_output_type: ModuleOutputType,
49
50    pub spirv_metadata: SpirvMetadata,
51
52    /// Whether to preserve `LinkageAttributes "..." Export` decorations,
53    /// even after resolving imports to exports.
54    ///
55    /// **Note**: currently only used for unit testing, and not exposed elsewhere.
56    pub keep_link_exports: bool,
57
58    // NOTE(eddyb) these are debugging options that used to be env vars
59    // (for more information see `docs/src/codegen-args.md`).
60    pub dump_post_merge: Option<PathBuf>,
61    pub dump_pre_inline: Option<PathBuf>,
62    pub dump_post_inline: Option<PathBuf>,
63    pub dump_post_split: Option<PathBuf>,
64    pub dump_spirt_passes: Option<PathBuf>,
65    pub spirt_strip_custom_debuginfo_from_dumps: bool,
66    pub spirt_keep_debug_sources_in_dumps: bool,
67    pub spirt_keep_unstructured_cfg_in_dumps: bool,
68    pub specializer_dump_instances: Option<PathBuf>,
69}
70
71pub enum LinkResult {
72    SingleModule(Box<Module>),
73    MultipleModules {
74        /// The "file stem" key is computed from the "entry name" in the value
75        /// (through `sanitize_filename`, replacing invalid chars with `-`),
76        /// but it's used as the map key because it *has to* be unique, even if
77        /// lossy sanitization could have erased distinctions between entry names.
78        file_stem_to_entry_name_and_module: BTreeMap<OsString, (String, Module)>,
79    },
80}
81
82fn id(header: &mut ModuleHeader) -> Word {
83    let result = header.bound;
84    header.bound += 1;
85    result
86}
87
88fn apply_rewrite_rules<'a>(
89    rewrite_rules: &FxHashMap<Word, Word>,
90    blocks: impl IntoIterator<Item = &'a mut Block>,
91) {
92    let all_ids_mut = blocks
93        .into_iter()
94        .flat_map(|b| b.label.iter_mut().chain(b.instructions.iter_mut()))
95        .flat_map(|inst| {
96            inst.result_id
97                .iter_mut()
98                .chain(inst.result_type.iter_mut())
99                .chain(
100                    inst.operands
101                        .iter_mut()
102                        .filter_map(|op| op.id_ref_any_mut()),
103                )
104        });
105    for id in all_ids_mut {
106        if let Some(&rewrite) = rewrite_rules.get(id) {
107            *id = rewrite;
108        }
109    }
110}
111
112fn get_names(module: &Module) -> FxHashMap<Word, &str> {
113    let entry_names = module
114        .entry_points
115        .iter()
116        .filter(|i| i.class.opcode == Op::EntryPoint)
117        .map(|i| {
118            (
119                i.operands[1].unwrap_id_ref(),
120                i.operands[2].unwrap_literal_string(),
121            )
122        });
123    let debug_names = module
124        .debug_names
125        .iter()
126        .filter(|i| i.class.opcode == Op::Name)
127        .map(|i| {
128            (
129                i.operands[0].unwrap_id_ref(),
130                i.operands[1].unwrap_literal_string(),
131            )
132        });
133    // items later on take priority
134    entry_names.chain(debug_names).collect()
135}
136
137fn get_name<'a>(names: &FxHashMap<Word, &'a str>, id: Word) -> Cow<'a, str> {
138    names.get(&id).map_or_else(
139        || Cow::Owned(format!("Unnamed function ID %{id}")),
140        |&s| Cow::Borrowed(s),
141    )
142}
143
144impl Options {
145    // FIXME(eddyb) using a method on this type seems a bit sketchy.
146    fn spirt_cleanup_for_dumping(&self, module: &mut spirt::Module) {
147        if self.spirt_strip_custom_debuginfo_from_dumps {
148            spirt_passes::debuginfo::convert_custom_debuginfo_to_spv(module);
149        }
150        if !self.spirt_keep_debug_sources_in_dumps {
151            const DOTS: &str = "⋯";
152            let dots_interned_str = module.cx().intern(DOTS);
153            let spirt::ModuleDebugInfo::Spv(debuginfo) = &mut module.debug_info;
154            for sources in debuginfo.source_languages.values_mut() {
155                for file in sources.file_contents.values_mut() {
156                    *file = DOTS.into();
157                }
158                sources.file_contents.insert(
159                    dots_interned_str,
160                    "sources hidden, to show them use \
161                     `RUSTGPU_CODEGEN_ARGS=--spirt-keep-debug-sources-in-dumps`"
162                        .into(),
163                );
164            }
165        }
166    }
167}
168
169pub fn link(
170    sess: &Session,
171    mut inputs: Vec<Module>,
172    opts: &Options,
173    outputs: &OutputFilenames,
174    disambiguated_crate_name_for_dumps: &OsStr,
175) -> Result<LinkResult> {
176    // HACK(eddyb) this is defined here to allow SPIR-T pretty-printing to apply
177    // to SPIR-V being dumped, outside of e.g. `--dump-spirt-passes`.
178    // FIXME(eddyb) this isn't used everywhere, sadly - to find those, search
179    // elsewhere for `.assemble()` and/or `spirv_tools::binary::from_binary`.
180    let spv_module_to_spv_words_and_spirt_module = |spv_module: &Module| {
181        let spv_words;
182        let spv_bytes = {
183            let _timer = sess.timer("assemble-to-spv_bytes-for-spirt");
184            spv_words = spv_module.assemble();
185            // FIXME(eddyb) this is wastefully cloning all the bytes, but also
186            // `spirt::Module` should have a method that takes `Vec<u32>`.
187            spirv_tools::binary::from_binary(&spv_words).to_vec()
188        };
189
190        // FIXME(eddyb) should've really been "spirt::Module::lower_from_spv_bytes".
191        let lower_from_spv_timer = sess.timer("spirt::Module::lower_from_spv_file");
192        let cx = std::rc::Rc::new(spirt::Context::new());
193        crate::custom_insts::register_to_spirt_context(&cx);
194        (
195            spv_words,
196            spirt::Module::lower_from_spv_bytes(cx, spv_bytes),
197            // HACK(eddyb) this is only returned for `SpirtDumpGuard`.
198            lower_from_spv_timer,
199        )
200    };
201
202    // FIXME(eddyb) deduplicate with `SpirtDumpGuard`.
203    let dump_spv_and_spirt = |spv_module: &Module, dump_file_path_stem: PathBuf| {
204        let (spv_words, spirt_module_or_err, _) =
205            spv_module_to_spv_words_and_spirt_module(spv_module);
206        std::fs::write(
207            dump_file_path_stem.with_extension("spv"),
208            spirv_tools::binary::from_binary(&spv_words),
209        )
210        .unwrap();
211
212        // FIXME(eddyb) reify SPIR-V -> SPIR-T errors so they're easier to debug.
213        if let Ok(mut module) = spirt_module_or_err {
214            // HACK(eddyb) avoid pretty-printing massive amounts of unused SPIR-T.
215            spirt::passes::link::minimize_exports(&mut module, |export_key| {
216                matches!(export_key, spirt::ExportKey::SpvEntryPoint { .. })
217            });
218
219            opts.spirt_cleanup_for_dumping(&mut module);
220
221            let pretty = spirt::print::Plan::for_module(&module).pretty_print();
222
223            // FIXME(eddyb) don't allocate whole `String`s here.
224            std::fs::write(
225                dump_file_path_stem.with_extension("spirt"),
226                pretty.to_string(),
227            )
228            .unwrap();
229            std::fs::write(
230                dump_file_path_stem.with_extension("spirt.html"),
231                pretty
232                    .render_to_html()
233                    .with_dark_mode_support()
234                    .to_html_doc(),
235            )
236            .unwrap();
237        }
238    };
239
240    let mut output = {
241        let _timer = sess.timer("link_merge");
242        // shift all the ids
243        let mut bound = inputs[0].header.as_ref().unwrap().bound - 1;
244        let version = inputs[0].header.as_ref().unwrap().version();
245
246        for module in inputs.iter_mut().skip(1) {
247            simple_passes::shift_ids(module, bound);
248            bound += module.header.as_ref().unwrap().bound - 1;
249            let this_version = module.header.as_ref().unwrap().version();
250            if version != this_version {
251                return Err(sess.dcx().err(format!(
252                    "cannot link two modules with different SPIR-V versions: v{}.{} and v{}.{}",
253                    version.0, version.1, this_version.0, this_version.1
254                )));
255            }
256        }
257
258        // merge the binaries
259        let mut loader = Loader::new();
260
261        for module in inputs {
262            module.all_inst_iter().for_each(|inst| {
263                loader.consume_instruction(inst.clone());
264            });
265        }
266
267        let mut output = loader.module();
268        let mut header = ModuleHeader::new(bound + 1);
269        header.set_version(version.0, version.1);
270        header.generator = 0x001B_0000;
271        output.header = Some(header);
272        output
273    };
274
275    if let Some(dir) = &opts.dump_post_merge {
276        dump_spv_and_spirt(&output, dir.join(disambiguated_crate_name_for_dumps));
277    }
278
279    // remove duplicates (https://github.com/KhronosGroup/SPIRV-Tools/blob/e7866de4b1dc2a7e8672867caeb0bdca49f458d3/source/opt/remove_duplicates_pass.cpp)
280    {
281        let _timer = sess.timer("link_remove_duplicates");
282        duplicates::remove_duplicate_extensions(&mut output);
283        duplicates::remove_duplicate_capabilities(&mut output);
284        duplicates::remove_duplicate_ext_inst_imports(&mut output);
285        duplicates::remove_duplicate_types(&mut output);
286        // jb-todo: strip identical OpDecoration / OpDecorationGroups
287    }
288
289    // find import / export pairs
290    {
291        let _timer = sess.timer("link_find_pairs");
292        import_export_link::run(opts, sess, &mut output)?;
293    }
294
295    {
296        let _timer = sess.timer("link_fragment_inst_check");
297        simple_passes::check_fragment_insts(sess, &output)?;
298    }
299
300    // HACK(eddyb) this has to run before the `report_and_remove_zombies` pass,
301    // so that any zombies that are passed as call arguments, but eventually unused,
302    // won't be (incorrectly) considered used.
303    {
304        let _timer = sess.timer("link_remove_unused_params");
305        output = param_weakening::remove_unused_params(output);
306    }
307
308    if opts.early_report_zombies {
309        // HACK(eddyb) `report_and_remove_zombies` is bad at determining whether
310        // some things are dead (such as whole blocks), and there's no reason to
311        // *not* run DCE, given SPIR-T exists and makes DCE mandatory, but we're
312        // still only going to do the minimum necessary ("block ordering").
313        {
314            let _timer = sess.timer("link_block_ordering_pass-before-report_and_remove_zombies");
315            for func in &mut output.functions {
316                simple_passes::block_ordering_pass(func);
317            }
318        }
319
320        let _timer = sess.timer("link_report_and_remove_zombies");
321        zombies::report_and_remove_zombies(sess, &mut output)?;
322    }
323
324    if opts.infer_storage_classes {
325        // HACK(eddyb) this is not the best approach, but storage class inference
326        // can still fail in entirely legitimate ways (i.e. mismatches in zombies).
327        if !opts.early_report_zombies {
328            let _timer = sess.timer("link_dce-before-specialize_generic_storage_class");
329            dce::dce(&mut output);
330        }
331
332        let _timer = sess.timer("specialize_generic_storage_class");
333        // HACK(eddyb) `specializer` requires functions' blocks to be in RPO order
334        // (i.e. `block_ordering_pass`) - this could be relaxed by using RPO visit
335        // inside `specializer`, but this is easier.
336        for func in &mut output.functions {
337            simple_passes::block_ordering_pass(func);
338        }
339        output = specializer::specialize(
340            opts,
341            output,
342            specializer::SimpleSpecialization {
343                specialize_operand: |operand| {
344                    matches!(operand, Operand::StorageClass(StorageClass::Generic))
345                },
346
347                // NOTE(eddyb) this can be anything that is guaranteed to pass
348                // validation - there are no constraints so this is either some
349                // unused pointer, or perhaps one created using `OpConstantNull`
350                // and simply never mixed with pointers that have a storage class.
351                // It would be nice to use `Generic` itself here so that we leave
352                // some kind of indication of it being unconstrained, but `Generic`
353                // requires additional capabilities, so we use `Function` instead.
354                // TODO(eddyb) investigate whether this can end up in a pointer
355                // type that's the value of a module-scoped variable, and whether
356                // `Function` is actually invalid! (may need `Private`)
357                concrete_fallback: Operand::StorageClass(StorageClass::Function),
358            },
359        );
360    }
361
362    // NOTE(eddyb) with SPIR-T, we can do `mem2reg` before inlining, too!
363    {
364        if opts.dce {
365            let _timer = sess.timer("link_dce-before-inlining");
366            dce::dce(&mut output);
367        }
368
369        let _timer = sess.timer("link_block_ordering_pass_and_mem2reg-before-inlining");
370        let mut pointer_to_pointee = FxHashMap::default();
371        let mut constants = FxHashMap::default();
372        let mut u32 = None;
373        for inst in &output.types_global_values {
374            match inst.class.opcode {
375                Op::TypePointer => {
376                    pointer_to_pointee
377                        .insert(inst.result_id.unwrap(), inst.operands[1].unwrap_id_ref());
378                }
379                Op::TypeInt
380                    if inst.operands[0].unwrap_literal_bit32() == 32
381                        && inst.operands[1].unwrap_literal_bit32() == 0 =>
382                {
383                    assert!(u32.is_none());
384                    u32 = Some(inst.result_id.unwrap());
385                }
386                Op::Constant if u32.is_some() && inst.result_type == u32 => {
387                    let value = inst.operands[0].unwrap_literal_bit32();
388                    constants.insert(inst.result_id.unwrap(), value);
389                }
390                _ => {}
391            }
392        }
393        for func in &mut output.functions {
394            simple_passes::block_ordering_pass(func);
395            // Note: mem2reg requires functions to be in RPO order (i.e. block_ordering_pass)
396            mem2reg::mem2reg(
397                output.header.as_mut().unwrap(),
398                &mut output.types_global_values,
399                &pointer_to_pointee,
400                &constants,
401                func,
402            );
403            destructure_composites::destructure_composites(func);
404        }
405    }
406
407    if opts.dce {
408        let _timer =
409            sess.timer("link_dce-and-remove_duplicate_debuginfo-after-mem2reg-before-inlining");
410        dce::dce(&mut output);
411        duplicates::remove_duplicate_debuginfo(&mut output);
412    }
413
414    if let Some(dir) = &opts.dump_pre_inline {
415        dump_spv_and_spirt(&output, dir.join(disambiguated_crate_name_for_dumps));
416    }
417
418    {
419        let _timer = sess.timer("link_inline");
420        inline::inline(sess, &mut output)?;
421    }
422
423    if opts.dce {
424        let _timer = sess.timer("link_dce-after-inlining");
425        dce::dce(&mut output);
426    }
427
428    // HACK(eddyb) this has to be after DCE, to not break SPIR-T w/ dead decorations.
429    if let Some(dir) = &opts.dump_post_inline {
430        dump_spv_and_spirt(&output, dir.join(disambiguated_crate_name_for_dumps));
431    }
432
433    {
434        let _timer = sess.timer("link_block_ordering_pass_and_mem2reg-after-inlining");
435        let mut pointer_to_pointee = FxHashMap::default();
436        let mut constants = FxHashMap::default();
437        let mut u32 = None;
438        for inst in &output.types_global_values {
439            match inst.class.opcode {
440                Op::TypePointer => {
441                    pointer_to_pointee
442                        .insert(inst.result_id.unwrap(), inst.operands[1].unwrap_id_ref());
443                }
444                Op::TypeInt
445                    if inst.operands[0].unwrap_literal_bit32() == 32
446                        && inst.operands[1].unwrap_literal_bit32() == 0 =>
447                {
448                    assert!(u32.is_none());
449                    u32 = Some(inst.result_id.unwrap());
450                }
451                Op::Constant if u32.is_some() && inst.result_type == u32 => {
452                    let value = inst.operands[0].unwrap_literal_bit32();
453                    constants.insert(inst.result_id.unwrap(), value);
454                }
455                _ => {}
456            }
457        }
458        for func in &mut output.functions {
459            simple_passes::block_ordering_pass(func);
460            // Note: mem2reg requires functions to be in RPO order (i.e. block_ordering_pass)
461            mem2reg::mem2reg(
462                output.header.as_mut().unwrap(),
463                &mut output.types_global_values,
464                &pointer_to_pointee,
465                &constants,
466                func,
467            );
468            destructure_composites::destructure_composites(func);
469        }
470    }
471
472    if opts.dce {
473        let _timer =
474            sess.timer("link_dce-and-remove_duplicate_debuginfo-after-mem2reg-after-inlining");
475        dce::dce(&mut output);
476        duplicates::remove_duplicate_debuginfo(&mut output);
477    }
478
479    {
480        let _timer = sess.timer("link_remove_non_uniform");
481        simple_passes::remove_non_uniform_decorations(sess, &mut output)?;
482    }
483
484    {
485        let _timer = sess.timer("link_remove_unused_type_capabilities");
486        simple_passes::remove_unused_type_capabilities(&mut output);
487    }
488
489    {
490        let _timer = sess.timer("link_type_capability_check");
491        simple_passes::check_type_capabilities(sess, &output)?;
492    }
493
494    // NOTE(eddyb) SPIR-T pipeline is entirely limited to this block.
495    {
496        let (spv_words, module_or_err, lower_from_spv_timer) =
497            spv_module_to_spv_words_and_spirt_module(&output);
498        let module = &mut module_or_err.map_err(|e| {
499            let spv_path = outputs.temp_path_for_diagnostic("spirt-lower-from-spv-input.spv");
500
501            let was_saved_msg =
502                match std::fs::write(&spv_path, spirv_tools::binary::from_binary(&spv_words)) {
503                    Ok(()) => format!("was saved to {}", spv_path.display()),
504                    Err(e) => format!("could not be saved: {e}"),
505                };
506
507            sess.dcx()
508                .struct_err(format!("{e}"))
509                .with_note("while lowering SPIR-V module to SPIR-T (spirt::spv::lower)")
510                .with_note(format!("input SPIR-V module {was_saved_msg}"))
511                .emit()
512        })?;
513
514        let mut dump_guard = SpirtDumpGuard {
515            sess,
516            linker_options: opts,
517            outputs,
518            disambiguated_crate_name_for_dumps,
519
520            module,
521            per_pass_module_for_dumping: vec![],
522            any_spirt_bugs: false,
523        };
524        let module = &mut *dump_guard.module;
525        // FIXME(eddyb) set the name into `dump_guard` to be able to access it on panic.
526        let before_pass = |pass| sess.timer(pass);
527        let mut after_pass = |pass, module: &spirt::Module, timer| {
528            drop(timer);
529            if opts.dump_spirt_passes.is_some() {
530                dump_guard
531                    .per_pass_module_for_dumping
532                    .push((pass, module.clone()));
533            }
534        };
535        // HACK(eddyb) don't dump the unstructured state if not requested, as
536        // after SPIR-T 0.4.0 it's extremely verbose (due to def-use hermeticity).
537        if opts.spirt_keep_unstructured_cfg_in_dumps || !opts.structurize {
538            after_pass("lower_from_spv", module, lower_from_spv_timer);
539        } else {
540            drop(lower_from_spv_timer);
541        }
542
543        // NOTE(eddyb) this *must* run on unstructured CFGs, to do its job.
544        // FIXME(eddyb) no longer relying on structurization, try porting this
545        // to replace custom aborts in `Block`s and inject `ExitInvocation`s
546        // after them (truncating the `Block` and/or parent region if necessary).
547        {
548            let _timer = before_pass(
549                "spirt_passes::controlflow::convert_custom_aborts_to_unstructured_returns_in_entry_points",
550            );
551            spirt_passes::controlflow::convert_custom_aborts_to_unstructured_returns_in_entry_points(opts, module);
552        }
553
554        if opts.structurize {
555            let timer = before_pass("spirt::legalize::structurize_func_cfgs");
556            spirt::passes::legalize::structurize_func_cfgs(module);
557            after_pass("structurize_func_cfgs", module, timer);
558        }
559
560        if !opts.spirt_passes.is_empty() {
561            // FIXME(eddyb) why does this focus on functions, it could just be module passes??
562            spirt_passes::run_func_passes(
563                module,
564                &opts.spirt_passes,
565                |name, _module| before_pass(name),
566                after_pass,
567            );
568        }
569
570        {
571            let _timer = before_pass("spirt_passes::diagnostics::report_diagnostics");
572            spirt_passes::diagnostics::report_diagnostics(sess, opts, module).map_err(
573                |spirt_passes::diagnostics::ReportedDiagnostics {
574                     rustc_errors_guarantee,
575                     any_errors_were_spirt_bugs,
576                 }| {
577                    dump_guard.any_spirt_bugs |= any_errors_were_spirt_bugs;
578                    rustc_errors_guarantee
579                },
580            )?;
581        }
582
583        // Replace our custom debuginfo instructions just before lifting to SPIR-V.
584        {
585            let _timer = before_pass("spirt_passes::debuginfo::convert_custom_debuginfo_to_spv");
586            spirt_passes::debuginfo::convert_custom_debuginfo_to_spv(module);
587        }
588
589        let spv_words = {
590            let _timer = before_pass("spirt::Module::lift_to_spv_module_emitter");
591            module.lift_to_spv_module_emitter().unwrap().words
592        };
593        // FIXME(eddyb) dump both SPIR-T and `spv_words` if there's an error here.
594        output = {
595            let _timer = sess.timer("parse-spv_words-from-spirt");
596            let mut loader = Loader::new();
597            rspirv::binary::parse_words(&spv_words, &mut loader).unwrap();
598            loader.module()
599        };
600    }
601
602    // Ensure that no references remain, to our custom "extended instruction set".
603    for inst in &output.ext_inst_imports {
604        assert_eq!(inst.class.opcode, Op::ExtInstImport);
605        let ext_inst_set = inst.operands[0].unwrap_literal_string();
606        if ext_inst_set.starts_with(custom_insts::CUSTOM_EXT_INST_SET_PREFIX) {
607            let expected = &custom_insts::CUSTOM_EXT_INST_SET[..];
608            if ext_inst_set == expected {
609                return Err(sess.dcx().err(format!(
610                    "`OpExtInstImport {ext_inst_set:?}` should not have been \
611                         left around after SPIR-T passes"
612                )));
613            } else {
614                return Err(sess.dcx().err(format!(
615                    "unsupported `OpExtInstImport {ext_inst_set:?}`
616                     (expected {expected:?} name - version mismatch?)"
617                )));
618            }
619        }
620    }
621
622    // FIXME(eddyb) rewrite these passes to SPIR-T ones, so we don't have to
623    // parse the output of `spirt::spv::lift` back into `rspirv` - also, for
624    // multi-module, it's much simpler with SPIR-T, just replace `module.exports`
625    // with a single-entry map, run `spirt::spv::lift` (or even `spirt::print`)
626    // on `module`, then put back the full original `module.exports` map.
627    {
628        let _timer = sess.timer("peephole_opts");
629        let types = peephole_opts::collect_types(&output);
630        for func in &mut output.functions {
631            peephole_opts::composite_construct(&types, func);
632            peephole_opts::vector_ops(output.header.as_mut().unwrap(), &types, func);
633            peephole_opts::bool_fusion(output.header.as_mut().unwrap(), &types, func);
634        }
635    }
636
637    {
638        let _timer = sess.timer("link_gather_all_interface_vars_from_uses");
639        entry_interface::gather_all_interface_vars_from_uses(&mut output);
640    }
641
642    if opts.spirv_metadata == SpirvMetadata::NameVariables {
643        let _timer = sess.timer("link_name_variables");
644        simple_passes::name_variables_pass(&mut output);
645    }
646
647    {
648        let _timer = sess.timer("link_sort_globals");
649        simple_passes::sort_globals(&mut output);
650    }
651
652    let mut output = if opts.module_output_type == ModuleOutputType::Multiple {
653        let mut file_stem_to_entry_name_and_module = BTreeMap::new();
654        for (i, entry) in output.entry_points.iter().enumerate() {
655            let mut module = output.clone();
656            module.entry_points.clear();
657            module.entry_points.push(entry.clone());
658            let entry_name = entry.operands[2].unwrap_literal_string().to_string();
659            let mut file_stem = OsString::from(
660                sanitize_filename::sanitize_with_options(
661                    &entry_name,
662                    sanitize_filename::Options {
663                        replacement: "-",
664                        ..Default::default()
665                    },
666                )
667                .replace("--", "-"),
668            );
669            // It's always possible to find an unambiguous `file_stem`, but it
670            // may take two tries (or more, in bizzare/adversarial cases).
671            let mut disambiguator = Some(i);
672            loop {
673                use std::collections::btree_map::Entry;
674                match file_stem_to_entry_name_and_module.entry(file_stem) {
675                    Entry::Vacant(entry) => {
676                        entry.insert((entry_name, module));
677                        break;
678                    }
679                    // FIXME(eddyb) false positive: `file_stem` was moved out of,
680                    // so assigning it is necessary, but clippy doesn't know that.
681                    #[allow(clippy::assigning_clones)]
682                    Entry::Occupied(entry) => {
683                        // FIXME(eddyb) there's no way to access the owned key
684                        // passed to `BTreeMap::entry` from `OccupiedEntry`.
685                        file_stem = entry.key().clone();
686                        file_stem.push(".");
687                        match disambiguator.take() {
688                            Some(d) => file_stem.push(d.to_string()),
689                            None => file_stem.push("next"),
690                        }
691                    }
692                }
693            }
694        }
695        LinkResult::MultipleModules {
696            file_stem_to_entry_name_and_module,
697        }
698    } else {
699        LinkResult::SingleModule(Box::new(output))
700    };
701
702    let output_module_iter = match &mut output {
703        LinkResult::SingleModule(m) => Either::Left(std::iter::once((None, &mut **m))),
704        LinkResult::MultipleModules {
705            file_stem_to_entry_name_and_module,
706        } => Either::Right(
707            file_stem_to_entry_name_and_module
708                .iter_mut()
709                .map(|(file_stem, (_, m))| (Some(file_stem), m)),
710        ),
711    };
712    for (file_stem, output) in output_module_iter {
713        if let Some(dir) = &opts.dump_post_split {
714            let mut file_name = disambiguated_crate_name_for_dumps.to_os_string();
715            if let Some(file_stem) = file_stem {
716                file_name.push(".");
717                file_name.push(file_stem);
718            }
719
720            dump_spv_and_spirt(output, dir.join(file_name));
721        }
722        // Run DCE again, even if module_output_type == ModuleOutputType::Multiple - the first DCE ran before
723        // structurization and mem2reg (for perf reasons), and mem2reg may remove references to
724        // invalid types, so we need to DCE again.
725        if opts.dce {
726            let _timer = sess.timer("link_dce_2");
727            dce::dce(output);
728        }
729
730        {
731            let _timer = sess.timer("link_remove_duplicate_debuginfo");
732            duplicates::remove_duplicate_debuginfo(output);
733        }
734
735        if opts.compact_ids {
736            let _timer = sess.timer("link_compact_ids");
737            // compact the ids https://github.com/KhronosGroup/SPIRV-Tools/blob/e02f178a716b0c3c803ce31b9df4088596537872/source/opt/compact_ids_pass.cpp#L43
738            output.header.as_mut().unwrap().bound = simple_passes::compact_ids(output);
739        };
740
741        // FIXME(eddyb) convert these into actual `OpLine`s with a SPIR-T pass,
742        // but that'd require keeping the modules in SPIR-T form (once lowered),
743        // and never loading them back into `rspirv` once lifted back to SPIR-V.
744        SrcLocDecoration::remove_all(output);
745
746        // FIXME(eddyb) might make more sense to rewrite these away on SPIR-T.
747        ZombieDecoration::remove_all(output);
748    }
749
750    Ok(output)
751}
752
753/// Helper for dumping SPIR-T on drop, which allows panics to also dump,
754/// not just successful compilation (i.e. via `--dump-spirt-passes`).
755struct SpirtDumpGuard<'a> {
756    sess: &'a Session,
757    linker_options: &'a Options,
758    outputs: &'a OutputFilenames,
759    disambiguated_crate_name_for_dumps: &'a OsStr,
760
761    module: &'a mut spirt::Module,
762    per_pass_module_for_dumping: Vec<(&'static str, spirt::Module)>,
763    any_spirt_bugs: bool,
764}
765
766impl Drop for SpirtDumpGuard<'_> {
767    fn drop(&mut self) {
768        self.any_spirt_bugs |= std::thread::panicking();
769
770        let mut dump_spirt_file_path =
771            self.linker_options
772                .dump_spirt_passes
773                .as_ref()
774                .map(|dump_dir| {
775                    dump_dir
776                        .join(self.disambiguated_crate_name_for_dumps)
777                        .with_extension("spirt")
778                });
779
780        // FIXME(eddyb) this won't allow seeing the individual passes, but it's
781        // better than nothing (theoretically the whole "SPIR-T pipeline" could
782        // be put in a loop so that everything is redone with per-pass tracking,
783        // but that requires keeping around e.g. the initial SPIR-V for longer,
784        // and probably invoking the "SPIR-T pipeline" here, as looping is hard).
785        if self.any_spirt_bugs && dump_spirt_file_path.is_none() {
786            if self.per_pass_module_for_dumping.is_empty() {
787                self.per_pass_module_for_dumping
788                    .push(("", self.module.clone()));
789            }
790            dump_spirt_file_path = Some(self.outputs.temp_path_for_diagnostic("spirt"));
791        }
792
793        if let Some(dump_spirt_file_path) = &dump_spirt_file_path {
794            for (_, module) in &mut self.per_pass_module_for_dumping {
795                self.linker_options.spirt_cleanup_for_dumping(module);
796            }
797
798            // FIXME(eddyb) catch panics during pretty-printing itself, and
799            // tell the user to use `--dump-spirt-passes` (and resolve the
800            // second FIXME below so it does anything) - also, that may need
801            // quieting the panic handler, likely controlled by a `thread_local!`
802            // (while the panic handler is global), but that would be useful
803            // for collecting a panic message (assuming any of this is worth it).
804            // FIXME(eddyb) when per-pass versions are available, try building
805            // plans for individual versions, or maybe repeat `Plan::for_versions`
806            // without the last version if it initially panicked?
807            let plan = spirt::print::Plan::for_versions(
808                self.module.cx_ref(),
809                self.per_pass_module_for_dumping
810                    .iter()
811                    .map(|(pass, module)| (format!("after {pass}"), module)),
812            );
813            let pretty = plan.pretty_print();
814
815            // FIXME(eddyb) don't allocate whole `String`s here.
816            std::fs::write(dump_spirt_file_path, pretty.to_string()).unwrap();
817            std::fs::write(
818                dump_spirt_file_path.with_extension("spirt.html"),
819                pretty
820                    .render_to_html()
821                    .with_dark_mode_support()
822                    .to_html_doc(),
823            )
824            .unwrap();
825            if self.any_spirt_bugs {
826                let mut note = self.sess.dcx().struct_note("SPIR-T bugs were encountered");
827                note.help(format!(
828                    "pretty-printed SPIR-T was saved to {}.html",
829                    dump_spirt_file_path.display()
830                ));
831                if self.linker_options.dump_spirt_passes.is_none() {
832                    note.help("re-run with `RUSTGPU_CODEGEN_ARGS=\"--dump-spirt-passes=$PWD\"` for more details");
833                }
834                note.note("pretty-printed SPIR-T is preferred when reporting Rust-GPU issues");
835                note.emit();
836            }
837        }
838    }
839}