Revert "WIP: mem2reg speedup"

This reverts commit efbf694.

Author: eddyb

crates/rustc_codegen_spirv/src/linker/dce.rs

@@ -7,10 +7,9 @@
 //! *references* a rooted thing is also rooted, not the other way around - but that's the basic
 //! concept.
 
-use rspirv::dr::{Block, Function, Instruction, Module, Operand};
+use rspirv::dr::{Function, Instruction, Module, Operand};
 use rspirv::spirv::{Decoration, LinkageType, Op, StorageClass, Word};
-use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
-use std::hash::Hash;
+use rustc_data_structures::fx::FxIndexSet;
 
 pub fn dce(module: &mut Module) {
     let mut rooted = collect_roots(module);
@@ -138,11 +137,11 @@ fn kill_unrooted(module: &mut Module, rooted: &FxIndexSet<Word>) {
     }
 }
 
-pub fn dce_phi(blocks: &mut FxIndexMap<impl Eq + Hash, &mut Block>) {
+pub fn dce_phi(func: &mut Function) {
     let mut used = FxIndexSet::default();
     loop {
         let mut changed = false;
-        for inst in blocks.values().flat_map(|block| &block.instructions) {
+        for inst in func.all_inst_iter() {
             if inst.class.opcode != Op::Phi || used.contains(&inst.result_id.unwrap()) {
                 for op in &inst.operands {
                     if let Some(id) = op.id_ref_any() {
@@ -155,7 +154,7 @@ pub fn dce_phi(blocks: &mut FxIndexMap<impl Eq + Hash, &mut Block>) {
             break;
         }
     }
-    for block in blocks.values_mut() {
+    for block in &mut func.blocks {
         block
             .instructions
             .retain(|inst| inst.class.opcode != Op::Phi || used.contains(&inst.result_id.unwrap()));

crates/rustc_codegen_spirv/src/linker/mem2reg.rs

@@ -13,31 +13,19 @@ use super::simple_passes::outgoing_edges;
 use super::{apply_rewrite_rules, id};
 use rspirv::dr::{Block, Function, Instruction, ModuleHeader, Operand};
 use rspirv::spirv::{Op, Word};
-use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexMap};
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_middle::bug;
 use std::collections::hash_map;
 
-// HACK(eddyb) newtype instead of type alias to avoid mistakes.
-#[derive(Copy, Clone, PartialEq, Eq, Hash)]
-struct LabelId(Word);
-
 pub fn mem2reg(
     header: &mut ModuleHeader,
     types_global_values: &mut Vec<Instruction>,
     pointer_to_pointee: &FxHashMap<Word, Word>,
     constants: &FxHashMap<Word, u32>,
     func: &mut Function,
 ) {
-    // HACK(eddyb) this ad-hoc indexing might be useful elsewhere as well, but
-    // it's made completely irrelevant by SPIR-T so only applies to legacy code.
-    let mut blocks: FxIndexMap<_, _> = func
-        .blocks
-        .iter_mut()
-        .map(|block| (LabelId(block.label_id().unwrap()), block))
-        .collect();
-
-    let reachable = compute_reachable(&blocks);
-    let preds = compute_preds(&blocks, &reachable);
+    let reachable = compute_reachable(&func.blocks);
+    let preds = compute_preds(&func.blocks, &reachable);
     let idom = compute_idom(&preds, &reachable);
     let dominance_frontier = compute_dominance_frontier(&preds, &idom);
     loop {
@@ -46,27 +34,31 @@ pub fn mem2reg(
             types_global_values,
             pointer_to_pointee,
             constants,
-            &mut blocks,
+            &mut func.blocks,
             &dominance_frontier,
         );
         if !changed {
             break;
         }
         // mem2reg produces minimal SSA form, not pruned, so DCE the dead ones
-        super::dce::dce_phi(&mut blocks);
+        super::dce::dce_phi(func);
     }
 }
 
-fn compute_reachable(blocks: &FxIndexMap<LabelId, &mut Block>) -> Vec<bool> {
-    fn recurse(blocks: &FxIndexMap<LabelId, &mut Block>, reachable: &mut [bool], block: usize) {
+fn label_to_index(blocks: &[Block], id: Word) -> usize {
+    blocks
+        .iter()
+        .position(|b| b.label_id().unwrap() == id)
+        .unwrap()
+}
+
+fn compute_reachable(blocks: &[Block]) -> Vec<bool> {
+    fn recurse(blocks: &[Block], reachable: &mut [bool], block: usize) {
         if !reachable[block] {
             reachable[block] = true;
-            for dest_id in outgoing_edges(blocks[block]) {
-                recurse(
-                    blocks,
-                    reachable,
-                    blocks.get_index_of(&LabelId(dest_id)).unwrap(),
-                );
+            for dest_id in outgoing_edges(&blocks[block]) {
+                let dest_idx = label_to_index(blocks, dest_id);
+                recurse(blocks, reachable, dest_idx);
             }
         }
     }
@@ -75,19 +67,17 @@ fn compute_reachable(blocks: &FxIndexMap<LabelId, &mut Block>) -> Vec<bool> {
     reachable
 }
 
-fn compute_preds(
-    blocks: &FxIndexMap<LabelId, &mut Block>,
-    reachable_blocks: &[bool],
-) -> Vec<Vec<usize>> {
+fn compute_preds(blocks: &[Block], reachable_blocks: &[bool]) -> Vec<Vec<usize>> {
     let mut result = vec![vec![]; blocks.len()];
     // Do not count unreachable blocks as valid preds of blocks
     for (source_idx, source) in blocks
-        .values()
+        .iter()
         .enumerate()
         .filter(|&(b, _)| reachable_blocks[b])
     {
         for dest_id in outgoing_edges(source) {
-            result[blocks.get_index_of(&LabelId(dest_id)).unwrap()].push(source_idx);
+            let dest_idx = label_to_index(blocks, dest_id);
+            result[dest_idx].push(source_idx);
         }
     }
     result
@@ -171,7 +161,7 @@ fn insert_phis_all(
     types_global_values: &mut Vec<Instruction>,
     pointer_to_pointee: &FxHashMap<Word, Word>,
     constants: &FxHashMap<Word, u32>,
-    blocks: &mut FxIndexMap<LabelId, &mut Block>,
+    blocks: &mut [Block],
     dominance_frontier: &[FxHashSet<usize>],
 ) -> bool {
     let var_maps_and_types = blocks[0]
@@ -208,11 +198,7 @@ fn insert_phis_all(
             rewrite_rules: FxHashMap::default(),
         };
         renamer.rename(0, None);
-        // FIXME(eddyb) shouldn't this full rescan of the function be done once?
-        apply_rewrite_rules(
-            &renamer.rewrite_rules,
-            blocks.values_mut().map(|block| &mut **block),
-        );
+        apply_rewrite_rules(&renamer.rewrite_rules, blocks);
         remove_nops(blocks);
     }
     remove_old_variables(blocks, &var_maps_and_types);
@@ -230,7 +216,7 @@ struct VarInfo {
 fn collect_access_chains(
     pointer_to_pointee: &FxHashMap<Word, Word>,
     constants: &FxHashMap<Word, u32>,
-    blocks: &FxIndexMap<LabelId, &mut Block>,
+    blocks: &[Block],
     base_var: Word,
     base_var_ty: Word,
 ) -> Option<FxHashMap<Word, VarInfo>> {
@@ -263,7 +249,7 @@ fn collect_access_chains(
     // Loop in case a previous block references a later AccessChain
     loop {
         let mut changed = false;
-        for inst in blocks.values().flat_map(|b| &b.instructions) {
+        for inst in blocks.iter().flat_map(|b| &b.instructions) {
             for (index, op) in inst.operands.iter().enumerate() {
                 if let Operand::IdRef(id) = op
                     && variables.contains_key(id)
@@ -317,10 +303,10 @@ fn collect_access_chains(
 // same var map (e.g. `s.x = s.y;`).
 fn split_copy_memory(
     header: &mut ModuleHeader,
-    blocks: &mut FxIndexMap<LabelId, &mut Block>,
+    blocks: &mut [Block],
     var_map: &FxHashMap<Word, VarInfo>,
 ) {
-    for block in blocks.values_mut() {
+    for block in blocks {
         let mut inst_index = 0;
         while inst_index < block.instructions.len() {
             let inst = &block.instructions[inst_index];
@@ -379,7 +365,7 @@ fn has_store(block: &Block, var_map: &FxHashMap<Word, VarInfo>) -> bool {
 }
 
 fn insert_phis(
-    blocks: &FxIndexMap<LabelId, &mut Block>,
+    blocks: &[Block],
     dominance_frontier: &[FxHashSet<usize>],
     var_map: &FxHashMap<Word, VarInfo>,
 ) -> FxHashSet<usize> {
@@ -388,7 +374,7 @@ fn insert_phis(
     let mut ever_on_work_list = FxHashSet::default();
     let mut work_list = Vec::new();
     let mut blocks_with_phi = FxHashSet::default();
-    for (block_idx, block) in blocks.values().enumerate() {
+    for (block_idx, block) in blocks.iter().enumerate() {
         if has_store(block, var_map) {
             ever_on_work_list.insert(block_idx);
             work_list.push(block_idx);
@@ -433,10 +419,10 @@ fn top_stack_or_undef(
     }
 }
 
-struct Renamer<'a, 'b> {
+struct Renamer<'a> {
     header: &'a mut ModuleHeader,
     types_global_values: &'a mut Vec<Instruction>,
-    blocks: &'a mut FxIndexMap<LabelId, &'b mut Block>,
+    blocks: &'a mut [Block],
     blocks_with_phi: FxHashSet<usize>,
     base_var_type: Word,
     var_map: &'a FxHashMap<Word, VarInfo>,
@@ -446,7 +432,7 @@ struct Renamer<'a, 'b> {
     rewrite_rules: FxHashMap<Word, Word>,
 }
 
-impl Renamer<'_, '_> {
+impl Renamer<'_> {
     // Returns the phi definition.
     fn insert_phi_value(&mut self, block: usize, from_block: usize) -> Word {
         let from_block_label = self.blocks[from_block].label_id().unwrap();
@@ -568,8 +554,9 @@ impl Renamer<'_, '_> {
             }
         }
 
-        for dest_id in outgoing_edges(self.blocks[block]).collect::<Vec<_>>() {
-            let dest_idx = self.blocks.get_index_of(&LabelId(dest_id)).unwrap();
+        for dest_id in outgoing_edges(&self.blocks[block]).collect::<Vec<_>>() {
+            // TODO: Don't do this find
+            let dest_idx = label_to_index(self.blocks, dest_id);
             self.rename(dest_idx, Some(block));
         }
 
@@ -579,16 +566,16 @@ impl Renamer<'_, '_> {
     }
 }
 
-fn remove_nops(blocks: &mut FxIndexMap<LabelId, &mut Block>) {
-    for block in blocks.values_mut() {
+fn remove_nops(blocks: &mut [Block]) {
+    for block in blocks {
         block
             .instructions
             .retain(|inst| inst.class.opcode != Op::Nop);
     }
 }
 
 fn remove_old_variables(
-    blocks: &mut FxIndexMap<LabelId, &mut Block>,
+    blocks: &mut [Block],
     var_maps_and_types: &[(FxHashMap<u32, VarInfo>, u32)],
 ) {
     blocks[0].instructions.retain(|inst| {
@@ -599,7 +586,7 @@ fn remove_old_variables(
                 .all(|(var_map, _)| !var_map.contains_key(&result_id))
         }
     });
-    for block in blocks.values_mut() {
+    for block in blocks {
         block.instructions.retain(|inst| {
             !matches!(inst.class.opcode, Op::AccessChain | Op::InBoundsAccessChain)
                 || inst.operands.iter().all(|op| {

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -85,12 +85,9 @@ fn id(header: &mut ModuleHeader) -> Word {
     result
 }
 
-fn apply_rewrite_rules<'a>(
-    rewrite_rules: &FxHashMap<Word, Word>,
-    blocks: impl IntoIterator<Item = &'a mut Block>,
-) {
+fn apply_rewrite_rules(rewrite_rules: &FxHashMap<Word, Word>, blocks: &mut [Block]) {
     let all_ids_mut = blocks
-        .into_iter()
+        .iter_mut()
         .flat_map(|b| b.label.iter_mut().chain(b.instructions.iter_mut()))
         .flat_map(|inst| {
             inst.result_id