feat(perf): Remove inc_rc/dec_rc instructions that follow a removed load in mem2reg

compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs

@@ -66,7 +66,7 @@
 //!
 //! Repeating this algorithm for each block in the function in program order should result in
 //! optimizing out most known loads. However, identifying all aliases correctly has been proven
 //! undecidable in general (Landi, 1992). So this pass will not always optimize out all loads
 //! that could theoretically be optimized out. This pass can be performed at any time in the
 //! SSA optimization pipeline, although it will be more successful the simpler the program's CFG is.
 //! This pass is currently performed several times to enable other passes - most notably being
@@ -154,7 +154,7 @@
     /// Load and Store instructions that should be removed at the end of the pass.
     ///
     /// We avoid removing individual instructions as we go since removing elements
     /// from the middle of Vecs many times will be slower than a single call to `retain`.
     instructions_to_remove: HashSet<InstructionId>,
 
     /// Track a value's last load across all blocks.
@@ -364,10 +364,12 @@
                 let result = self.inserter.function.dfg.instruction_results(instruction)[0];
                 references.remember_dereference(self.inserter.function, address, result);
 
+                let mut known_value_exists = false;
                 // If the load is known, replace it with the known value and remove the load
                 if let Some(value) = references.get_known_value(address) {
                     self.inserter.map_value(result, value);
                     self.instructions_to_remove.insert(instruction);
+                    known_value_exists = true;
                 } else {
                     references.mark_value_used(address, self.inserter.function);
 
@@ -420,12 +422,13 @@
                         self.inserter.map_value(result, previous_result);
                         self.instructions_to_remove.insert(instruction);
                     }
+                } else if !known_value_exists {
+                    // We want to set the load for every load even if the address has a known value
+                    // and the previous load instruction was removed.
+                    // We are safe to still remove a repeat load in this case as we are mapping from the current load's
+                    // result to the previous load, which if it was removed should already have a mapping to the known value.
+                    references.set_last_load(address, instruction);
                 }
-                // We want to set the load for every load even if the address has a known value
-                // and the previous load instruction was removed.
-                // We are safe to still remove a repeat load in this case as we are mapping from the current load's
-                // result to the previous load, which if it was removed should already have a mapping to the known value.
-                references.set_last_load(address, instruction);
             }
             Instruction::Store { address, value } => {
                 let address = self.inserter.function.dfg.resolve(*address);
@@ -547,6 +550,33 @@
                 }
                 self.mark_all_unknown(arguments, references);
             }
+            Instruction::IncrementRc { value } | Instruction::DecrementRc { value } => {
+                // Account for whether a load has been removed.
+                // If a load has been removed, there may potentially be an inc_rc/dec_rc instruction
+                // that follows the load instruction. With the load gone, we can also remove the inc_rc/dec_rc.
+                let instructions = self.inserter.function.dfg[block_id].instructions();
+                if instructions.len() < 2 {
+                    return;
+                }
+                // We subtract two as we always push an instruction before analyzing it and we want the instruction
+                // before this increment or decrement rc instruction.
+                let last_instruction = instructions[instructions.len() - 2];
+                let value = self.inserter.function.dfg.resolve(*value);
+                if self.instructions_to_remove.contains(&last_instruction) {
+                    if let Instruction::Load { address } =
+                        self.inserter.function.dfg[last_instruction]
+                    {
+                        let result =
+                            self.inserter.function.dfg.instruction_results(last_instruction)[0];
+                        let result = self.inserter.resolve(result);
+                        // We only want to remove the inc/dec rc if the removed load came from a repeated load
+                        if result == value && references.last_loads.get(&address).is_some() {
+                            self.instructions_to_remove.insert(instruction);
+                        }
+                    }
+                }
+            }
+
             _ => (),
         }
     }