mem2reg: expressions -> aliases indirection might not be needed? #9482
Description
In mem2reg, in each block we track what Expression a value maps to, and then what aliases that Expression has:
noir/compiler/noirc_evaluator/src/ssa/opt/mem2reg/block.rs
Lines 19 to 28 in b51b616
However, it's not clearly explained why this indirection is needed. Would mapping each value to an AliasSet be enough? My guess is "no" because if we detect that a value's related AliasSet needs to be modified, we'd need to modify it too for every value in that AliasSet. That's when the Expression indirection comes into play: we map all the aliased values to a single expression, then when we modify that expression's AliasSet it's automatically modified for all values that map to that expression.
Or maybe it's done to avoid having multiple identical AliasSets for a list of values, that is, for compilation performance reasons? It might be good to document this design decision. (though I checked the time and memory to compile token_contract and ram_blowup_regression and there was no noticeable change)
What's interesting is that if we remove this indirection all tests still pass (and no bytecode changes are produced). That means that either the indirection is not needed (at least not for correctness) or that a test is missing proving that this indirection is in fact needed. If the indirection is not needed we could consider removing it as it slightly simplifies the mem2reg logic.
Here's the patch to remove the indirection in case you want to try running all tests:
Patch to remove indirection
diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs
index 6dd27aa282..b5af9ecd0a 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs
@@ -97,7 +97,7 @@ use crate::ssa::{
};
use self::alias_set::AliasSet;
-use self::block::{Block, Expression};
+use self::block::Block;
impl Ssa {
/// Attempts to remove any load instructions that recover values that are already available in
@@ -192,12 +192,10 @@ impl<'f> PerFunctionContext<'f> {
// then that reference gets to keep its last store.
let typ = self.inserter.function.dfg.type_of_value(value);
if Self::contains_references(&typ) {
- if let Some(expression) = references.expressions.get(&value) {
- if let Some(aliases) = references.aliases.get(expression) {
- aliases.for_each(|alias| {
- all_terminator_values.insert(alias);
- });
- }
+ if let Some(aliases) = references.aliases.get(&value) {
+ aliases.for_each(|alias| {
+ all_terminator_values.insert(alias);
+ });
}
}
});
@@ -216,10 +214,7 @@ impl<'f> PerFunctionContext<'f> {
&per_func_block_params,
);
- let is_dereference = block
- .expressions
- .get(store_address)
- .is_some_and(|expression| matches!(expression, Expression::Dereference(_)));
+ let is_dereference = block.dereferences.contains(store_address);
if !self.last_loads.contains_key(store_address)
&& !store_alias_used
@@ -244,47 +239,45 @@ impl<'f> PerFunctionContext<'f> {
let reference_parameters = self.reference_parameters();
// Check whether the store address has an alias that crosses an entry point boundary (e.g. a Call or Return)
- if let Some(expression) = block.expressions.get(store_address) {
- if let Some(aliases) = block.aliases.get(expression) {
- let allocation_aliases_parameter =
- aliases.any(|alias| reference_parameters.contains(&alias));
- if allocation_aliases_parameter == Some(true) {
- return true;
- }
-
- let allocation_aliases_parameter =
- aliases.any(|alias| per_func_block_params.contains(&alias));
- if allocation_aliases_parameter == Some(true) {
- return true;
- }
+ if let Some(aliases) = block.aliases.get(store_address) {
+ let allocation_aliases_parameter =
+ aliases.any(|alias| reference_parameters.contains(&alias));
+ if allocation_aliases_parameter == Some(true) {
+ return true;
+ }
- // Is any alias of this address an input to some function call, or a return value?
- let allocation_aliases_instr_input =
- aliases.any(|alias| self.instruction_input_references.contains(&alias));
- if allocation_aliases_instr_input == Some(true) {
- return true;
- }
+ let allocation_aliases_parameter =
+ aliases.any(|alias| per_func_block_params.contains(&alias));
+ if allocation_aliases_parameter == Some(true) {
+ return true;
+ }
- // Is any alias of this address used in a block terminator?
- let allocation_aliases_terminator_args =
- aliases.any(|alias| all_terminator_values.contains(&alias));
- if allocation_aliases_terminator_args == Some(true) {
- return true;
- }
+ // Is any alias of this address an input to some function call, or a return value?
+ let allocation_aliases_instr_input =
+ aliases.any(|alias| self.instruction_input_references.contains(&alias));
+ if allocation_aliases_instr_input == Some(true) {
+ return true;
+ }
- // Check whether there are any aliases whose instructions are not all marked for removal.
- // If there is any alias marked to survive, we should not remove its last store.
- let has_alias_not_marked_for_removal = aliases.any(|alias| {
- if let Some(alias_instructions) = self.aliased_references.get(&alias) {
- !alias_instructions.is_subset(&self.instructions_to_remove)
- } else {
- false
- }
- });
+ // Is any alias of this address used in a block terminator?
+ let allocation_aliases_terminator_args =
+ aliases.any(|alias| all_terminator_values.contains(&alias));
+ if allocation_aliases_terminator_args == Some(true) {
+ return true;
+ }
- if has_alias_not_marked_for_removal == Some(true) {
- return true;
+ // Check whether there are any aliases whose instructions are not all marked for removal.
+ // If there is any alias marked to survive, we should not remove its last store.
+ let has_alias_not_marked_for_removal = aliases.any(|alias| {
+ if let Some(alias_instructions) = self.aliased_references.get(&alias) {
+ !alias_instructions.is_subset(&self.instructions_to_remove)
+ } else {
+ false
}
+ });
+
+ if has_alias_not_marked_for_removal == Some(true) {
+ return true;
}
}
@@ -384,15 +377,8 @@ impl<'f> PerFunctionContext<'f> {
}
for aliases in aliases.into_values() {
- let first = aliases.first();
- let first = first.expect("All parameters alias at least themselves or we early return");
-
- let expression = Expression::Other(first);
- let previous = references.aliases.insert(expression.clone(), aliases.clone());
- assert!(previous.is_none());
-
aliases.for_each(|alias| {
- let previous = references.expressions.insert(alias, expression.clone());
+ let previous = references.aliases.insert(alias, aliases.clone());
assert!(previous.is_none());
});
}
@@ -404,14 +390,12 @@ impl<'f> PerFunctionContext<'f> {
let reference_parameters = self.reference_parameters();
for (allocation, instruction) in &references.last_stores {
- if let Some(expression) = references.expressions.get(allocation) {
- if let Some(aliases) = references.aliases.get(expression) {
- let allocation_aliases_parameter =
- aliases.any(|alias| reference_parameters.contains(&alias));
- // If `allocation_aliases_parameter` is known to be false
- if allocation_aliases_parameter == Some(false) {
- self.instructions_to_remove.insert(*instruction);
- }
+ if let Some(aliases) = references.aliases.get(allocation) {
+ let allocation_aliases_parameter =
+ aliases.any(|alias| reference_parameters.contains(&alias));
+ // If `allocation_aliases_parameter` is known to be false
+ if allocation_aliases_parameter == Some(false) {
+ self.instructions_to_remove.insert(*instruction);
}
}
}
@@ -517,8 +501,7 @@ impl<'f> PerFunctionContext<'f> {
Instruction::Allocate => {
// Register the new reference
let result = self.inserter.function.dfg.instruction_results(instruction)[0];
- references.expressions.insert(result, Expression::Other(result));
- references.aliases.insert(Expression::Other(result), AliasSet::known(result));
+ references.aliases.insert(result, AliasSet::known(result));
}
Instruction::ArrayGet { array, .. } => {
let result = self.inserter.function.dfg.instruction_results(instruction)[0];
@@ -531,9 +514,7 @@ impl<'f> PerFunctionContext<'f> {
});
references.mark_value_used(array, self.inserter.function);
- let expression = Expression::ArrayElement(Box::new(Expression::Other(array)));
-
- if let Some(aliases) = references.aliases.get_mut(&expression) {
+ if let Some(aliases) = references.aliases.get_mut(&array) {
aliases.insert(result);
}
}
@@ -543,13 +524,9 @@ impl<'f> PerFunctionContext<'f> {
let element_type = self.inserter.function.dfg.type_of_value(*value);
if Self::contains_references(&element_type) {
- let result = self.inserter.function.dfg.instruction_results(instruction)[0];
let array = *array;
- let expression = Expression::ArrayElement(Box::new(Expression::Other(array)));
-
- let mut aliases = if let Some(aliases) = references.aliases.get_mut(&expression)
- {
+ let mut aliases = if let Some(aliases) = references.aliases.get_mut(&array) {
aliases.clone()
} else if let Some((elements, _)) =
self.inserter.function.dfg.get_array_constant(array)
@@ -563,8 +540,7 @@ impl<'f> PerFunctionContext<'f> {
aliases.unify(&references.get_aliases_for_value(*value));
- references.expressions.insert(result, expression.clone());
- references.aliases.insert(expression, aliases);
+ references.aliases.insert(array, aliases);
// Similar to how we remember that we used a value in a `Store` instruction,
// take note that it was used in the `ArraySet`. If this instruction is not
@@ -590,9 +566,8 @@ impl<'f> PerFunctionContext<'f> {
if Self::contains_references(typ) {
let array = self.inserter.function.dfg.instruction_results(instruction)[0];
- let expr = Expression::ArrayElement(Box::new(Expression::Other(array)));
- references.expressions.insert(array, expr.clone());
- let aliases = references.aliases.entry(expr).or_insert(AliasSet::known_empty());
+ let aliases =
+ references.aliases.entry(array).or_insert(AliasSet::known_empty());
self.add_array_aliases(elements, aliases);
}
@@ -602,10 +577,8 @@ impl<'f> PerFunctionContext<'f> {
let result_type = self.inserter.function.dfg.type_of_value(result);
if Self::contains_references(&result_type) {
- let expr = Expression::Other(result);
- references.expressions.insert(result, expr.clone());
references.aliases.insert(
- expr,
+ result,
AliasSet::known_multiple(vec![*then_value, *else_value].into()),
);
}
@@ -638,9 +611,7 @@ impl<'f> PerFunctionContext<'f> {
}
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();
+ let aliases = references.aliases.entry(address).or_default();
*aliases = new_aliases;
}
@@ -696,18 +667,16 @@ impl<'f> PerFunctionContext<'f> {
if self.inserter.function.dfg.value_is_reference(*parameter) {
let argument = *argument;
- if let Some(expression) = references.expressions.get(&argument) {
- if let Some(aliases) = references.aliases.get_mut(expression) {
- // The argument reference is possibly aliased by this block parameter
- aliases.insert(*parameter);
-
- // Check if we have seen the same argument
- let seen_parameters =
- arg_set.entry(argument).or_insert_with(VecSet::empty);
- // Add the current parameter to the parameters we have seen for this argument.
- // The previous parameters and the current one alias one another.
- seen_parameters.insert(*parameter);
- }
+ if let Some(aliases) = references.aliases.get_mut(&argument) {
+ // The argument reference is possibly aliased by this block parameter
+ aliases.insert(*parameter);
+
+ // Check if we have seen the same argument
+ let seen_parameters =
+ arg_set.entry(argument).or_insert_with(VecSet::empty);
+ // Add the current parameter to the parameters we have seen for this argument.
+ // The previous parameters and the current one alias one another.
+ seen_parameters.insert(*parameter);
}
}
}
diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs
index 443b21cfd1..f845ef3454 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/mem2reg/alias_set.rs
@@ -89,11 +89,4 @@ impl AliasSet {
}
}
}
-
- /// Return the first ValueId in the alias set as long as there is at least one.
- /// The ordering is arbitrary (by lowest ValueId) so this method should only be
- /// used when you need an arbitrary ValueId from the alias set.
- pub(super) fn first(&self) -> Option<ValueId> {
- self.aliases.as_ref().and_then(|aliases| aliases.iter().next().copied())
- }
}
diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg/block.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg/block.rs
index 25a0b230f8..011568937b 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/mem2reg/block.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/mem2reg/block.rs
@@ -16,16 +16,13 @@ use super::alias_set::AliasSet;
/// 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: im::OrdMap<ValueId, Expression>,
+ pub(super) dereferences: im::OrdSet<ValueId>,
- /// Each expression is tracked as to how many aliases it
+ /// Each value 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: im::OrdMap<Expression, AliasSet>,
+ pub(super) aliases: im::OrdMap<ValueId, AliasSet>,
/// Each allocate instruction result (and some reference block parameters)
/// will map to a Reference value which tracks whether the last value stored
@@ -39,16 +36,6 @@ pub(super) struct Block {
pub(super) last_loads: im::OrdMap<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, Hash)]
-pub(super) 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)]
pub(super) enum ReferenceValue {
@@ -65,14 +52,12 @@ impl ReferenceValue {
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<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 let Some(alias) = aliases.single_alias() {
- if let Some(ReferenceValue::Known(value)) = self.references.get(&alias) {
- return Some(*value);
- }
+ if let Some(aliases) = self.aliases.get(&address) {
+ // 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);
}
}
}
@@ -89,8 +74,7 @@ impl Block {
}
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();
+ let aliases = self.aliases.entry(address).or_default();
if aliases.is_unknown() {
// uh-oh, we don't know at all what this reference refers to, could be anything.
@@ -115,14 +99,6 @@ impl Block {
}
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 {
// If nothing would change, then don't call `.entry(...).and_modify(...)` as it involves creating more `Arc`s.
if let Some(aliases) = self.aliases.get(expression) {
@@ -130,10 +106,9 @@ impl Block {
continue;
}
}
- let expression = expression.clone();
self.aliases
- .entry(expression)
+ .entry(*expression)
.and_modify(|aliases| aliases.unify(new_aliases))
.or_insert_with(|| new_aliases.clone());
}
@@ -158,16 +133,11 @@ impl Block {
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.
- }
+ if function.dfg.value_is_reference(result) && self.get_known_value(address).is_none() {
+ self.dereferences.insert(result);
+ // 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.
}
}
@@ -177,12 +147,10 @@ impl Block {
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);
- });
- }
+ if let Some(aliases) = self.aliases.get(&address).cloned() {
+ aliases.for_each(|alias| {
+ f(self, alias);
+ });
}
}
@@ -231,10 +199,8 @@ impl Block {
}
pub(super) fn get_aliases_for_value(&self, value: ValueId) -> Cow<AliasSet> {
- if let Some(expression) = self.expressions.get(&value) {
- if let Some(aliases) = self.aliases.get(expression) {
- return Cow::Borrowed(aliases);
- }
+ if let Some(aliases) = self.aliases.get(&value) {
+ return Cow::Borrowed(aliases);
}
Cow::Owned(AliasSet::unknown())