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): mem2reg handles muts only #1654

Closed
wants to merge 3 commits into from
Closed

fix(ssa refactor): mem2reg handles muts only #1654

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8370156
fix(ssa refactor): mem2reg handles muts only
joss-aztec Jun 13, 2023
bf79349
chore(ssa refactor): tidy up
joss-aztec Jun 13, 2023
bd54fb4
Merge branch 'master' into joss/revise-mem2reg
joss-aztec 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
198 changes: 38 additions & 160 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
Expand Up @@ -3,71 +3,44 @@
//! mutable variables into values that are easier to manipulate.
use std::collections::{BTreeMap, BTreeSet};

use acvm::FieldElement;

use crate::ssa_refactor::{
ir::{
basic_block::BasicBlockId,
dfg::DataFlowGraph,
instruction::{BinaryOp, Instruction, InstructionId, TerminatorInstruction},
types::Type,
value::{Value, ValueId},
instruction::{Instruction, InstructionId, TerminatorInstruction},
value::ValueId,
},
ssa_gen::Ssa,
};

impl Ssa {
/// Attempts to remove any load instructions that recover values that already available in
/// scope, and attempts to remove store that are subsequently redundant, as long as they are
/// not stores on memory that will be passed into a function call.
///
/// This pass also assumes that constant folding has been run, such that all addresses given
/// as input to store/load instructions are represented as one of:
/// - a value that directly resolves to an allocate instruction
/// - a value that directly resolves to a binary add instruction which has a allocate
/// instruction and a numeric constant as its operands
/// Attempts to remove any load instructions that recover values that are already available in
/// scope. Also attempts to remove store instructions if the function contains only a single
/// block.
pub(crate) fn mem2reg(mut self) -> Ssa {
let mut first_context = None;

for function in self.functions.values_mut() {
let blocks = function.reachable_blocks();
let is_single_block = blocks.len() == 1;
for block in function.reachable_blocks() {
Comment on lines +22 to 24
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The second call to reachable_blocks can be replaced with blocks

let mut context = PerBlockContext::new(block);
context.eliminate_known_loads(&mut function.dfg);
first_context = Some(context);
if is_single_block {
// If this function has only a single block, we know that the side effects of a
// store instruction only have bearing within the scope of the block.
context.remove_unused_stores(&mut function.dfg);
}
}
}

// 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 {
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<ValueId, ValueId>,
store_ids: Vec<InstructionId>,
failed_substitutes: BTreeSet<Address>,
alloc_ids_used_externally: BTreeSet<InstructionId>,
failed_substitutes: BTreeSet<ValueId>,
}

impl PerBlockContext {
Expand All @@ -77,7 +50,6 @@ impl PerBlockContext {
last_stores: BTreeMap::new(),
store_ids: Vec::new(),
failed_substitutes: BTreeSet::new(),
alloc_ids_used_externally: BTreeSet::new(),
}
}

Expand All @@ -90,30 +62,19 @@ impl PerBlockContext {
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.push((*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());
}
for value in arguments {
assert!(!self.last_stores.contains_key(value), "Mutable vars are loaded before being passed as function arguments - if this pattern changes, so do our safety assumptions.");
}
}
_ => {
Expand All @@ -122,29 +83,24 @@ 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());
}
assert!(!self.last_stores.contains_key(value), "Mutable vars are loaded before being returned - if this pattern changes, so do our safety assumptions.");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd like to retain the ability to pass in or return mutable references to functions for implementing references later on. Does handling this cause any bugs that were present previously?

}
}

// Substitute load result values
for (instruction_id, new_value) in &loads_to_substitute {
let result_value = *dfg
.instruction_results(*instruction_id)
.first()
.expect("ICE: Load instructions should have single result");
dfg.set_value_from_id(result_value, *new_value);
let result_values = dfg.instruction_results(*instruction_id);
assert_eq!(result_values.len(), 1);
dfg.set_value_from_id(result_values[0], *new_value);
}

// 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);
for (instruction_id, _) in &loads_to_substitute {
// Technically we could leave this removal to the DIE pass, but the debug print is
// easier to read if we remove it now.
block.remove_instruction(*instruction_id);
}
}

Expand All @@ -156,13 +112,8 @@ impl PerBlockContext {
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) {
stores_to_remove.push(*instruction_id);
}
}

Expand All @@ -172,56 +123,6 @@ impl PerBlockContext {
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
}
}
}

#[cfg(test)]
Expand Down Expand Up @@ -270,8 +171,6 @@ mod tests {

let ssa = builder.finish().mem2reg().fold_constants();

println!("{ssa}");

let func = ssa.main();
let block_id = func.entry_block();

Expand All @@ -286,7 +185,8 @@ mod tests {
}

#[test]
fn test_simple_with_call() {
#[should_panic]
fn test_calls_disallowed() {
// fn func {
// b0():
// v0 = allocate
Expand All @@ -295,6 +195,7 @@ mod tests {
// call f0(v0)
// return v1
// }
// Passing a memory address as function arguments is unsupported

let func_id = Id::test_new(0);
let mut builder = FunctionBuilder::new("func".into(), func_id, RuntimeType::Acir);
Expand All @@ -306,23 +207,12 @@ mod tests {
builder.insert_call(f0, vec![v0], vec![]);
builder.terminate_with_return(vec![v1]);

let ssa = builder.finish().mem2reg();

let func = ssa.main();
let block_id = func.entry_block();

assert_eq!(count_loads(block_id, &func.dfg), 0);
assert_eq!(count_stores(block_id, &func.dfg), 1);

let ret_val_id = match func.dfg[block_id].terminator().unwrap() {
TerminatorInstruction::Return { return_values } => return_values.first().unwrap(),
_ => unreachable!(),
};
assert_eq!(func.dfg[*ret_val_id], func.dfg[one]);
builder.finish().mem2reg();
}

#[test]
fn test_simple_with_return() {
#[should_panic]
fn test_return_disallowed() {
// fn func {
// b0():
// v0 = allocate
Expand All @@ -337,19 +227,7 @@ mod tests {
builder.insert_store(v0, const_one);
builder.terminate_with_return(vec![v0]);

let ssa = builder.finish().mem2reg();

let func = ssa.main();
let block_id = func.entry_block();

// Store affects outcome of returned array, and can't be removed
assert_eq!(count_stores(block_id, &func.dfg), 1);

let ret_val_id = match func.dfg[block_id].terminator().unwrap() {
TerminatorInstruction::Return { return_values } => return_values.first().unwrap(),
_ => unreachable!(),
};
assert_eq!(func.dfg[*ret_val_id], func.dfg[v0]);
builder.finish().mem2reg();
}

fn count_stores(block: BasicBlockId, dfg: &DataFlowGraph) -> usize {
Expand Down