feat(perf): Remove last store in return block if last load is before that store

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

@@ -56,7 +56,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
@@ -111,12 +111,12 @@
     /// 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: BTreeSet<InstructionId>,
 
     /// Track a value's last load across all blocks.
     /// If a value is not used in anymore loads we can remove the last store to that value.
-    last_loads: HashMap<ValueId, InstructionId>,
+    last_loads: HashMap<ValueId, (InstructionId, BasicBlockId)>,
 }
 
 impl<'f> PerFunctionContext<'f> {
@@ -152,9 +152,31 @@
         // This rule does not apply to reference parameters, which we must also check for before removing these stores.
         for (block_id, block) in self.blocks.iter() {
             let block_params = self.inserter.function.dfg.block_parameters(*block_id);
-            for (value, store_instruction) in block.last_stores.iter() {
-                let is_reference_param = block_params.contains(value);
-                if self.last_loads.get(value).is_none() && !is_reference_param {
+            for (store_address, store_instruction) in block.last_stores.iter() {
+                let is_reference_param = block_params.contains(store_address);
+                let terminator = self.inserter.function.dfg[*block_id].unwrap_terminator();
+
+                let is_return = matches!(terminator, TerminatorInstruction::Return { .. });
+                let remove_load = if is_return {
+                    // Determine whether the last store is used in the return value
+                    let mut is_return_value = false;
+                    terminator.for_each_value(|return_value| {
+                        is_return_value = return_value == *store_address || is_return_value;
+                    });
+
+                    // If the last load of a store is not part of the block with a return terminator,
+                    // we can safely remove this store.
+                    let last_load_not_in_return = self
+                        .last_loads
+                        .get(store_address)
+                        .map(|(_, last_load_block)| *last_load_block != *block_id)
+                        .unwrap_or(true);
+                    !is_return_value && last_load_not_in_return
+                } else {
+                    self.last_loads.get(store_address).is_none()
+                };
+
+                if remove_load && !is_reference_param {
                     self.instructions_to_remove.insert(*store_instruction);
                 }
             }
@@ -259,7 +281,7 @@
                 } else {
                     references.mark_value_used(address, self.inserter.function);
 
-                    self.last_loads.insert(address, instruction);
+                    self.last_loads.insert(address, (instruction, block_id));
                 }
             }
             Instruction::Store { address, value } => {