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

Jump to bottom

chore: Cleanup mem2reg pass #2531

Merged
merged 3 commits into from
Sep 5, 2023
Merged

chore: Cleanup mem2reg pass #2531

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d4892ae
Create AliasSet concept
jfecher Aug 31, 2023
0993d1c
Formatting
jfecher Sep 1, 2023
97ed693
Move new types into modules
jfecher Sep 1, 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
10 changes: 4 additions & 6 deletions crates/acvm_backend_barretenberg/src/proof_system.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,9 +76,8 @@ impl ProofSystemCompiler for Barretenberg {
let temp_dir_path_str = temp_directory.to_str().unwrap();

// Create a temporary file for the witness
let serialized_witnesses: Vec<u8> = witness_values
.try_into()
.expect("could not serialize witness map");
let serialized_witnesses: Vec<u8> =
witness_values.try_into().expect("could not serialize witness map");
let witness_path = temp_directory.join("witness").with_extension("tr");
write_to_file(&serialized_witnesses, &witness_path);

Expand Down Expand Up @@ -232,9 +231,8 @@ fn prepend_public_inputs(proof: Vec<u8>, public_inputs: Vec<FieldElement>) -> Ve
return proof;
}

let public_inputs_bytes = public_inputs
.into_iter()
.flat_map(|assignment| assignment.to_be_bytes());
let public_inputs_bytes =
public_inputs.into_iter().flat_map(|assignment| assignment.to_be_bytes());

public_inputs_bytes.chain(proof.into_iter()).collect()
}
Expand Down
264 changes: 31 additions & 233 deletions crates/noirc_evaluator/src/ssa/opt/mem2reg.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,11 +56,14 @@
//!
//! 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

Check warning on line 59 in crates/noirc_evaluator/src/ssa/opt/mem2reg.rs

View workflow job for this annotation

GitHub Actions / Spellcheck / Spellcheck 

Unknown word (Landi)
//! 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
//! performed before loop unrolling to try to allow for mutable variables used for loop indices.
mod alias_set;
mod block;

use std::collections::{BTreeMap, BTreeSet};

use crate::ssa::{
Expand All @@ -78,6 +81,9 @@
ssa_gen::Ssa,
};

use self::alias_set::AliasSet;
use self::block::{Block, Expression};

impl Ssa {
/// Attempts to remove any load instructions that recover values that are already available in
/// scope, and attempts to remove stores that are subsequently redundant.
Expand All @@ -103,49 +109,10 @@
/// 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`.

Check warning on line 112 in crates/noirc_evaluator/src/ssa/opt/mem2reg.rs

View workflow job for this annotation

GitHub Actions / Spellcheck / Spellcheck 

Unknown word (Vecs)
instructions_to_remove: BTreeSet<InstructionId>,
}

#[derive(Debug, Default, Clone)]
struct Block {
/// Maps a ValueId to the Expression it represents.
/// Multiple ValueIds can map to the same Expression, e.g.
/// dereferences to the same allocation.
expressions: BTreeMap<ValueId, Expression>,

/// Each expression is tracked as to how many aliases it
/// may have. If there is only 1, we can attempt to optimize
/// out any known loads to that alias. Note that "alias" here
/// includes the original reference as well.
aliases: BTreeMap<Expression, BTreeSet<ValueId>>,

/// Each allocate instruction result (and some reference block parameters)
/// will map to a Reference value which tracks whether the last value stored
/// to the reference is known.
references: BTreeMap<ValueId, ReferenceValue>,

/// The last instance of a `Store` instruction to each address in this block
last_stores: BTreeMap<ValueId, InstructionId>,
}

/// An `Expression` here is used to represent a canonical key
/// into the aliases map since otherwise two dereferences of the
/// same address will be given different ValueIds.
#[derive(Debug, Clone, PartialOrd, Ord, PartialEq, Eq)]
enum Expression {
Dereference(Box<Expression>),
ArrayElement(Box<Expression>),
Other(ValueId),
}

/// Every reference's value is either Known and can be optimized away, or Unknown.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum ReferenceValue {
Unknown,
Known(ValueId),
}

impl<'f> PerFunctionContext<'f> {
fn new(function: &'f mut Function) -> Self {
let cfg = ControlFlowGraph::with_function(function);
Expand Down Expand Up @@ -234,9 +201,10 @@
if let Some(expression) = references.expressions.get(allocation) {
if let Some(aliases) = references.aliases.get(expression) {
let allocation_aliases_parameter =
aliases.iter().any(|alias| reference_parameters.contains(alias));
aliases.any(|alias| reference_parameters.contains(&alias));

if !aliases.is_empty() && !allocation_aliases_parameter {
// If `allocation_aliases_parameter` is known to be false
if allocation_aliases_parameter == Some(false) {
self.instructions_to_remove.insert(*instruction);
}
}
Expand Down Expand Up @@ -290,17 +258,11 @@
references.set_known_value(address, value);
references.last_stores.insert(address, instruction);
}
Instruction::Call { arguments, .. } => {
self.mark_all_unknown(arguments, references);
}
Instruction::Allocate => {
// Register the new reference
let result = self.inserter.function.dfg.instruction_results(instruction)[0];
references.expressions.insert(result, Expression::Other(result));

let mut aliases = BTreeSet::new();
aliases.insert(result);
references.aliases.insert(Expression::Other(result), aliases);
references.aliases.insert(Expression::Other(result), AliasSet::known(result));
}
Instruction::ArrayGet { array, .. } => {
let result = self.inserter.function.dfg.instruction_results(instruction)[0];
Expand All @@ -317,28 +279,34 @@
}
Instruction::ArraySet { array, value, .. } => {
references.mark_value_used(*array, self.inserter.function);
let element_type = self.inserter.function.dfg.type_of_value(*value);

if self.inserter.function.dfg.value_is_reference(*value) {
if Self::contains_references(&element_type) {
let result = self.inserter.function.dfg.instruction_results(instruction)[0];
let array = self.inserter.function.dfg.resolve(*array);

let expression = Expression::ArrayElement(Box::new(Expression::Other(array)));

if let Some(aliases) = references.aliases.get_mut(&expression) {
aliases.insert(result);
let mut aliases = if let Some(aliases) = references.aliases.get_mut(&expression)
{
aliases.clone()
} else if let Some((elements, _)) =
self.inserter.function.dfg.get_array_constant(array)
{
// TODO: This should be a unification of each alias set
// If any are empty, the whole should be as well.
for reference in elements {
self.try_add_alias(references, reference, array);
}
}
let aliases = references.collect_all_aliases(elements);
self.set_aliases(references, array, aliases.clone());
aliases
} else {
AliasSet::unknown()
};

aliases.unify(&references.get_aliases_for_value(*value));

references.expressions.insert(result, expression);
references.expressions.insert(result, expression.clone());
references.aliases.insert(expression, aliases);
}
}
Instruction::Call { arguments, .. } => self.mark_all_unknown(arguments, references),
_ => (),
}
}
Expand Down Expand Up @@ -369,12 +337,11 @@
}
}

fn try_add_alias(&self, references: &mut Block, reference: ValueId, alias: ValueId) {
if let Some(expression) = references.expressions.get(&reference) {
if let Some(aliases) = references.aliases.get_mut(expression) {
aliases.insert(alias);
}
}
fn set_aliases(&self, references: &mut Block, address: ValueId, new_aliases: AliasSet) {
let expression =
references.expressions.entry(address).or_insert(Expression::Other(address));
let aliases = references.aliases.entry(expression.clone()).or_default();
*aliases = new_aliases;
}

fn mark_all_unknown(&self, values: &[ValueId], references: &mut Block) {
Expand Down Expand Up @@ -432,175 +399,6 @@
}
}

impl Block {
/// If the given reference id points to a known value, return the value
fn get_known_value(&self, address: ValueId) -> Option<ValueId> {
if let Some(expression) = self.expressions.get(&address) {
if let Some(aliases) = self.aliases.get(expression) {
// We could allow multiple aliases if we check that the reference
// value in each is equal.
if aliases.len() == 1 {
let alias = aliases.first().expect("There should be exactly 1 alias");

if let Some(ReferenceValue::Known(value)) = self.references.get(alias) {
return Some(*value);
}
}
}
}
None
}

/// If the given address is known, set its value to `ReferenceValue::Known(value)`.
fn set_known_value(&mut self, address: ValueId, value: ValueId) {
self.set_value(address, ReferenceValue::Known(value));
}

fn set_unknown(&mut self, address: ValueId) {
self.set_value(address, ReferenceValue::Unknown);
}

fn set_value(&mut self, address: ValueId, value: ReferenceValue) {
let expression = self.expressions.entry(address).or_insert(Expression::Other(address));
let aliases = self.aliases.entry(expression.clone()).or_default();

if aliases.is_empty() {
// uh-oh, we don't know at all what this reference refers to, could be anything.
// Now we have to invalidate every reference we know of
self.invalidate_all_references();
} else if aliases.len() == 1 {
let alias = aliases.first().expect("There should be exactly 1 alias");
self.references.insert(*alias, value);
} else {
// More than one alias. We're not sure which it refers to so we have to
// conservatively invalidate all references it may refer to.
for alias in aliases.iter() {
if let Some(reference_value) = self.references.get_mut(alias) {
*reference_value = ReferenceValue::Unknown;
}
}
}
}

fn invalidate_all_references(&mut self) {
for reference_value in self.references.values_mut() {
*reference_value = ReferenceValue::Unknown;
}

self.last_stores.clear();
}

fn unify(mut self, other: &Self) -> Self {
for (value_id, expression) in &other.expressions {
if let Some(existing) = self.expressions.get(value_id) {
assert_eq!(existing, expression, "Expected expressions for {value_id} to be equal");
} else {
self.expressions.insert(*value_id, expression.clone());
}
}

for (expression, new_aliases) in &other.aliases {
let expression = expression.clone();

self.aliases
.entry(expression)
.and_modify(|aliases| {
for alias in new_aliases {
aliases.insert(*alias);
}
})
.or_insert_with(|| new_aliases.clone());
}

// Keep only the references present in both maps.
let mut intersection = BTreeMap::new();
for (value_id, reference) in &other.references {
if let Some(existing) = self.references.get(value_id) {
intersection.insert(*value_id, existing.unify(*reference));
}
}
self.references = intersection;

self
}

/// Remember that `result` is the result of dereferencing `address`. This is important to
/// track aliasing when references are stored within other references.
fn remember_dereference(&mut self, function: &Function, address: ValueId, result: ValueId) {
if function.dfg.value_is_reference(result) {
if let Some(known_address) = self.get_known_value(address) {
self.expressions.insert(result, Expression::Other(known_address));
} else {
let expression = Expression::Dereference(Box::new(Expression::Other(address)));
self.expressions.insert(result, expression);
// No known aliases to insert for this expression... can we find an alias
// even if we don't have a known address? If not we'll have to invalidate all
// known references if this reference is ever stored to.
}
}
}

/// Iterate through each known alias of the given address and apply the function `f` to each.
fn for_each_alias_of<T>(
&mut self,
address: ValueId,
mut f: impl FnMut(&mut Self, ValueId) -> T,
) {
if let Some(expr) = self.expressions.get(&address) {
if let Some(aliases) = self.aliases.get(expr).cloned() {
for alias in aliases {
f(self, alias);
}
}
}
}

fn keep_last_stores_for(&mut self, address: ValueId, function: &Function) {
let address = function.dfg.resolve(address);
self.keep_last_store(address, function);
self.for_each_alias_of(address, |t, alias| t.keep_last_store(alias, function));
}

fn keep_last_store(&mut self, address: ValueId, function: &Function) {
let address = function.dfg.resolve(address);

if let Some(instruction) = self.last_stores.remove(&address) {
// Whenever we decide we want to keep a store instruction, we also need
// to go through its stored value and mark that used as well.
match &function.dfg[instruction] {
Instruction::Store { value, .. } => {
self.mark_value_used(*value, function);
}
other => {
unreachable!("last_store held an id of a non-store instruction: {other:?}")
}
}
}
}

fn mark_value_used(&mut self, value: ValueId, function: &Function) {
self.keep_last_stores_for(value, function);

// We must do a recursive check for arrays since they're the only Values which may contain
// other ValueIds.
if let Some((array, _)) = function.dfg.get_array_constant(value) {
for value in array {
self.mark_value_used(value, function);
}
}
}
}

impl ReferenceValue {
fn unify(self, other: Self) -> Self {
if self == other {
self
} else {
ReferenceValue::Unknown
}
}
}

#[cfg(test)]
mod tests {
use std::rc::Rc;
Expand Down Expand Up @@ -730,8 +528,8 @@

// Since the mem2reg pass simplifies as it goes, the id of the allocate instruction result
// is most likely no longer v0. We have to retrieve the new id here.
let alloca_id = func.dfg.instruction_results(instructions[0])[0];

Check warning on line 531 in crates/noirc_evaluator/src/ssa/opt/mem2reg.rs

View workflow job for this annotation

GitHub Actions / Spellcheck / Spellcheck 

Unknown word (alloca)
assert_eq!(ret_val_id, alloca_id);

Check warning on line 532 in crates/noirc_evaluator/src/ssa/opt/mem2reg.rs

View workflow job for this annotation

GitHub Actions / Spellcheck / Spellcheck 

Unknown word (alloca)
}

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