linker/test: prepare for SPIR-T compatibility.

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -515,6 +515,7 @@ impl CodegenArgs {
             // FIXME(eddyb) deduplicate between `CodegenArgs` and `linker::Options`.
             emit_multiple_modules: module_output_type == ModuleOutputType::Multiple,
             spirv_metadata,
+            keep_link_exports: false,
 
             // NOTE(eddyb) these are debugging options that used to be env vars
             // (for more information see `docs/src/codegen-args.md`).

crates/rustc_codegen_spirv/src/linker/dce.rs

@@ -8,7 +8,7 @@
 //! concept.
 
 use rspirv::dr::{Function, Instruction, Module, Operand};
-use rspirv::spirv::{Op, StorageClass, Word};
+use rspirv::spirv::{Decoration, LinkageType, Op, StorageClass, Word};
 use rustc_data_structures::fx::FxHashSet;
 
 pub fn dce(module: &mut Module) {
@@ -19,9 +19,23 @@ pub fn dce(module: &mut Module) {
 
 pub fn collect_roots(module: &Module) -> FxHashSet<Word> {
     let mut rooted = FxHashSet::default();
+
     for inst in &module.entry_points {
         root(inst, &mut rooted);
     }
+
+    // NOTE(eddyb) such "link exports" roots are only relevant when `Options`'s
+    // `keep_link_export`s field is used to request that `Export`s are left in
+    // (primarily for unit testing - see also its doc comment).
+    for inst in &module.annotations {
+        if inst.class.opcode == Op::Decorate
+            && inst.operands[1].unwrap_decoration() == Decoration::LinkageAttributes
+            && inst.operands[3].unwrap_linkage_type() == LinkageType::Export
+        {
+            root(inst, &mut rooted);
+        }
+    }
+
     rooted
 }
 

crates/rustc_codegen_spirv/src/linker/import_export_link.rs

@@ -5,10 +5,10 @@ use rspirv::spirv::{Capability, Decoration, LinkageType, Op, Word};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_session::Session;
 
-pub fn run(sess: &Session, module: &mut Module) -> Result<()> {
+pub fn run(opts: &super::Options, sess: &Session, module: &mut Module) -> Result<()> {
     let (rewrite_rules, killed_parameters) =
         find_import_export_pairs_and_killed_params(sess, module)?;
-    kill_linkage_instructions(module, &rewrite_rules);
+    kill_linkage_instructions(opts, module, &rewrite_rules);
     import_kill_annotations_and_debug(module, &rewrite_rules, &killed_parameters);
     replace_all_uses_with(module, &rewrite_rules);
     Ok(())
@@ -215,7 +215,11 @@ fn replace_all_uses_with(module: &mut Module, rules: &FxHashMap<u32, u32>) {
     });
 }
 
-fn kill_linkage_instructions(module: &mut Module, rewrite_rules: &FxHashMap<u32, u32>) {
+fn kill_linkage_instructions(
+    opts: &super::Options,
+    module: &mut Module,
+    rewrite_rules: &FxHashMap<u32, u32>,
+) {
     // drop imported functions
     module
         .functions
@@ -227,16 +231,27 @@ fn kill_linkage_instructions(module: &mut Module, rewrite_rules: &FxHashMap<u32,
             .map_or(true, |v| !rewrite_rules.contains_key(&v))
     });
 
+    // NOTE(eddyb) `Options`'s `keep_link_export`s field requests that `Export`s
+    // are left in (primarily for unit testing - see also its doc comment).
+    let mut kept_any_linkage_decorations = false;
     module.annotations.retain(|inst| {
-        inst.class.opcode != Op::Decorate
-            || inst.operands[1].unwrap_decoration() != Decoration::LinkageAttributes
-    });
-
-    // drop OpCapability Linkage
-    module.capabilities.retain(|inst| {
-        inst.class.opcode != Op::Capability
-            || inst.operands[0].unwrap_capability() != Capability::Linkage
+        !(inst.class.opcode == Op::Decorate
+            && inst.operands[1].unwrap_decoration() == Decoration::LinkageAttributes
+            && match inst.operands[3].unwrap_linkage_type() {
+                LinkageType::Export if opts.keep_link_exports => {
+                    kept_any_linkage_decorations = true;
+                    false
+                }
+                _ => true,
+            })
     });
+    if !kept_any_linkage_decorations {
+        // drop OpCapability Linkage
+        module.capabilities.retain(|inst| {
+            inst.class.opcode != Op::Capability
+                || inst.operands[0].unwrap_capability() != Capability::Linkage
+        });
+    }
 }
 
 fn import_kill_annotations_and_debug(

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -38,6 +38,12 @@ pub struct Options {
     pub emit_multiple_modules: bool,
     pub spirv_metadata: SpirvMetadata,
 
+    /// Whether to preserve `LinkageAttributes "..." Export` decorations,
+    /// even after resolving imports to exports.
+    ///
+    /// **Note**: currently only used for unit testing, and not exposed elsewhere.
+    pub keep_link_exports: bool,
+
     // NOTE(eddyb) these are debugging options that used to be env vars
     // (for more information see `docs/src/codegen-args.md`).
     pub dump_post_merge: Option<PathBuf>,
@@ -176,7 +182,7 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
     // find import / export pairs
     {
         let _timer = sess.timer("link_find_pairs");
-        import_export_link::run(sess, &mut output)?;
+        import_export_link::run(opts, sess, &mut output)?;
     }
 
     {