Add name_variables option to spirv-builder (#660)

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -277,6 +277,8 @@ pub struct CodegenArgs {
     pub disassemble_entry: Option<String>,
     pub disassemble_globals: bool,
 
+    pub name_variables: bool,
+
     // spirv-val flags
     pub relax_struct_store: bool,
     pub relax_logical_pointer: bool,
@@ -313,6 +315,13 @@ impl CodegenArgs {
         );
         opts.optflagopt("", "disassemble-globals", "print globals to stderr", "");
 
+        opts.optflagopt(
+            "",
+            "name-variables",
+            "Keep OpName for OpVariables, strip all others.",
+            "",
+        );
+
         opts.optflagopt("", "relax-struct-store", "Allow store from one struct type to a different type with compatible layout and members.", "");
         opts.optflagopt("", "relax-logical-pointer", "Allow allocating an object of a pointer type and returning a pointer value from a function in logical addressing mode", "");
         opts.optflagopt("", "relax-block-layout", "Enable VK_KHR_relaxed_block_layout when checking standard uniform, storage buffer, and push constant layouts. This is the default when targeting Vulkan 1.1 or later.", "");
@@ -328,6 +337,8 @@ impl CodegenArgs {
         let disassemble_entry = matches.opt_str("disassemble-entry");
         let disassemble_globals = matches.opt_present("disassemble-globals");
 
+        let name_variables = matches.opt_present("name-variables");
+
         let relax_struct_store = matches.opt_present("relax-struct-store");
         let relax_logical_pointer = matches.opt_present("relax-logical-pointer");
         let relax_block_layout = matches.opt_present("relax-block-layout");
@@ -344,6 +355,8 @@ impl CodegenArgs {
             disassemble_entry,
             disassemble_globals,
 
+            name_variables,
+
             relax_struct_store,
             relax_logical_pointer,
             relax_block_layout,

crates/rustc_codegen_spirv/src/link.rs

@@ -137,10 +137,17 @@ fn link_exe(
         codegen_results,
     );
 
-    let spv_binary = do_link(sess, &objects, &rlibs, legalize, emit_multiple_modules);
-
     let cg_args = crate::codegen_cx::CodegenArgs::from_session(sess);
 
+    let spv_binary = do_link(
+        sess,
+        &cg_args,
+        &objects,
+        &rlibs,
+        legalize,
+        emit_multiple_modules,
+    );
+
     let mut root_file_name = out_filename.file_name().unwrap().to_owned();
     root_file_name.push(".dir");
     let out_dir = out_filename.with_file_name(root_file_name);
@@ -221,10 +228,11 @@ fn post_link_single_module(
         preserve_spec_constants: false,
     };
 
-    let spv_binary = if sess.opts.optimize != OptLevel::No || sess.opts.debuginfo == DebugInfo::None
+    let spv_binary = if sess.opts.optimize != OptLevel::No
+        || (sess.opts.debuginfo == DebugInfo::None && !cg_args.name_variables)
     {
         let _timer = sess.timer("link_spirv_opt");
-        do_spirv_opt(sess, spv_binary, out_filename, opt_options)
+        do_spirv_opt(sess, cg_args, spv_binary, out_filename, opt_options)
     } else {
         spv_binary
     };
@@ -249,6 +257,7 @@ fn post_link_single_module(
 
 fn do_spirv_opt(
     sess: &Session,
+    cg_args: &crate::codegen_cx::CodegenArgs,
     spv_binary: Vec<u32>,
     filename: &Path,
     options: spirv_tools::opt::Options,
@@ -270,7 +279,7 @@ fn do_spirv_opt(
         }
     }
 
-    if sess.opts.debuginfo == DebugInfo::None {
+    if sess.opts.debuginfo == DebugInfo::None && !cg_args.name_variables {
         optimizer
             .register_pass(opt::Passes::EliminateDeadConstant)
             .register_pass(opt::Passes::StripDebugInfo);
@@ -484,6 +493,7 @@ pub fn read_metadata(rlib: &Path) -> Result<MetadataRef, String> {
 /// shenanigans to collect all the object files we need to link.
 fn do_link(
     sess: &Session,
+    cg_args: &crate::codegen_cx::CodegenArgs,
     objects: &[PathBuf],
     rlibs: &[PathBuf],
     legalize: bool,
@@ -540,6 +550,7 @@ fn do_link(
         mem2reg: legalize,
         structurize: env::var("NO_STRUCTURIZE").is_err(),
         emit_multiple_modules,
+        name_variables: cg_args.name_variables,
     };
 
     let link_result = linker::link(sess, modules, &options);

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -29,6 +29,7 @@ pub struct Options {
     pub mem2reg: bool,
     pub structurize: bool,
     pub emit_multiple_modules: bool,
+    pub name_variables: bool,
 }
 
 pub enum LinkResult {
@@ -251,6 +252,11 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
         duplicates::remove_duplicate_lines(&mut output);
     }
 
+    if opts.name_variables {
+        let _timer = sess.timer("link_name_variables");
+        simple_passes::name_variables_pass(&mut output);
+    }
+
     {
         let _timer = sess.timer("link_sort_globals");
         simple_passes::sort_globals(&mut output);

crates/rustc_codegen_spirv/src/linker/simple_passes.rs

@@ -129,3 +129,25 @@ pub fn sort_globals(module: &mut Module) {
     // have a function declaration without a body in a fully linked module?
     module.functions.sort_by_key(|f| !f.blocks.is_empty());
 }
+
+pub fn name_variables_pass(module: &mut Module) {
+    let variables = module
+        .types_global_values
+        .iter()
+        .filter(|inst| inst.class.opcode == Op::Variable)
+        .map(|inst| inst.result_id.unwrap())
+        .collect::<FxHashSet<Word>>();
+    module
+        .debug_names
+        .retain(|inst| variables.contains(&inst.operands[0].unwrap_id_ref()));
+    module
+        .types_global_values
+        .retain(|inst| inst.class.opcode != Op::Line);
+    for func in &mut module.functions {
+        for block in &mut func.blocks {
+            block
+                .instructions
+                .retain(|inst| inst.class.opcode != Op::Line);
+        }
+    }
+}

crates/rustc_codegen_spirv/src/linker/test.rs

@@ -95,6 +95,7 @@ fn assemble_and_link(binaries: &[&[u8]]) -> Result<Module, String> {
                 mem2reg: false,
                 structurize: false,
                 emit_multiple_modules: false,
+                name_variables: false,
             },
         );
         assert_eq!(compiler.session().has_errors(), res.is_err());

crates/spirv-builder/src/lib.rs

@@ -133,6 +133,7 @@ pub struct SpirvBuilder {
     target: String,
     bindless: bool,
     multimodule: bool,
+    name_variables: bool,
     capabilities: Vec<Capability>,
     extensions: Vec<String>,
 
@@ -154,6 +155,7 @@ impl SpirvBuilder {
             target: target.into(),
             bindless: false,
             multimodule: false,
+            name_variables: false,
             capabilities: Vec::new(),
             extensions: Vec::new(),
 
@@ -193,6 +195,15 @@ impl SpirvBuilder {
         self
     }
 
+    /// Keep `OpName` reflection information for global `OpVariable`s (which means things like
+    /// uniforms and shader input/outputs) but strip `OpName`s for everything else (functions,
+    /// types, and so on). This is useful if you want a small binary size without debugging
+    /// information, but need variable name reflection to work.
+    pub fn name_variables(mut self, v: bool) -> Self {
+        self.name_variables = v;
+        self
+    }
+
     /// Adds a capability to the SPIR-V module. Checking if a capability is enabled in code can be
     /// done via `#[cfg(target_feature = "TheCapability")]`.
     pub fn capability(mut self, capability: Capability) -> Self {
@@ -352,6 +363,9 @@ fn invoke_rustc(builder: &SpirvBuilder) -> Result<PathBuf, SpirvBuilderError> {
     if builder.multimodule {
         llvm_args.push("--module-output=multiple");
     }
+    if builder.name_variables {
+        llvm_args.push("--name-variables");
+    }
     if builder.relax_struct_store {
         llvm_args.push("--relax-struct-store");
     }