diff --git a/crates/acvm_backend_barretenberg/src/proof_system.rs b/crates/acvm_backend_barretenberg/src/proof_system.rs index 33ec8457a43..8a24559c867 100644 --- a/crates/acvm_backend_barretenberg/src/proof_system.rs +++ b/crates/acvm_backend_barretenberg/src/proof_system.rs @@ -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 = witness_values - .try_into() - .expect("could not serialize witness map"); + let serialized_witnesses: Vec = + 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); @@ -232,9 +231,8 @@ fn prepend_public_inputs(proof: Vec, public_inputs: Vec) -> 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() } diff --git a/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs b/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs index 84af1053ed7..6039c3b55ec 100644 --- a/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs +++ b/crates/noirc_evaluator/src/ssa/opt/mem2reg.rs @@ -61,6 +61,9 @@ //! 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::{ @@ -78,6 +81,9 @@ use crate::ssa::{ 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. @@ -107,45 +113,6 @@ struct PerFunctionContext<'f> { instructions_to_remove: BTreeSet, } -#[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, - - /// 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>, - - /// 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, - - /// The last instance of a `Store` instruction to each address in this block - last_stores: BTreeMap, -} - -/// 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), - ArrayElement(Box), - 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); @@ -234,9 +201,10 @@ impl<'f> PerFunctionContext<'f> { 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); } } @@ -290,17 +258,11 @@ impl<'f> PerFunctionContext<'f> { 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]; @@ -317,28 +279,34 @@ impl<'f> PerFunctionContext<'f> { } 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), _ => (), } } @@ -369,12 +337,11 @@ impl<'f> PerFunctionContext<'f> { } } - 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) { @@ -432,175 +399,6 @@ impl<'f> PerFunctionContext<'f> { } } -impl Block { - /// If the given reference id points to a known value, return the value - fn get_known_value(&self, address: ValueId) -> Option { - 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( - &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; diff --git a/crates/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs b/crates/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs new file mode 100644 index 00000000000..5477025e429 --- /dev/null +++ b/crates/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs @@ -0,0 +1,76 @@ +use std::collections::BTreeSet; + +use crate::ssa::ir::value::ValueId; + +/// A set of possible aliases. Each ValueId in this set represents one possible value the reference +/// holding this AliasSet may be aliased to. This struct wrapper is provided so that when we take +/// the union of multiple alias sets, the result should be empty if any individual set is empty. +/// +/// Note that we distinguish between "definitely has no aliases" - `Some(BTreeSet::new())`, and +/// "unknown which aliases this may refer to" - `None`. +#[derive(Debug, Default, Clone)] +pub(super) struct AliasSet { + aliases: Option>, +} + +impl AliasSet { + pub(super) fn unknown() -> AliasSet { + Self { aliases: None } + } + + pub(super) fn known(value: ValueId) -> AliasSet { + let mut aliases = BTreeSet::new(); + aliases.insert(value); + Self { aliases: Some(aliases) } + } + + /// In rare cases, such as when creating an empty array of references, the set of aliases for a + /// particular value will be known to be zero, which is distinct from being unknown and + /// possibly referring to any alias. + pub(super) fn known_empty() -> AliasSet { + Self { aliases: Some(BTreeSet::new()) } + } + + pub(super) fn is_unknown(&self) -> bool { + self.aliases.is_none() + } + + /// Return the single known alias if there is exactly one. + /// Otherwise, return None. + pub(super) fn single_alias(&self) -> Option { + self.aliases + .as_ref() + .and_then(|aliases| (aliases.len() == 1).then(|| *aliases.first().unwrap())) + } + + /// Unify this alias set with another. The result of this set is empty if either set is empty. + /// Otherwise, it is the union of both alias sets. + pub(super) fn unify(&mut self, other: &Self) { + if let (Some(self_aliases), Some(other_aliases)) = (&mut self.aliases, &other.aliases) { + self_aliases.extend(other_aliases); + } else { + self.aliases = None; + } + } + + /// Inserts a new alias into this set if it is not unknown + pub(super) fn insert(&mut self, new_alias: ValueId) { + if let Some(aliases) = &mut self.aliases { + aliases.insert(new_alias); + } + } + + /// Returns `Some(true)` if `f` returns true for any known alias in this set. + /// If this alias set is unknown, None is returned. + pub(super) fn any(&self, f: impl FnMut(ValueId) -> bool) -> Option { + self.aliases.as_ref().map(|aliases| aliases.iter().copied().any(f)) + } + + pub(super) fn for_each(&self, mut f: impl FnMut(ValueId)) { + if let Some(aliases) = &self.aliases { + for alias in aliases { + f(*alias); + } + } + } +} diff --git a/crates/noirc_evaluator/src/ssa/opt/mem2reg/block.rs b/crates/noirc_evaluator/src/ssa/opt/mem2reg/block.rs new file mode 100644 index 00000000000..22c5705b723 --- /dev/null +++ b/crates/noirc_evaluator/src/ssa/opt/mem2reg/block.rs @@ -0,0 +1,243 @@ +use std::{borrow::Cow, collections::BTreeMap}; + +use crate::ssa::ir::{ + function::Function, + instruction::{Instruction, InstructionId}, + value::ValueId, +}; + +use super::alias_set::AliasSet; + +/// A `Block` acts as a per-block context for the mem2reg pass. +/// Most notably, it contains the current alias set thought to track each +/// reference value if known, and it contains the expected ReferenceValue +/// for each ValueId. When a block is finished, the final values of these +/// are expected to match the values held by each ValueId at the very end +/// of a block. +#[derive(Debug, Default, Clone)] +pub(super) 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. + pub(super) expressions: BTreeMap, + + /// 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. + pub(super) aliases: BTreeMap, + + /// 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. + pub(super) references: BTreeMap, + + /// The last instance of a `Store` instruction to each address in this block + pub(super) last_stores: BTreeMap, +} + +/// 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)] +pub(super) enum Expression { + Dereference(Box), + ArrayElement(Box), + Other(ValueId), +} + +/// Every reference's value is either Known and can be optimized away, or Unknown. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub(super) enum ReferenceValue { + Unknown, + Known(ValueId), +} + +impl ReferenceValue { + fn unify(self, other: Self) -> Self { + if self == other { + self + } else { + ReferenceValue::Unknown + } + } +} + +impl Block { + /// If the given reference id points to a known value, return the value + pub(super) fn get_known_value(&self, address: ValueId) -> Option { + 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 let Some(alias) = aliases.single_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)`. + pub(super) fn set_known_value(&mut self, address: ValueId, value: ValueId) { + self.set_value(address, ReferenceValue::Known(value)); + } + + pub(super) 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_unknown() { + // 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 let Some(alias) = aliases.single_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. + aliases.for_each(|alias| { + 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(); + } + + pub(super) 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| aliases.unify(new_aliases)) + .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. + pub(super) 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( + &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() { + aliases.for_each(|alias| { + 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:?}") + } + } + } + } + + pub(super) 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); + } + } + } + + /// Collect all aliases used by the given value list + pub(super) fn collect_all_aliases( + &self, + values: impl IntoIterator, + ) -> AliasSet { + let mut aliases = AliasSet::known_empty(); + for value in values { + aliases.unify(&self.get_aliases_for_value(value)); + } + aliases + } + + pub(super) fn get_aliases_for_value(&self, value: ValueId) -> Cow { + if let Some(expression) = self.expressions.get(&value) { + if let Some(aliases) = self.aliases.get(expression) { + return Cow::Borrowed(aliases); + } + } + + Cow::Owned(AliasSet::unknown()) + } +}