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    // NOTE(eddyb) SPIR-T pipeline is entirely limited to this block.
485    {
486        let (spv_words, module_or_err, lower_from_spv_timer) =
487            spv_module_to_spv_words_and_spirt_module(&output);
488        let module = &mut module_or_err.map_err(|e| {
489            let spv_path = outputs.temp_path_for_diagnostic("spirt-lower-from-spv-input.spv");
490
491            let was_saved_msg =
492                match std::fs::write(&spv_path, spirv_tools::binary::from_binary(&spv_words)) {
493                    Ok(()) => format!("was saved to {}", spv_path.display()),
494                    Err(e) => format!("could not be saved: {e}"),
495                };
496
497            sess.dcx()
498                .struct_err(format!("{e}"))
499                .with_note("while lowering SPIR-V module to SPIR-T (spirt::spv::lower)")
500                .with_note(format!("input SPIR-V module {was_saved_msg}"))
501                .emit()
502        })?;
503
504        let mut dump_guard = SpirtDumpGuard {
505            sess,
506            linker_options: opts,
507            outputs,
508            disambiguated_crate_name_for_dumps,
509
510            module,
511            per_pass_module_for_dumping: vec![],
512            any_spirt_bugs: false,
513        };
514        let module = &mut *dump_guard.module;
515        // FIXME(eddyb) set the name into `dump_guard` to be able to access it on panic.
516        let before_pass = |pass| sess.timer(pass);
517        let mut after_pass = |pass, module: &spirt::Module, timer| {
518            drop(timer);
519            if opts.dump_spirt_passes.is_some() {
520                dump_guard
521                    .per_pass_module_for_dumping
522                    .push((pass, module.clone()));
523            }
524        };
525        // HACK(eddyb) don't dump the unstructured state if not requested, as
526        // after SPIR-T 0.4.0 it's extremely verbose (due to def-use hermeticity).
527        if opts.spirt_keep_unstructured_cfg_in_dumps || !opts.structurize {
528            after_pass("lower_from_spv", module, lower_from_spv_timer);
529        } else {
530            drop(lower_from_spv_timer);
531        }
532
533        // NOTE(eddyb) this *must* run on unstructured CFGs, to do its job.
534        // FIXME(eddyb) no longer relying on structurization, try porting this
535        // to replace custom aborts in `Block`s and inject `ExitInvocation`s
536        // after them (truncating the `Block` and/or parent region if necessary).
537        {
538            let _timer = before_pass(
539                "spirt_passes::controlflow::convert_custom_aborts_to_unstructured_returns_in_entry_points",
540            );
541            spirt_passes::controlflow::convert_custom_aborts_to_unstructured_returns_in_entry_points(opts, module);
542        }
543
544        if opts.structurize {
545            let timer = before_pass("spirt::legalize::structurize_func_cfgs");
546            spirt::passes::legalize::structurize_func_cfgs(module);
547            after_pass("structurize_func_cfgs", module, timer);
548        }
549
550        if !opts.spirt_passes.is_empty() {
551            // FIXME(eddyb) why does this focus on functions, it could just be module passes??
552            spirt_passes::run_func_passes(
553                module,
554                &opts.spirt_passes,
555                |name, _module| before_pass(name),
556                &mut after_pass,
557            );
558        }
559
560        {
561            let timer = before_pass("spirt_passes::validate");
562            spirt_passes::validate::validate(module);
563            after_pass("validate", module, timer);
564        }
565
566        {
567            let _timer = before_pass("spirt_passes::diagnostics::report_diagnostics");
568            spirt_passes::diagnostics::report_diagnostics(sess, opts, module).map_err(
569                |spirt_passes::diagnostics::ReportedDiagnostics {
570                     rustc_errors_guarantee,
571                     any_errors_were_spirt_bugs,
572                 }| {
573                    dump_guard.any_spirt_bugs |= any_errors_were_spirt_bugs;
574                    rustc_errors_guarantee
575                },
576            )?;
577        }
578
579        // Replace our custom debuginfo instructions just before lifting to SPIR-V.
580        {
581            let _timer = before_pass("spirt_passes::debuginfo::convert_custom_debuginfo_to_spv");
582            spirt_passes::debuginfo::convert_custom_debuginfo_to_spv(module);
583        }
584
585        let spv_words = {
586            let _timer = before_pass("spirt::Module::lift_to_spv_module_emitter");
587            module.lift_to_spv_module_emitter().unwrap().words
588        };
589        // FIXME(eddyb) dump both SPIR-T and `spv_words` if there's an error here.
590        output = {
591            let _timer = sess.timer("parse-spv_words-from-spirt");
592            let mut loader = Loader::new();
593            rspirv::binary::parse_words(&spv_words, &mut loader).unwrap();
594            loader.module()
595        };
596    }
597
598    // Ensure that no references remain, to our custom "extended instruction set".
599    for inst in &output.ext_inst_imports {
600        assert_eq!(inst.class.opcode, Op::ExtInstImport);
601        let ext_inst_set = inst.operands[0].unwrap_literal_string();
602        if ext_inst_set.starts_with(custom_insts::CUSTOM_EXT_INST_SET_PREFIX) {
603            let expected = &custom_insts::CUSTOM_EXT_INST_SET[..];
604            if ext_inst_set == expected {
605                return Err(sess.dcx().err(format!(
606                    "`OpExtInstImport {ext_inst_set:?}` should not have been \
607                         left around after SPIR-T passes"
608                )));
609            } else {
610                return Err(sess.dcx().err(format!(
611                    "unsupported `OpExtInstImport {ext_inst_set:?}`
612                     (expected {expected:?} name - version mismatch?)"
613                )));
614            }
615        }
616    }
617
618    // FIXME(eddyb) rewrite these passes to SPIR-T ones, so we don't have to
619    // parse the output of `spirt::spv::lift` back into `rspirv` - also, for
620    // multi-module, it's much simpler with SPIR-T, just replace `module.exports`
621    // with a single-entry map, run `spirt::spv::lift` (or even `spirt::print`)
622    // on `module`, then put back the full original `module.exports` map.
623    {
624        let _timer = sess.timer("peephole_opts");
625        let types = peephole_opts::collect_types(&output);
626        for func in &mut output.functions {
627            peephole_opts::composite_construct(&types, func);
628            peephole_opts::vector_ops(output.header.as_mut().unwrap(), &types, func);
629            peephole_opts::bool_fusion(output.header.as_mut().unwrap(), &types, func);
630        }
631    }
632
633    {
634        let _timer = sess.timer("link_remove_unused_type_capabilities");
635        simple_passes::remove_unused_type_capabilities(&mut output);
636    }
637
638    {
639        let _timer = sess.timer("link_gather_all_interface_vars_from_uses");
640        entry_interface::gather_all_interface_vars_from_uses(&mut output);
641    }
642
643    if opts.spirv_metadata == SpirvMetadata::NameVariables {
644        let _timer = sess.timer("link_name_variables");
645        simple_passes::name_variables_pass(&mut output);
646    }
647
648    {
649        let _timer = sess.timer("link_sort_globals");
650        simple_passes::sort_globals(&mut output);
651    }
652
653    let mut output = if opts.module_output_type == ModuleOutputType::Multiple {
654        let mut file_stem_to_entry_name_and_module = BTreeMap::new();
655        for (i, entry) in output.entry_points.iter().enumerate() {
656            let mut module = output.clone();
657            module.entry_points.clear();
658            module.entry_points.push(entry.clone());
659            let entry_name = entry.operands[2].unwrap_literal_string().to_string();
660            let mut file_stem = OsString::from(
661                sanitize_filename::sanitize_with_options(
662                    &entry_name,
663                    sanitize_filename::Options {
664                        replacement: "-",
665                        ..Default::default()
666                    },
667                )
668                .replace("--", "-"),
669            );
670            // It's always possible to find an unambiguous `file_stem`, but it
671            // may take two tries (or more, in bizzare/adversarial cases).
672            let mut disambiguator = Some(i);
673            loop {
674                use std::collections::btree_map::Entry;
675                match file_stem_to_entry_name_and_module.entry(file_stem) {
676                    Entry::Vacant(entry) => {
677                        entry.insert((entry_name, module));
678                        break;
679                    }
680                    // FIXME(eddyb) false positive: `file_stem` was moved out of,
681                    // so assigning it is necessary, but clippy doesn't know that.
682                    #[allow(clippy::assigning_clones)]
683                    Entry::Occupied(entry) => {
684                        // FIXME(eddyb) there's no way to access the owned key
685                        // passed to `BTreeMap::entry` from `OccupiedEntry`.
686                        file_stem = entry.key().clone();
687                        file_stem.push(".");
688                        match disambiguator.take() {
689                            Some(d) => file_stem.push(d.to_string()),
690                            None => file_stem.push("next"),
691                        }
692                    }
693                }
694            }
695        }
696        LinkResult::MultipleModules {
697            file_stem_to_entry_name_and_module,
698        }
699    } else {
700        LinkResult::SingleModule(Box::new(output))
701    };
702
703    let output_module_iter = match &mut output {
704        LinkResult::SingleModule(m) => Either::Left(std::iter::once((None, &mut **m))),
705        LinkResult::MultipleModules {
706            file_stem_to_entry_name_and_module,
707        } => Either::Right(
708            file_stem_to_entry_name_and_module
709                .iter_mut()
710                .map(|(file_stem, (_, m))| (Some(file_stem), m)),
711        ),
712    };
713    for (file_stem, output) in output_module_iter {
714        if let Some(dir) = &opts.dump_post_split {
715            let mut file_name = disambiguated_crate_name_for_dumps.to_os_string();
716            if let Some(file_stem) = file_stem {
717                file_name.push(".");
718                file_name.push(file_stem);
719            }
720
721            dump_spv_and_spirt(output, dir.join(file_name));
722        }
723        // Run DCE again, even if module_output_type == ModuleOutputType::Multiple - the first DCE ran before
724        // structurization and mem2reg (for perf reasons), and mem2reg may remove references to
725        // invalid types, so we need to DCE again.
726        if opts.dce {
727            let _timer = sess.timer("link_dce_2");
728            dce::dce(output);
729        }
730
731        {
732            let _timer = sess.timer("link_remove_duplicate_debuginfo");
733            duplicates::remove_duplicate_debuginfo(output);
734        }
735
736        if opts.compact_ids {
737            let _timer = sess.timer("link_compact_ids");
738            // compact the ids https://github.com/KhronosGroup/SPIRV-Tools/blob/e02f178a716b0c3c803ce31b9df4088596537872/source/opt/compact_ids_pass.cpp#L43
739            output.header.as_mut().unwrap().bound = simple_passes::compact_ids(output);
740        };
741
742        // FIXME(eddyb) convert these into actual `OpLine`s with a SPIR-T pass,
743        // but that'd require keeping the modules in SPIR-T form (once lowered),
744        // and never loading them back into `rspirv` once lifted back to SPIR-V.
745        SrcLocDecoration::remove_all(output);
746
747        // FIXME(eddyb) might make more sense to rewrite these away on SPIR-T.
748        ZombieDecoration::remove_all(output);
749    }
750
751    Ok(output)
752}
753
754/// Helper for dumping SPIR-T on drop, which allows panics to also dump,
755/// not just successful compilation (i.e. via `--dump-spirt-passes`).
756struct SpirtDumpGuard<'a> {
757    sess: &'a Session,
758    linker_options: &'a Options,
759    outputs: &'a OutputFilenames,
760    disambiguated_crate_name_for_dumps: &'a OsStr,
761
762    module: &'a mut spirt::Module,
763    per_pass_module_for_dumping: Vec<(&'static str, spirt::Module)>,
764    any_spirt_bugs: bool,
765}
766
767impl Drop for SpirtDumpGuard<'_> {
768    fn drop(&mut self) {
769        self.any_spirt_bugs |= std::thread::panicking();
770
771        let mut dump_spirt_file_path =
772            self.linker_options
773                .dump_spirt_passes
774                .as_ref()
775                .map(|dump_dir| {
776                    dump_dir
777                        .join(self.disambiguated_crate_name_for_dumps)
778                        .with_extension("spirt")
779                });
780
781        // FIXME(eddyb) this won't allow seeing the individual passes, but it's
782        // better than nothing (theoretically the whole "SPIR-T pipeline" could
783        // be put in a loop so that everything is redone with per-pass tracking,
784        // but that requires keeping around e.g. the initial SPIR-V for longer,
785        // and probably invoking the "SPIR-T pipeline" here, as looping is hard).
786        if self.any_spirt_bugs && dump_spirt_file_path.is_none() {
787            if self.per_pass_module_for_dumping.is_empty() {
788                self.per_pass_module_for_dumping
789                    .push(("", self.module.clone()));
790            }
791            dump_spirt_file_path = Some(self.outputs.temp_path_for_diagnostic("spirt"));
792        }
793
794        if let Some(dump_spirt_file_path) = &dump_spirt_file_path {
795            for (_, module) in &mut self.per_pass_module_for_dumping {
796                self.linker_options.spirt_cleanup_for_dumping(module);
797            }
798
799            // FIXME(eddyb) catch panics during pretty-printing itself, and
800            // tell the user to use `--dump-spirt-passes` (and resolve the
801            // second FIXME below so it does anything) - also, that may need
802            // quieting the panic handler, likely controlled by a `thread_local!`
803            // (while the panic handler is global), but that would be useful
804            // for collecting a panic message (assuming any of this is worth it).
805            // FIXME(eddyb) when per-pass versions are available, try building
806            // plans for individual versions, or maybe repeat `Plan::for_versions`
807            // without the last version if it initially panicked?
808            let plan = spirt::print::Plan::for_versions(
809                self.module.cx_ref(),
810                self.per_pass_module_for_dumping
811                    .iter()
812                    .map(|(pass, module)| (format!("after {pass}"), module)),
813            );
814            let pretty = plan.pretty_print();
815
816            // FIXME(eddyb) don't allocate whole `String`s here.
817            std::fs::write(dump_spirt_file_path, pretty.to_string()).unwrap();
818            std::fs::write(
819                dump_spirt_file_path.with_extension("spirt.html"),
820                pretty
821                    .render_to_html()
822                    .with_dark_mode_support()
823                    .to_html_doc(),
824            )
825            .unwrap();
826            if self.any_spirt_bugs {
827                let mut note = self.sess.dcx().struct_note("SPIR-T bugs were encountered");
828                note.help(format!(
829                    "pretty-printed SPIR-T was saved to {}.html",
830                    dump_spirt_file_path.display()
831                ));
832                if self.linker_options.dump_spirt_passes.is_none() {
833                    note.help("re-run with `RUSTGPU_CODEGEN_ARGS=\"--dump-spirt-passes=$PWD\"` for more details");
834                }
835                note.note("pretty-printed SPIR-T is preferred when reporting Rust-GPU issues");
836                note.emit();
837            }
838        }
839    }
840}