Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(ssa refactor): Fix mem2reg pass not always removing unused stores #1677

Merged
merged 3 commits into from
Jun 14, 2023
Merged

fix(ssa refactor): Fix mem2reg pass not always removing unused stores #1677

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fb6acd0
Fix mem2reg pass not always removing unused stores
jfecher Jun 13, 2023
a2f9e0b
Merge branch 'master' into jf/mem2reg
jfecher Jun 13, 2023
038aecd
Remove instructions more efficiently
jfecher Jun 13, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 41 additions & 107 deletions crates/noirc_evaluator/src/ssa_refactor/opt/mem2reg.rs
View file
Open in desktop
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