qptr/lower: more aggressively strip Offset(0) and remove unused instructions.

Author: eddyb

src/qptr/lower.rs

@@ -8,11 +8,14 @@ use crate::qptr::{shapes, QPtrAttr, QPtrOp};
 use crate::transform::{InnerInPlaceTransform, Transformed, Transformer};
 use crate::{
     spv, AddrSpace, AttrSet, AttrSetDef, Const, ConstCtor, ConstDef, Context, ControlNode,
-    ControlNodeKind, DataInst, DataInstDef, DataInstForm, DataInstFormDef, DataInstKind, Diag,
-    FuncDecl, GlobalVarDecl, OrdAssertEq, Type, TypeCtor, TypeCtorArg, TypeDef, Value,
+    ControlNodeKind, DataInst, DataInstDef, DataInstForm, DataInstFormDef, DataInstKind, DeclDef,
+    Diag, EntityOrientedDenseMap, FuncDecl, GlobalVarDecl, OrdAssertEq, Type, TypeCtor,
+    TypeCtorArg, TypeDef, Value,
 };
+use rustc_hash::FxHashMap;
 use smallvec::SmallVec;
 use std::cell::Cell;
+use std::mem;
 use std::num::NonZeroU32;
 use std::rc::Rc;
 
@@ -143,7 +146,13 @@ impl<'a> LowerFromSpvPtrs<'a> {
         // separately - so `LowerFromSpvPtrInstsInFunc` will leave all value defs
         // (including replaced instructions!) with unchanged `OpTypePointer`
         // types, that only `EraseSpvPtrs`, later, replaces with `QPtr`.
-        LowerFromSpvPtrInstsInFunc { lowerer: self }.in_place_transform_func_decl(func_decl);
+        LowerFromSpvPtrInstsInFunc {
+            lowerer: self,
+            data_inst_use_counts: Default::default(),
+            remove_if_dead_inst_and_parent_block: Default::default(),
+            noop_offsets_to_base_ptr: Default::default(),
+        }
+        .in_place_transform_func_decl(func_decl);
         EraseSpvPtrs { lowerer: self }.in_place_transform_func_decl(func_decl);
     }
 
@@ -258,6 +267,19 @@ impl Transformer for EraseSpvPtrs<'_> {
 
 struct LowerFromSpvPtrInstsInFunc<'a> {
     lowerer: &'a LowerFromSpvPtrs<'a>,
+
+    // FIXME(eddyb) consider removing this and just do a full second traversal.
+    data_inst_use_counts: EntityOrientedDenseMap<DataInst, NonZeroU32>,
+
+    // HACK(eddyb) this acts as a "queue" for `qptr`-producing instructions,
+    // which may end up dead because they're unused (either unused originally,
+    // in SPIR-V, or because of offset folding).
+    remove_if_dead_inst_and_parent_block: Vec<(DataInst, ControlNode)>,
+
+    // FIXME(eddyb) this is redundant with a few other things and only here
+    // because it needs to be available from `transform_value`, which doesn't
+    // have access to a `FuncAt` to look up anything.
+    noop_offsets_to_base_ptr: FxHashMap<DataInst, Value>,
 }
 
 /// One `QPtr`->`QPtr` step used in the lowering of `Op*AccessChain`.
@@ -413,7 +435,7 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
     }
 
     fn try_lower_data_inst_def(
-        &self,
+        &mut self,
         mut func_at_data_inst: FuncAtMut<'_, DataInst>,
         parent_block: ControlNode,
     ) -> Result<Transformed<DataInstDef>, LowerError> {
@@ -438,18 +460,25 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
             _ => return Ok(Transformed::Unchanged),
         };
 
-        // Map `ptr` to its base & offset, if it points to a `QPtrOp::Offset`.
-        let ptr_to_base_ptr_and_offset = |ptr| match ptr {
-            Value::DataInstOutput(ptr_inst) => {
+        // Flatten `QPtrOp::Offset`s behind `ptr` into a base pointer and offset.
+        let flatten_offsets = |mut ptr| {
+            let mut offset = 0;
+            while let Value::DataInstOutput(ptr_inst) = ptr {
                 let ptr_inst_def = func.at(ptr_inst).def();
                 match cx[ptr_inst_def.form].kind {
                     DataInstKind::QPtr(QPtrOp::Offset(ptr_offset)) => {
-                        Some((ptr_inst_def.inputs[0], ptr_offset))
+                        match ptr_offset.checked_add(offset) {
+                            Some(combined_offset) => {
+                                ptr = ptr_inst_def.inputs[0];
+                                offset = combined_offset;
+                            }
+                            None => break,
+                        }
                     }
-                    _ => None,
+                    _ => break,
                 }
             }
-            _ => None,
+            (ptr, offset)
         };
 
         let replacement_kind_and_inputs = if spv_inst.opcode == wk.OpVariable {
@@ -478,7 +507,7 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
 
             let ptr = data_inst_def.inputs[0];
 
-            let (ptr, offset) = ptr_to_base_ptr_and_offset(ptr).unwrap_or((ptr, 0));
+            let (ptr, offset) = flatten_offsets(ptr);
 
             (QPtrOp::Load { offset }.into(), [ptr].into_iter().collect())
         } else if spv_inst.opcode == wk.OpStore {
@@ -491,7 +520,7 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
             let ptr = data_inst_def.inputs[0];
             let value = data_inst_def.inputs[1];
 
-            let (ptr, offset) = ptr_to_base_ptr_and_offset(ptr).unwrap_or((ptr, 0));
+            let (ptr, offset) = flatten_offsets(ptr);
 
             (
                 QPtrOp::Store { offset }.into(),
@@ -608,11 +637,10 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
                 dyn_idx: None,
             }) = steps.first_mut()
             {
-                if let Some((ptr_base_ptr, ptr_offset)) = ptr_to_base_ptr_and_offset(ptr) {
-                    if let Some(new_first_offset) = first_offset.checked_add(ptr_offset) {
-                        ptr = ptr_base_ptr;
-                        *first_offset = new_first_offset;
-                    }
+                let (ptr_base_ptr, ptr_offset) = flatten_offsets(ptr);
+                if let Some(new_first_offset) = first_offset.checked_add(ptr_offset) {
+                    ptr = ptr_base_ptr;
+                    *first_offset = new_first_offset;
                 }
             }
 
@@ -648,6 +676,11 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
                 match &mut func.control_nodes[parent_block].kind {
                     ControlNodeKind::Block { insts } => {
                         insts.insert_before(step_data_inst, data_inst, func.data_insts);
+
+                        // HACK(eddyb) this tracking is kind of ad-hoc but should
+                        // easily cover everything we care about for now.
+                        self.remove_if_dead_inst_and_parent_block
+                            .push((step_data_inst, parent_block));
                     }
                     _ => unreachable!(),
                 }
@@ -762,9 +795,51 @@ impl LowerFromSpvPtrInstsInFunc<'_> {
             func_at_data_inst.def().attrs = cx.intern(attrs);
         }
     }
+
+    // FIXME(eddyb) these are only this whacky because an `u32` is being
+    // encoded as `Option<NonZeroU32>` for (dense) map entry reasons.
+    fn add_value_uses(&mut self, values: &[Value]) {
+        for &v in values {
+            if let Value::DataInstOutput(data_inst) = v {
+                let count = self.data_inst_use_counts.entry(data_inst);
+                *count = Some(
+                    NonZeroU32::new(count.map_or(0, |c| c.get()).checked_add(1).unwrap()).unwrap(),
+                );
+            }
+        }
+    }
+    fn remove_value_uses(&mut self, values: &[Value]) {
+        for &v in values {
+            if let Value::DataInstOutput(data_inst) = v {
+                let count = self.data_inst_use_counts.entry(data_inst);
+                *count = NonZeroU32::new(count.unwrap().get() - 1);
+            }
+        }
+    }
 }
 
 impl Transformer for LowerFromSpvPtrInstsInFunc<'_> {
+    // NOTE(eddyb) it's important that this only gets invoked on already lowered
+    // `Value`s, so we can rely on e.g. `noop_offsets_to_base_ptr` being filled.
+    fn transform_value_use(&mut self, v: &Value) -> Transformed<Value> {
+        let mut v = *v;
+
+        let transformed = match v {
+            Value::DataInstOutput(inst) => self
+                .noop_offsets_to_base_ptr
+                .get(&inst)
+                .copied()
+                .map_or(Transformed::Unchanged, Transformed::Changed),
+
+            _ => Transformed::Unchanged,
+        };
+
+        transformed.apply_to(&mut v);
+        self.add_value_uses(&[v]);
+
+        transformed
+    }
+
     // HACK(eddyb) while we want to transform `DataInstDef`s, we can't inject
     // adjacent instructions without access to the parent `ControlNodeKind::Block`,
     // and to fix this would likely require list nodes to carry some handle to
@@ -777,16 +852,44 @@ impl Transformer for LowerFromSpvPtrInstsInFunc<'_> {
         &mut self,
         mut func_at_control_node: FuncAtMut<'_, ControlNode>,
     ) {
-        func_at_control_node
-            .reborrow()
-            .inner_in_place_transform_with(self);
-
         let control_node = func_at_control_node.position;
         if let ControlNodeKind::Block { insts } = func_at_control_node.reborrow().def().kind {
             let mut func_at_inst_iter = func_at_control_node.reborrow().at(insts).into_iter();
             while let Some(mut func_at_inst) = func_at_inst_iter.next() {
                 match self.try_lower_data_inst_def(func_at_inst.reborrow(), control_node) {
                     Ok(Transformed::Changed(new_def)) => {
+                        // HACK(eddyb) this tracking is kind of ad-hoc but should
+                        // easily cover everything we care about for now.
+                        if let DataInstKind::QPtr(op) = &self.lowerer.cx[new_def.form].kind {
+                            match op {
+                                QPtrOp::HandleArrayIndex
+                                | QPtrOp::BufferData
+                                | QPtrOp::BufferDynLen { .. }
+                                | QPtrOp::Offset(_)
+                                | QPtrOp::DynOffset { .. } => {
+                                    self.remove_if_dead_inst_and_parent_block
+                                        .push((func_at_inst.position, control_node));
+                                }
+
+                                QPtrOp::FuncLocalVar(_)
+                                | QPtrOp::Load { .. }
+                                | QPtrOp::Store { .. } => {}
+                            }
+
+                            if let QPtrOp::Offset(0) = op {
+                                let mut base_ptr = new_def.inputs[0];
+                                if let Value::DataInstOutput(base_ptr_inst) = base_ptr {
+                                    if let Some(&base_ptr_base_ptr) =
+                                        self.noop_offsets_to_base_ptr.get(&base_ptr_inst)
+                                    {
+                                        base_ptr = base_ptr_base_ptr;
+                                    }
+                                }
+                                self.noop_offsets_to_base_ptr
+                                    .insert(func_at_inst.position, base_ptr);
+                            }
+                        }
+
                         *func_at_inst.def() = new_def;
                     }
                     result @ (Ok(Transformed::Unchanged) | Err(_)) => {
@@ -795,5 +898,41 @@ impl Transformer for LowerFromSpvPtrInstsInFunc<'_> {
                 }
             }
         }
+
+        // NOTE(eddyb) this is done last so that `transform_value_use` only sees
+        // the lowered `Value`s, not the original ones.
+        func_at_control_node
+            .reborrow()
+            .inner_in_place_transform_with(self);
+    }
+
+    fn in_place_transform_func_decl(&mut self, func_decl: &mut FuncDecl) {
+        func_decl.inner_in_place_transform_with(self);
+
+        // Apply all `remove_if_dead_inst_and_parent_block` removals, that are truly unused.
+        if let DeclDef::Present(func_def_body) = &mut func_decl.def {
+            let remove_if_dead_inst_and_parent_block =
+                mem::take(&mut self.remove_if_dead_inst_and_parent_block);
+            // NOTE(eddyb) reverse order is important, as each removal can reduce
+            // use counts of an earlier definition, allowing further removal.
+            for (inst, parent_block) in remove_if_dead_inst_and_parent_block.into_iter().rev() {
+                if self.data_inst_use_counts.get(inst).is_none() {
+                    // HACK(eddyb) can't really use helpers like `FuncAtMut::def`,
+                    // due to the need to borrow `control_nodes` and `data_insts`
+                    // at the same time - perhaps some kind of `FuncAtMut` position
+                    // types for "where a list is in a parent entity" could be used
+                    // to make this more ergonomic, although the potential need for
+                    // an actual list entity of its own, should be considered.
+                    match &mut func_def_body.control_nodes[parent_block].kind {
+                        ControlNodeKind::Block { insts } => {
+                            insts.remove(inst, &mut func_def_body.data_insts);
+                        }
+                        _ => unreachable!(),
+                    }
+
+                    self.remove_value_uses(&func_def_body.at(inst).def().inputs);
+                }
+            }
+        }
     }
 }