diff --git a/compiler/noirc_evaluator/src/ssa/ir/dom.rs b/compiler/noirc_evaluator/src/ssa/ir/dom.rs index c1a7f14e0d..504eecf480 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/dom.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/dom.rs @@ -119,6 +119,29 @@ impl DominatorTree { } } + /// Walk up the dominator tree until we find a block for which `f` returns `Some` value. + /// Otherwise return `None` when we reach the top. + /// + /// Similar to `Iterator::filter_map` but only returns the first hit. + pub(crate) fn find_map_dominator( + &self, + mut block_id: BasicBlockId, + f: impl Fn(BasicBlockId) -> Option, + ) -> Option { + if !self.is_reachable(block_id) { + return None; + } + loop { + if let Some(value) = f(block_id) { + return Some(value); + } + block_id = match self.immediate_dominator(block_id) { + Some(immediate_dominator) => immediate_dominator, + None => return None, + } + } + } + /// Allocate and compute a dominator tree from a pre-computed control flow graph and /// post-order counterpart. pub(crate) fn with_cfg_and_post_order(cfg: &ControlFlowGraph, post_order: &PostOrder) -> Self { @@ -448,4 +471,22 @@ mod tests { assert!(dt.dominates(block2_id, block1_id)); assert!(dt.dominates(block2_id, block2_id)); } + + #[test] + fn test_find_map_dominator() { + let (dt, b0, b1, b2, _b3) = unreachable_node_setup(); + + assert_eq!( + dt.find_map_dominator(b2, |b| if b == b0 { Some("root") } else { None }), + Some("root") + ); + assert_eq!( + dt.find_map_dominator(b1, |b| if b == b0 { Some("unreachable") } else { None }), + None + ); + assert_eq!( + dt.find_map_dominator(b1, |b| if b == b1 { Some("not part of tree") } else { None }), + None + ); + } } diff --git a/compiler/noirc_evaluator/src/ssa/ir/instruction.rs b/compiler/noirc_evaluator/src/ssa/ir/instruction.rs index b48c755dbe..6737b335b7 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/instruction.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/instruction.rs @@ -362,6 +362,44 @@ impl Instruction { matches!(self.result_type(), InstructionResultType::Unknown) } + /// Indicates if the instruction has a side effect, ie. it can fail, or it interacts with memory. + /// + /// This is similar to `can_be_deduplicated`, but it doesn't depend on whether the caller takes + /// constraints into account, because it might not use it to isolate the side effects across branches. + pub(crate) fn has_side_effects(&self, dfg: &DataFlowGraph) -> bool { + use Instruction::*; + + match self { + // These either have side-effects or interact with memory + EnableSideEffectsIf { .. } + | Allocate + | Load { .. } + | Store { .. } + | IncrementRc { .. } + | DecrementRc { .. } => true, + + Call { func, .. } => match dfg[*func] { + Value::Intrinsic(intrinsic) => intrinsic.has_side_effects(), + _ => true, // Be conservative and assume other functions can have side effects. + }, + + // These can fail. + Constrain(..) | RangeCheck { .. } => true, + + // This should never be side-effectful + MakeArray { .. } => false, + + // These can have different behavior depending on the EnableSideEffectsIf context. + Binary(_) + | Cast(_, _) + | Not(_) + | Truncate { .. } + | IfElse { .. } + | ArrayGet { .. } + | ArraySet { .. } => self.requires_acir_gen_predicate(dfg), + } + } + /// Indicates if the instruction can be safely replaced with the results of another instruction with the same inputs. /// If `deduplicate_with_predicate` is set, we assume we're deduplicating with the instruction /// and its predicate, rather than just the instruction. Setting this means instructions that diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index 9668380404..7236dba76d 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -149,7 +149,8 @@ impl Function { use_constraint_info: bool, brillig_info: Option, ) { - let mut context = Context::new(self, use_constraint_info, brillig_info); + let mut context = Context::new(use_constraint_info, brillig_info); + let mut dom = DominatorTree::with_function(self); context.block_queue.push_back(self.entry_block()); while let Some(block) = context.block_queue.pop_front() { @@ -158,7 +159,7 @@ impl Function { } context.visited_blocks.insert(block); - context.fold_constants_in_block(self, block); + context.fold_constants_in_block(&mut self.dfg, &mut dom, block); } } } @@ -177,12 +178,10 @@ struct Context<'a> { /// We partition the maps of constrained values according to the side-effects flag at the point /// at which the values are constrained. This prevents constraints which are only sometimes enforced /// being used to modify the rest of the program. - constraint_simplification_mappings: HashMap>, + constraint_simplification_mappings: ConstraintSimplificationCache, // Cache of instructions without any side-effects along with their outputs. cached_instruction_results: InstructionResultCache, - - dom: DominatorTree, } #[derive(Copy, Clone)] @@ -191,6 +190,50 @@ pub(crate) struct BrilligInfo<'a> { brillig_functions: &'a BTreeMap, } +/// Records a simplified equivalents of an [`Instruction`] in the blocks +/// where the constraint that advised the simplification has been encountered. +/// +/// For more information see [`ConstraintSimplificationCache`]. +#[derive(Default)] +struct SimplificationCache { + /// Simplified expressions where we found them. + /// + /// It will always have at least one value because `add` is called + /// after the default is constructed. + simplifications: HashMap, +} + +impl SimplificationCache { + /// Called with a newly encountered simplification. + fn add(&mut self, dfg: &DataFlowGraph, simple: ValueId, block: BasicBlockId) { + self.simplifications + .entry(block) + .and_modify(|existing| { + // `SimplificationCache` may already hold a simplification in this block + // so we check whether `simple` is a better simplification than the current one. + if let Some((_, simpler)) = simplify(dfg, *existing, simple) { + *existing = simpler; + }; + }) + .or_insert(simple); + } + + /// Try to find a simplification in a visible block. + fn get(&self, block: BasicBlockId, dom: &DominatorTree) -> Option { + // Deterministically walk up the dominator chain until we encounter a block that contains a simplification. + dom.find_map_dominator(block, |b| self.simplifications.get(&b).cloned()) + } +} + +/// HashMap from `(side_effects_enabled_var, Instruction)` to a simplified expression that it can +/// be replaced with based on constraints that testify to their equivalence, stored together +/// with the set of blocks at which this constraint has been observed. +/// +/// Only blocks dominated by one in the cache should have access to this information, otherwise +/// we create a sort of time paradox where we replace an instruction with a constant we believe +/// it _should_ equal to, without ever actually producing and asserting the value. +type ConstraintSimplificationCache = HashMap>; + /// HashMap from `(Instruction, side_effects_enabled_var)` to the results of the instruction. /// Stored as a two-level map to avoid cloning Instructions during the `.get` call. /// @@ -210,11 +253,7 @@ struct ResultCache { } impl<'brillig> Context<'brillig> { - fn new( - function: &Function, - use_constraint_info: bool, - brillig_info: Option>, - ) -> Self { + fn new(use_constraint_info: bool, brillig_info: Option>) -> Self { Self { use_constraint_info, brillig_info, @@ -222,49 +261,57 @@ impl<'brillig> Context<'brillig> { block_queue: Default::default(), constraint_simplification_mappings: Default::default(), cached_instruction_results: Default::default(), - dom: DominatorTree::with_function(function), } } - fn fold_constants_in_block(&mut self, function: &mut Function, block: BasicBlockId) { - let instructions = function.dfg[block].take_instructions(); + fn fold_constants_in_block( + &mut self, + dfg: &mut DataFlowGraph, + dom: &mut DominatorTree, + block: BasicBlockId, + ) { + let instructions = dfg[block].take_instructions(); // Default side effect condition variable with an enabled state. - let mut side_effects_enabled_var = - function.dfg.make_constant(FieldElement::one(), Type::bool()); + let mut side_effects_enabled_var = dfg.make_constant(FieldElement::one(), Type::bool()); for instruction_id in instructions { self.fold_constants_into_instruction( - &mut function.dfg, + dfg, + dom, block, instruction_id, &mut side_effects_enabled_var, ); } - self.block_queue.extend(function.dfg[block].successors()); + self.block_queue.extend(dfg[block].successors()); } fn fold_constants_into_instruction( &mut self, dfg: &mut DataFlowGraph, + dom: &mut DominatorTree, mut block: BasicBlockId, id: InstructionId, side_effects_enabled_var: &mut ValueId, ) { let constraint_simplification_mapping = self.get_constraint_map(*side_effects_enabled_var); - let instruction = Self::resolve_instruction(id, dfg, constraint_simplification_mapping); + + let instruction = + Self::resolve_instruction(id, block, dfg, dom, constraint_simplification_mapping); + let old_results = dfg.instruction_results(id).to_vec(); // If a copy of this instruction exists earlier in the block, then reuse the previous results. if let Some(cache_result) = - self.get_cached(dfg, &instruction, *side_effects_enabled_var, block) + self.get_cached(dfg, dom, &instruction, *side_effects_enabled_var, block) { match cache_result { CacheResult::Cached(cached) => { Self::replace_result_ids(dfg, &old_results, cached); return; } - CacheResult::NeedToHoistToCommonBlock(dominator, _cached) => { + CacheResult::NeedToHoistToCommonBlock(dominator) => { // Just change the block to insert in the common dominator instead. // This will only move the current instance of the instruction right now. // When constant folding is run a second time later on, it'll catch @@ -314,8 +361,10 @@ impl<'brillig> Context<'brillig> { /// Fetches an [`Instruction`] by its [`InstructionId`] and fully resolves its inputs. fn resolve_instruction( instruction_id: InstructionId, + block: BasicBlockId, dfg: &DataFlowGraph, - constraint_simplification_mapping: &HashMap, + dom: &mut DominatorTree, + constraint_simplification_mapping: &HashMap, ) -> Instruction { let instruction = dfg[instruction_id].clone(); @@ -325,20 +374,29 @@ impl<'brillig> Context<'brillig> { // This allows us to reach a stable final `ValueId` for each instruction input as we add more // constraints to the cache. fn resolve_cache( + block: BasicBlockId, dfg: &DataFlowGraph, - cache: &HashMap, + dom: &mut DominatorTree, + cache: &HashMap, value_id: ValueId, ) -> ValueId { let resolved_id = dfg.resolve(value_id); match cache.get(&resolved_id) { - Some(cached_value) => resolve_cache(dfg, cache, *cached_value), + Some(simplification_cache) => { + if let Some(simplified) = simplification_cache.get(block, dom) { + resolve_cache(block, dfg, dom, cache, simplified) + } else { + resolved_id + } + } None => resolved_id, } } // Resolve any inputs to ensure that we're comparing like-for-like instructions. - instruction - .map_values(|value_id| resolve_cache(dfg, constraint_simplification_mapping, value_id)) + instruction.map_values(|value_id| { + resolve_cache(block, dfg, dom, constraint_simplification_mapping, value_id) + }) } /// Pushes a new [`Instruction`] into the [`DataFlowGraph`] which applies any optimizations @@ -384,26 +442,11 @@ impl<'brillig> Context<'brillig> { // to map from the more complex to the simpler value. if let Instruction::Constrain(lhs, rhs, _) = instruction { // These `ValueId`s should be fully resolved now. - match (&dfg[lhs], &dfg[rhs]) { - // Ignore trivial constraints - (Value::NumericConstant { .. }, Value::NumericConstant { .. }) => (), - - // Prefer replacing with constants where possible. - (Value::NumericConstant { .. }, _) => { - self.get_constraint_map(side_effects_enabled_var).insert(rhs, lhs); - } - (_, Value::NumericConstant { .. }) => { - self.get_constraint_map(side_effects_enabled_var).insert(lhs, rhs); - } - // Otherwise prefer block parameters over instruction results. - // This is as block parameters are more likely to be a single witness rather than a full expression. - (Value::Param { .. }, Value::Instruction { .. }) => { - self.get_constraint_map(side_effects_enabled_var).insert(rhs, lhs); - } - (Value::Instruction { .. }, Value::Param { .. }) => { - self.get_constraint_map(side_effects_enabled_var).insert(lhs, rhs); - } - (_, _) => (), + if let Some((complex, simple)) = simplify(dfg, lhs, rhs) { + self.get_constraint_map(side_effects_enabled_var) + .entry(complex) + .or_default() + .add(dfg, simple, block); } } } @@ -430,7 +473,7 @@ impl<'brillig> Context<'brillig> { fn get_constraint_map( &mut self, side_effects_enabled_var: ValueId, - ) -> &mut HashMap { + ) -> &mut HashMap { self.constraint_simplification_mappings.entry(side_effects_enabled_var).or_default() } @@ -445,9 +488,11 @@ impl<'brillig> Context<'brillig> { } } + /// Get a cached result if it can be used in this context. fn get_cached( - &mut self, + &self, dfg: &DataFlowGraph, + dom: &mut DominatorTree, instruction: &Instruction, side_effects_enabled_var: ValueId, block: BasicBlockId, @@ -457,7 +502,7 @@ impl<'brillig> Context<'brillig> { let predicate = self.use_constraint_info && instruction.requires_acir_gen_predicate(dfg); let predicate = predicate.then_some(side_effects_enabled_var); - results_for_instruction.get(&predicate)?.get(block, &mut self.dom) + results_for_instruction.get(&predicate)?.get(block, dom, instruction.has_side_effects(dfg)) } /// Checks if the given instruction is a call to a brillig function with all constant arguments. @@ -635,14 +680,21 @@ impl ResultCache { /// We require that the cached instruction's block dominates `block` in order to avoid /// cycles causing issues (e.g. two instructions being replaced with the results of each other /// such that neither instruction exists anymore.) - fn get(&self, block: BasicBlockId, dom: &mut DominatorTree) -> Option { - self.result.as_ref().map(|(origin_block, results)| { + fn get( + &self, + block: BasicBlockId, + dom: &mut DominatorTree, + has_side_effects: bool, + ) -> Option { + self.result.as_ref().and_then(|(origin_block, results)| { if dom.dominates(*origin_block, block) { - CacheResult::Cached(results) - } else { + Some(CacheResult::Cached(results)) + } else if !has_side_effects { // Insert a copy of this instruction in the common dominator let dominator = dom.common_dominator(*origin_block, block); - CacheResult::NeedToHoistToCommonBlock(dominator, results) + Some(CacheResult::NeedToHoistToCommonBlock(dominator)) + } else { + None } }) } @@ -650,7 +702,7 @@ impl ResultCache { enum CacheResult<'a> { Cached(&'a [ValueId]), - NeedToHoistToCommonBlock(BasicBlockId, &'a [ValueId]), + NeedToHoistToCommonBlock(BasicBlockId), } /// Result of trying to evaluate an instruction (any instruction) in this pass. @@ -697,6 +749,25 @@ fn value_id_to_calldata(value_id: ValueId, dfg: &DataFlowGraph, calldata: &mut V panic!("Expected ValueId to be numeric constant or array constant"); } +/// Check if one expression is simpler than the other. +/// Returns `Some((complex, simple))` if a simplification was found, otherwise `None`. +/// Expects the `ValueId`s to be fully resolved. +fn simplify(dfg: &DataFlowGraph, lhs: ValueId, rhs: ValueId) -> Option<(ValueId, ValueId)> { + match (&dfg[lhs], &dfg[rhs]) { + // Ignore trivial constraints + (Value::NumericConstant { .. }, Value::NumericConstant { .. }) => None, + + // Prefer replacing with constants where possible. + (Value::NumericConstant { .. }, _) => Some((rhs, lhs)), + (_, Value::NumericConstant { .. }) => Some((lhs, rhs)), + // Otherwise prefer block parameters over instruction results. + // This is as block parameters are more likely to be a single witness rather than a full expression. + (Value::Param { .. }, Value::Instruction { .. }) => Some((rhs, lhs)), + (Value::Instruction { .. }, Value::Param { .. }) => Some((lhs, rhs)), + (_, _) => None, + } +} + #[cfg(test)] mod test { use std::sync::Arc; @@ -1352,6 +1423,52 @@ mod test { assert_normalized_ssa_equals(ssa, expected); } + #[test] + fn does_not_use_cached_constrain_in_block_that_is_not_dominated() { + let src = " + brillig(inline) fn main f0 { + b0(v0: Field, v1: Field): + v3 = eq v0, Field 0 + jmpif v3 then: b1, else: b2 + b1(): + v5 = eq v1, Field 1 + constrain v1 == Field 1 + jmp b2() + b2(): + v6 = eq v1, Field 0 + constrain v1 == Field 0 + return + } + "; + let ssa = Ssa::from_str(src).unwrap(); + let ssa = ssa.fold_constants_using_constraints(); + assert_normalized_ssa_equals(ssa, src); + } + + #[test] + fn does_not_hoist_constrain_to_common_ancestor() { + let src = " + brillig(inline) fn main f0 { + b0(v0: Field, v1: Field): + v3 = eq v0, Field 0 + jmpif v3 then: b1, else: b2 + b1(): + constrain v1 == Field 1 + jmp b2() + b2(): + jmpif v0 then: b3, else: b4 + b3(): + constrain v1 == Field 1 // This was incorrectly hoisted to b0 but this condition is not valid when going b0 -> b2 -> b4 + jmp b4() + b4(): + return + } + "; + let ssa = Ssa::from_str(src).unwrap(); + let ssa = ssa.fold_constants_using_constraints(); + assert_normalized_ssa_equals(ssa, src); + } + #[test] fn deduplicates_side_effecting_intrinsics() { let src = "