Added a transformation that gets rid of temporary composites. (#690)

* Added an optimization that gets rid of temporary composites.

Those temporary composites result from inlining of multi-argument
closures. Not only are they rather useless, they're also sometimes
invalid, when an argument to said closure is e.g. a pointer.

* Correctness fixes to transitive unused removal:

- delay only if the instruction is in reference set
- properly mark composites being inserted into composites as used

* cargo fmt

* clippy

* Make transformation per-function & rely on DCE for eliminating dead constructs.

* Forgot to mark CompositeInsert as pure & additional line cleaning

* Rustfmt

* Remove duplicate lines only once

crates/rustc_codegen_spirv/src/link.rs

@@ -541,6 +541,7 @@ fn do_link(
         dce: env::var("NO_DCE").is_err(),
         compact_ids: env::var("NO_COMPACT_IDS").is_err(),
         inline: legalize,
+        destructure: legalize,
         mem2reg: legalize,
         structurize: env::var("NO_STRUCTURIZE").is_err(),
         emit_multiple_modules: cg_args.module_output_type == ModuleOutputType::Multiple,

crates/rustc_codegen_spirv/src/linker/dce.rs

@@ -162,6 +162,7 @@ fn instruction_is_pure(inst: &Instruction) -> bool {
         | InBoundsPtrAccessChain
         | CompositeConstruct
         | CompositeExtract
+        | CompositeInsert
         | CopyObject
         | Transpose
         | ConvertFToU

crates/rustc_codegen_spirv/src/linker/destructure_composites.rs

@@ -0,0 +1,74 @@
+//! Simplify `OpCompositeExtract` pointing to `OpCompositeConstruct`s / `OpCompositeInsert`s.
+//! Such constructions arise after inlining, when using multi-argument closures
+//! (and other `Fn*` trait implementations). These composites can frequently be invalid,
+//! containing pointers, `OpFunctionArgument`s, etc. After simplification, components
+//! will become valid targets for `OpLoad`/`OpStore`.
+use super::apply_rewrite_rules;
+use rspirv::dr::{Function, Instruction};
+use rspirv::spirv::Op;
+use rustc_data_structures::fx::FxHashMap;
+
+pub fn destructure_composites(function: &mut Function) {
+    let mut rewrite_rules = FxHashMap::default();
+    let reference: FxHashMap<_, _> = function
+        .all_inst_iter()
+        .filter_map(|inst| match inst.class.opcode {
+            Op::CompositeConstruct => Some((inst.result_id.unwrap(), inst.clone())),
+            Op::CompositeInsert if inst.operands.len() == 3 => {
+                Some((inst.result_id.unwrap(), inst.clone()))
+            }
+            _ => None,
+        })
+        .collect();
+    for inst in function.all_inst_iter_mut() {
+        if inst.class.opcode == Op::CompositeExtract && inst.operands.len() == 2 {
+            let mut composite = inst.operands[0].unwrap_id_ref();
+            let index = inst.operands[1].unwrap_literal_int32();
+
+            let origin = loop {
+                if let Some(inst) = reference.get(&composite) {
+                    match inst.class.opcode {
+                        Op::CompositeInsert => {
+                            let insert_index = inst.operands[2].unwrap_literal_int32();
+                            if insert_index == index {
+                                break Some(inst.operands[0].unwrap_id_ref());
+                            }
+                            composite = inst.operands[1].unwrap_id_ref();
+                        }
+                        Op::CompositeConstruct => {
+                            break inst.operands.get(index as usize).map(|o| o.unwrap_id_ref());
+                        }
+                        _ => unreachable!(),
+                    }
+                } else {
+                    break None;
+                }
+            };
+
+            if let Some(origin_id) = origin {
+                rewrite_rules.insert(
+                    inst.result_id.unwrap(),
+                    rewrite_rules.get(&origin_id).map_or(origin_id, |id| *id),
+                );
+                *inst = Instruction::new(Op::Nop, None, None, vec![]);
+                continue;
+            }
+        }
+    }
+
+    // Transitive closure computation
+    let mut closed_rewrite_rules = rewrite_rules.clone();
+    for (_, value) in closed_rewrite_rules.iter_mut() {
+        while let Some(next) = rewrite_rules.get(value) {
+            *value = *next;
+        }
+    }
+
+    // Remove instructions replaced by NOPs, as well as unused composite values.
+    for block in function.blocks.iter_mut() {
+        block
+            .instructions
+            .retain(|inst| inst.class.opcode != Op::Nop);
+    }
+    apply_rewrite_rules(&closed_rewrite_rules, &mut function.blocks);
+}

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -2,6 +2,7 @@
 mod test;
 
 mod dce;
+mod destructure_composites;
 mod duplicates;
 mod import_export_link;
 mod inline;
@@ -27,6 +28,7 @@ pub struct Options {
     pub dce: bool,
     pub inline: bool,
     pub mem2reg: bool,
+    pub destructure: bool,
     pub structurize: bool,
     pub emit_multiple_modules: bool,
     pub name_variables: bool,
@@ -228,6 +230,10 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
                 // mem2reg produces minimal SSA form, not pruned, so DCE the dead ones
                 dce::dce_phi(func);
             }
+            if opts.destructure {
+                let _timer = sess.timer("link_destructure");
+                destructure_composites::destructure_composites(func);
+            }
         }
     }
 
@@ -240,11 +246,6 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
         }
     }
 
-    {
-        let _timer = sess.timer("link_remove_duplicate_lines");
-        duplicates::remove_duplicate_lines(&mut output);
-    }
-
     if opts.name_variables {
         let _timer = sess.timer("link_name_variables");
         simple_passes::name_variables_pass(&mut output);
@@ -289,6 +290,11 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
             dce::dce(output);
         }
 
+        {
+            let _timer = sess.timer("link_remove_duplicate_lines");
+            duplicates::remove_duplicate_lines(output);
+        }
+
         if opts.compact_ids {
             let _timer = sess.timer("link_compact_ids");
             // compact the ids https://github.com/KhronosGroup/SPIRV-Tools/blob/e02f178a716b0c3c803ce31b9df4088596537872/source/opt/compact_ids_pass.cpp#L43

crates/rustc_codegen_spirv/src/linker/test.rs

@@ -92,6 +92,7 @@ fn assemble_and_link(binaries: &[&[u8]]) -> Result<Module, String> {
                 compact_ids: true,
                 dce: false,
                 inline: false,
+                destructure: false,
                 mem2reg: false,
                 structurize: false,
                 emit_multiple_modules: false,

tests/ui/dis/index_user_dst.stderr

@@ -3,34 +3,32 @@
 OpLine %5 7 12
 %6 = OpAccessChain  %7  %8 %9
 %10 = OpArrayLength  %11  %8 0
-OpLine %5 7 0
-%12 = OpCompositeInsert  %13  %6 %14 0
 OpLine %5 8 21
-%15 = OpULessThan  %16  %9 %10
+%12 = OpULessThan  %13  %9 %10
 OpLine %5 8 21
-OpSelectionMerge %17 None
+OpSelectionMerge %14 None
-OpBranchConditional %15 %18 %19
+OpBranchConditional %12 %15 %16
-%18 = OpLabel
+%15 = OpLabel
 OpLine %5 8 21
-%20 = OpInBoundsAccessChain  %21  %6 %9
+%17 = OpInBoundsAccessChain  %18  %6 %9
-%22 = OpLoad  %23  %20
+%19 = OpLoad  %20  %17
 OpLine %5 10 1
 OpReturn
-%19 = OpLabel
+%16 = OpLabel
 OpLine %5 8 21
-OpBranch %24
+OpBranch %21
-%24 = OpLabel
+%21 = OpLabel
+OpBranch %22
+%22 = OpLabel
+%23 = OpPhi  %13  %24 %21 %24 %25
+OpLoopMerge %26 %25 None
+OpBranchConditional %23 %27 %26
+%27 = OpLabel
 OpBranch %25
 %25 = OpLabel
-%26 = OpPhi  %16  %27 %24 %27 %28
+OpBranch %22
-OpLoopMerge %29 %28 None
+%26 = OpLabel
-OpBranchConditional %26 %30 %29
-%30 = OpLabel
-OpBranch %28
-%28 = OpLabel
-OpBranch %25
-%29 = OpLabel
 OpUnreachable
-%17 = OpLabel
+%14 = OpLabel
 OpUnreachable
 OpFunctionEnd

tests/ui/lang/control_flow/closure_multi.rs

@@ -0,0 +1,16 @@
+// build-pass
+
+use spirv_std;
+
+fn closure_user<F: FnMut(&u32, u32)>(ptr: &u32, xmax: u32, mut callback: F) {
+    for i in 0..xmax {
+        callback(ptr, i);
+    }
+}
+
+#[spirv(fragment)]
+pub fn main(ptr: &mut u32) {
+    closure_user(ptr, 10, |ptr, i| {
+        if *ptr == i { spirv_std::arch::kill(); }
+    });
+}