Skip to content

Commit

Permalink
fix(ssa refactor): Fix mem2reg pass not always removing unused stores (…
Browse files Browse the repository at this point in the history 
…#1677)

* Fix mem2reg pass not always removing unused stores

* Remove instructions more efficiently
  • Loading branch information
@jfecher
jfecher committed Jun 14, 2023
1 parent 5d4dd75 commit 8310544
Showing 1 changed file with 41 additions and 107 deletions.
148 changes: 41 additions & 107 deletions crates/noirc_evaluator/src/ssa_refactor/opt/mem2reg.rs
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,13 @@
//! mem2reg implements a pass for promoting values stored in memory to values in registers where
//! possible. This is particularly important for converting our memory-based representation of
//! mutable variables into values that are easier to manipulate.
use std::collections::{BTreeMap, BTreeSet};

use acvm::FieldElement;
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};

use crate::ssa_refactor::{
ir::{
basic_block::BasicBlockId,
dfg::DataFlowGraph,
instruction::{BinaryOp, Instruction, InstructionId, TerminatorInstruction},
types::Type,
instruction::{Instruction, InstructionId, TerminatorInstruction},
value::{Value, ValueId},
},
ssa_gen::Ssa,
Expand Down Expand Up @@ -39,37 +36,27 @@ impl Ssa {

// If there is only one block in total, remove any unused stores as well since we
// know there is no other block they can impact.
if self.functions.len() == 1 && self.main().dfg.basic_blocks_iter().len() == 1 {
if self.functions.len() == 1 && self.main().reachable_blocks().len() == 1 {
first_context.unwrap().remove_unused_stores(&mut self.main_mut().dfg);
}

self
}
}

#[derive(PartialEq, PartialOrd, Eq, Ord)]
enum Address {
Zeroth(InstructionId),
Offset(InstructionId, FieldElement),
}

impl Address {
fn alloc_id(&self) -> InstructionId {
match self {
Address::Zeroth(alloc_id) => *alloc_id,
Address::Offset(alloc_id, _) => *alloc_id,
}
}
}

struct PerBlockContext {
block_id: BasicBlockId,
last_stores: BTreeMap<Address, ValueId>,
last_stores: BTreeMap<AllocId, ValueId>,
store_ids: Vec<InstructionId>,
failed_substitutes: BTreeSet<Address>,
alloc_ids_used_externally: BTreeSet<InstructionId>,
failed_substitutes: BTreeSet<AllocId>,
alloc_ids_used_externally: BTreeSet<AllocId>,
}

/// An AllocId is the ValueId returned from an allocate instruction. E.g. v0 in v0 = allocate.
/// This type alias is used to help signal where the only valid ValueIds are those that are from
/// an allocate instruction.
type AllocId = ValueId;

impl PerBlockContext {
fn new(block_id: BasicBlockId) -> Self {
PerBlockContext {
Expand All @@ -84,35 +71,26 @@ impl PerBlockContext {
// Attempts to remove load instructions for which the result is already known from previous
// store instructions to the same address in the same block.
fn eliminate_known_loads(&mut self, dfg: &mut DataFlowGraph) {
let mut loads_to_substitute = Vec::new();
let mut loads_to_substitute = HashMap::new();
let block = &dfg[self.block_id];

for instruction_id in block.instructions() {
match &dfg[*instruction_id] {
Instruction::Store { address, value } => {
if let Some(address) = self.try_const_address(*address, dfg) {
// We can only track the address if it is a constant offset from an
// allocation. A previous constant folding pass should make such addresses
// possible to identify.
self.last_stores.insert(address, *value);
}
// TODO: Consider if it's worth falling back to storing addresses by their
// value id such we can shallowly check for dynamic address reuse.
self.last_stores.insert(*address, *value);
self.store_ids.push(*instruction_id);
}
Instruction::Load { address } => {
if let Some(address) = self.try_const_address(*address, dfg) {
if let Some(last_value) = self.last_stores.get(&address) {
loads_to_substitute.push((*instruction_id, *last_value));
} else {
self.failed_substitutes.insert(address);
}
if let Some(last_value) = self.last_stores.get(address) {
loads_to_substitute.insert(*instruction_id, *last_value);
} else {
self.failed_substitutes.insert(*address);
}
}
Instruction::Call { arguments, .. } => {
for arg in arguments {
if let Some(address) = self.try_const_address(*arg, dfg) {
self.alloc_ids_used_externally.insert(address.alloc_id());
if Self::value_is_from_allocation(*arg, dfg) {
self.alloc_ids_used_externally.insert(*arg);
}
}
}
Expand All @@ -125,8 +103,8 @@ impl PerBlockContext {
// Identify any arrays that are returned from this function
if let TerminatorInstruction::Return { return_values } = block.unwrap_terminator() {
for value in return_values {
if let Some(address) = self.try_const_address(*value, dfg) {
self.alloc_ids_used_externally.insert(address.alloc_id());
if Self::value_is_from_allocation(*value, dfg) {
self.alloc_ids_used_externally.insert(*value);
}
}
}
Expand All @@ -142,85 +120,41 @@ impl PerBlockContext {

// Delete load instructions
// TODO: should we let DCE do this instead?
let block = &mut dfg[self.block_id];
for (instruction_id, _) in loads_to_substitute {
block.remove_instruction(instruction_id);
dfg[self.block_id]
.instructions_mut()
.retain(|instruction| !loads_to_substitute.contains_key(instruction));
}

fn value_is_from_allocation(value: ValueId, dfg: &DataFlowGraph) -> bool {
match &dfg[value] {
Value::Instruction { instruction, .. } => {
matches!(&dfg[*instruction], Instruction::Allocate)
}
_ => false,
}
}

fn remove_unused_stores(self, dfg: &mut DataFlowGraph) {
// Scan for unused stores
let mut stores_to_remove: Vec<InstructionId> = Vec::new();
let mut stores_to_remove = HashSet::new();

for instruction_id in &self.store_ids {
let address = match &dfg[*instruction_id] {
Instruction::Store { address, .. } => *address,
_ => unreachable!("store_ids should contain only store instructions"),
};

if let Some(address) = self.try_const_address(address, dfg) {
if !self.failed_substitutes.contains(&address)
&& !self.alloc_ids_used_externally.contains(&address.alloc_id())
{
stores_to_remove.push(*instruction_id);
}
if !self.failed_substitutes.contains(&address)
&& !self.alloc_ids_used_externally.contains(&address)
{
stores_to_remove.insert(*instruction_id);
}
}

// Delete unused stores
let block = &mut dfg[self.block_id];
for instruction_id in stores_to_remove {
block.remove_instruction(instruction_id);
}
}

// Attempts to normalize the given value into a const address
fn try_const_address(&self, value_id: ValueId, dfg: &DataFlowGraph) -> Option<Address> {
if dfg.type_of_value(value_id) != Type::Reference {
return None;
}
let value = &dfg[value_id];
let instruction_id = match value {
Value::Instruction { instruction, .. } => *instruction,
_ => return None,
};
let instruction = &dfg[instruction_id];
match instruction {
// Arrays can be returned by allocations and function calls
Instruction::Allocate { .. } | Instruction::Call { .. } => {
Some(Address::Zeroth(instruction_id))
}
Instruction::Binary(binary) => {
if binary.operator != BinaryOp::Add {
return None;
}
let lhs = &dfg[binary.lhs];
let rhs = &dfg[binary.rhs];
self.try_const_address_offset(lhs, rhs, dfg)
.or_else(|| self.try_const_address_offset(rhs, lhs, dfg))
}
_ => None,
}
}

// Tries val1 as an allocation instruction id and val2 as a constant offset
fn try_const_address_offset(
&self,
val1: &Value,
val2: &Value,
dfg: &DataFlowGraph,
) -> Option<Address> {
let alloc_id = match val1 {
Value::Instruction { instruction, .. } => match &dfg[*instruction] {
Instruction::Allocate { .. } => *instruction,
_ => return None,
},
_ => return None,
};
if let Value::NumericConstant { constant, .. } = val2 {
Some(Address::Offset(alloc_id, *constant))
} else {
None
}
dfg[self.block_id]
.instructions_mut()
.retain(|instruction| !stores_to_remove.contains(instruction));
}
}

Expand Down

0 comments on commit 8310544

Please sign in to comment.