dedup builtins: add linker pass remove_duplicate_builtin_variables

Author: fluffysquirrels

crates/rustc_codegen_spirv/src/linker/duplicates.rs

@@ -1,7 +1,7 @@
 use crate::custom_insts::{self, CustomOp};
 use rspirv::binary::Assemble;
 use rspirv::dr::{Instruction, Module, Operand};
-use rspirv::spirv::{Op, Word};
+use rspirv::spirv::{BuiltIn, Decoration, Op, Word};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_middle::bug;
 use smallvec::SmallVec;
@@ -104,6 +104,7 @@ fn gather_annotations(annotations: &[Instruction]) -> FxHashMap<Word, Vec<u32>>
         .collect()
 }
 
+/// Returns a map from an ID to its debug name (given by `OpName`).
 fn gather_names(debug_names: &[Instruction]) -> FxHashMap<Word, String> {
     debug_names
         .iter()
@@ -173,17 +174,34 @@ fn make_dedupe_key(
 }
 
 fn rewrite_inst_with_rules(inst: &mut Instruction, rules: &FxHashMap<u32, u32>) {
-    if let Some(ref mut id) = inst.result_type {
+    if let Some(ref mut id) = inst.result_id {
         // If the rewrite rules contain this ID, replace with the mapped value, otherwise don't touch it.
         *id = rules.get(id).copied().unwrap_or(*id);
     }
+    if let Some(ref mut type_id) = inst.result_type {
+        *type_id = rules.get(type_id).copied().unwrap_or(*type_id);
+    }
     for op in &mut inst.operands {
         if let Some(id) = op.id_ref_any_mut() {
             *id = rules.get(id).copied().unwrap_or(*id);
         }
     }
 }
 
+/// Remove duplicate `OpName` and `OpMemberName` instructions from module debug names section.
+fn remove_duplicate_debug_names(debug_names: &mut Vec<Instruction>) {
+    let mut name_ids = FxHashSet::default();
+    let mut member_name_ids = FxHashSet::default();
+    debug_names.retain(|inst| {
+        (inst.class.opcode != Op::Name || name_ids.insert(inst.operands[0].unwrap_id_ref()))
+            && (inst.class.opcode != Op::MemberName
+                || member_name_ids.insert((
+                    inst.operands[0].unwrap_id_ref(),
+                    inst.operands[1].unwrap_literal_bit32(),
+                )))
+    });
+}
+
 pub fn remove_duplicate_types(module: &mut Module) {
     // Keep in mind, this algorithm requires forward type references to not exist - i.e. it's a valid spir-v module.
 
@@ -259,17 +277,9 @@ pub fn remove_duplicate_types(module: &mut Module) {
     module
         .annotations
         .retain(|inst| anno_set.insert(inst.assemble()));
-    // Same thing with OpName
-    let mut name_ids = FxHashSet::default();
-    let mut member_name_ids = FxHashSet::default();
-    module.debug_names.retain(|inst| {
-        (inst.class.opcode != Op::Name || name_ids.insert(inst.operands[0].unwrap_id_ref()))
-            && (inst.class.opcode != Op::MemberName
-                || member_name_ids.insert((
-                    inst.operands[0].unwrap_id_ref(),
-                    inst.operands[1].unwrap_literal_bit32(),
-                )))
-    });
+
+    // Same thing with debug names section
+    remove_duplicate_debug_names(&mut module.debug_names);
 }
 
 pub fn remove_duplicate_debuginfo(module: &mut Module) {
@@ -541,3 +551,87 @@ pub fn remove_duplicate_debuginfo(module: &mut Module) {
         }
     }
 }
+
+pub fn remove_duplicate_builtin_variables(module: &mut Module) {
+    // Find the variables decorated as builtins, and any duplicates of each builtin.
+
+    // A map from deleted duplicate variable ID to the new, de-duplicated ID.
+    let duplicate_vars: FxHashMap<Word, Word> = {
+        // A map from builtin to a de-duplicated variable decorated as that builtin.
+        let mut builtin_to_var_id = FxHashMap::<BuiltIn, Word>::default();
+
+        let mut duplicate_vars = FxHashMap::<Word, Word>::default();
+
+        for inst in module.annotations.iter() {
+            if inst.class.opcode == Op::Decorate
+                && let [
+                    Operand::IdRef(var_id),
+                    Operand::Decoration(Decoration::BuiltIn),
+                    Operand::BuiltIn(builtin),
+                ] = inst.operands[..]
+            {
+                match builtin_to_var_id.entry(builtin) {
+                    // first variable we've seen for this builtin,
+                    // record it in the builtins map.
+                    hash_map::Entry::Vacant(vacant) => {
+                        vacant.insert(var_id);
+                    }
+
+                    // this builtin already has a variable,
+                    // record it in the duplicates map.
+                    hash_map::Entry::Occupied(occupied) => {
+                        duplicate_vars.insert(var_id, *occupied.get());
+                    }
+                };
+            }
+        }
+        // Rebind as immutable in fn scope.
+        duplicate_vars
+    };
+
+    // Rewrite entry points, removing duplicate variables.
+    for entry in &mut module.entry_points {
+        if entry.class.opcode != Op::EntryPoint {
+            continue;
+        }
+
+        entry
+            .operands
+            .retain(|op| !matches!(op, Operand::IdRef(id) if duplicate_vars.contains_key(id)));
+    }
+
+    // Remove duplicate debug names after merging variables.
+    remove_duplicate_debug_names(&mut module.debug_names);
+
+    // Rewrite annotations for duplicates to point at de-duplicated variables;
+    // this will merge the annotations sets but produce duplicate annotations
+    // (which we will remove next).
+    for inst in &mut module.annotations {
+        rewrite_inst_with_rules(inst, &duplicate_vars);
+    }
+
+    // Merge the annotations for duplicate vars.
+    {
+        let mut annotations_set = FxHashSet::default();
+        module
+            .annotations
+            // Note: insert returns true when an annotation is inserted for the first time.
+            .retain(|inst| annotations_set.insert(inst.assemble()));
+    }
+
+    // Remove the duplicate variable definitions.
+    module.types_global_values.retain(|inst| {
+        !matches!(inst.result_id,
+                                   Some(id) if duplicate_vars.contains_key(&id))
+    });
+
+    // Rewrite function blocks to use de-duplicated variables.
+    for inst in &mut module
+        .functions
+        .iter_mut()
+        .flat_map(|f| &mut f.blocks)
+        .flat_map(|b| &mut b.instructions)
+    {
+        rewrite_inst_with_rules(inst, &duplicate_vars);
+    }
+}

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -290,6 +290,7 @@ pub fn link(
         duplicates::remove_duplicate_capabilities(&mut output);
         duplicates::remove_duplicate_ext_inst_imports(&mut output);
         duplicates::remove_duplicate_types(&mut output);
+        duplicates::remove_duplicate_builtin_variables(&mut output);
         // jb-todo: strip identical OpDecoration / OpDecorationGroups
     }
 

tests/compiletests/ui/entry/builtin_duplicates.stderr

@@ -1,73 +1,63 @@
 OpCapability Shader
 OpMemoryModel Logical Simple
-OpEntryPoint GLCompute %1 "entry_1" %2 %3
-OpEntryPoint GLCompute %4 "entry_2" %2 %3
+OpEntryPoint GLCompute %1 "entry_1" %2
+OpEntryPoint GLCompute %3 "entry_2" %2
 OpExecutionMode %1 LocalSize 1 1 1
-OpExecutionMode %4 LocalSize 1 1 1
-OpName %6 "builtin_duplicates::entry_1"
-OpName %7 "builtin_duplicates::sub_1"
-OpName %8 "builtin_duplicates::sub_2"
-OpName %9 "builtin_duplicates::entry_2"
+OpExecutionMode %3 LocalSize 1 1 1
+OpName %5 "builtin_duplicates::entry_1"
+OpName %6 "builtin_duplicates::sub_1"
+OpName %7 "builtin_duplicates::sub_2"
+OpName %8 "builtin_duplicates::entry_2"
 OpDecorate %2 BuiltIn LocalInvocationIndex
-OpDecorate %3 BuiltIn LocalInvocationIndex
-%10 = OpTypeInt 32 0
-%11 = OpTypePointer Input %10
-%12 = OpTypeVoid
-%13 = OpTypeFunction %12
-%2 = OpVariable  %11  Input
-%14 = OpTypeFunction %12 %10
-%3 = OpVariable  %11  Input
-%15 = OpTypeFunction %10
-%1 = OpFunction  %12  None %13
-%16 = OpLabel
-%17 = OpLoad  %10  %2
-%18 = OpFunctionCall  %12  %6 %17
+%9 = OpTypeInt 32 0
+%10 = OpTypePointer Input %9
+%11 = OpTypeVoid
+%12 = OpTypeFunction %11
+%2 = OpVariable  %10  Input
+%13 = OpTypeFunction %11 %9
+%14 = OpTypeFunction %9
+%1 = OpFunction  %11  None %12
+%15 = OpLabel
+%16 = OpLoad  %9  %2
+%17 = OpFunctionCall  %11  %5 %16
 OpNoLine
 OpReturn
 OpFunctionEnd
-%6 = OpFunction  %12  None %14
-%19 = OpFunctionParameter  %10
-%20 = OpLabel
-%21 = OpLoad  %10  %3
-%22 = OpLoad  %10  %3
-%23 = OpFunctionCall  %10  %7
-%24 = OpFunctionCall  %10  %8
+%5 = OpFunction  %11  None %13
+%18 = OpFunctionParameter  %9
+%19 = OpLabel
+%20 = OpLoad  %9  %2
+%21 = OpLoad  %9  %2
+%22 = OpFunctionCall  %9  %6
+%23 = OpFunctionCall  %9  %7
 OpNoLine
 OpReturn
 OpFunctionEnd
-%7 = OpFunction  %10  DontInline %15
-%25 = OpLabel
-%26 = OpLoad  %10  %3
+%6 = OpFunction  %9  DontInline %14
+%24 = OpLabel
+%25 = OpLoad  %9  %2
 OpNoLine
-OpReturnValue %26
+OpReturnValue %25
 OpFunctionEnd
-%8 = OpFunction  %10  DontInline %15
-%27 = OpLabel
-%28 = OpLoad  %10  %3
+%7 = OpFunction  %9  DontInline %14
+%26 = OpLabel
+%27 = OpLoad  %9  %2
 OpNoLine
-OpReturnValue %28
+OpReturnValue %27
 OpFunctionEnd
-%4 = OpFunction  %12  None %13
-%29 = OpLabel
-%30 = OpLoad  %10  %2
-%31 = OpFunctionCall  %12  %9 %30
+%3 = OpFunction  %11  None %12
+%28 = OpLabel
+%29 = OpLoad  %9  %2
+%30 = OpFunctionCall  %11  %8 %29
 OpNoLine
 OpReturn
 OpFunctionEnd
-%9 = OpFunction  %12  None %14
-%32 = OpFunctionParameter  %10
-%33 = OpLabel
-%34 = OpLoad  %10  %3
-%35 = OpFunctionCall  %10  %7
-%36 = OpFunctionCall  %10  %8
+%8 = OpFunction  %11  None %13
+%31 = OpFunctionParameter  %9
+%32 = OpLabel
+%33 = OpLoad  %9  %2
+%34 = OpFunctionCall  %9  %6
+%35 = OpFunctionCall  %9  %7
 OpNoLine
 OpReturn
 OpFunctionEnd
-error: error:0:0 - [VUID-StandaloneSpirv-OpEntryPoint-09658] OpEntryPoint contains duplicate input variables with LocalInvocationIndex builtin
-         %gl_LocalInvocationIndex = OpVariable %_ptr_Input_uint Input
-   |
-   = note: spirv-val failed
-   = note: module `$TEST_BUILD_DIR/entry/builtin_duplicates.vulkan1.2`
-
-error: aborting due to 1 previous error
-