From 25b44cd1b239a72e362976d86ad4e1c9f5eb1bca Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Tue, 26 Nov 2024 12:08:43 -0300 Subject: [PATCH 1/5] Add a test that fails --- .../src/ssa/opt/constant_folding.rs | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index 019bace33a..78f1b06e82 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -1341,4 +1341,26 @@ mod test { let ssa = ssa.fold_constants_with_brillig(&brillig); 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); + } } From 2b05814e3726ecdf750f4aca6d1df642c4a2168b Mon Sep 17 00:00:00 2001 From: Akosh Farkash Date: Tue, 26 Nov 2024 15:51:54 +0000 Subject: [PATCH 2/5] Add simplify function --- .../src/ssa/opt/constant_folding.rs | 41 ++++++++++--------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index 78f1b06e82..a669cfb55a 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -377,26 +377,8 @@ 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).insert(complex, simple); } } } @@ -687,6 +669,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; From 8005a80f652d634427d523a2e770bf5f56c8ad07 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Tue, 26 Nov 2024 13:39:58 -0300 Subject: [PATCH 3/5] Only simplify a constrained value to values in dominant blocks --- .../src/ssa/opt/constant_folding.rs | 56 ++++++++++++++----- 1 file changed, 42 insertions(+), 14 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index a669cfb55a..b812854bb3 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -169,12 +169,17 @@ struct Context<'a> { /// Contains sets of values which are constrained to be equivalent to each other. /// - /// The mapping's structure is `side_effects_enabled_var => (constrained_value => simplified_value)`. + /// The mapping's structure is `side_effects_enabled_var => (constrained_value => [(block, simplified_value)])`. /// /// 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>, + /// + /// We also keep track of how a value was simplified to other values per block. That is, + /// a same ValueId could have been simplified to one value in one block and to another value + /// in another block. + constraint_simplification_mappings: + HashMap>>, // Cache of instructions without any side-effects along with their outputs. cached_instruction_results: InstructionResultCache, @@ -244,8 +249,15 @@ impl<'brillig> Context<'brillig> { 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 constraint_simplification_mapping = + self.constraint_simplification_mappings.get(side_effects_enabled_var); + let instruction = Self::resolve_instruction( + id, + block, + dfg, + &mut self.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. @@ -307,8 +319,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: Option<&HashMap>>, ) -> Instruction { let instruction = dfg[instruction_id].clone(); @@ -319,19 +333,30 @@ impl<'brillig> Context<'brillig> { // constraints to the cache. fn resolve_cache( dfg: &DataFlowGraph, - cache: &HashMap, + dom: &mut DominatorTree, + cache: Option<&HashMap>>, value_id: ValueId, + block: BasicBlockId, ) -> ValueId { let resolved_id = dfg.resolve(value_id); - match cache.get(&resolved_id) { - Some(cached_value) => resolve_cache(dfg, cache, *cached_value), - None => resolved_id, + let Some(cached_values) = cache.and_then(|cache| cache.get(&resolved_id)) else { + return resolved_id; + }; + + for (cached_block, cached_value) in cached_values { + // We can only use the simplified value if it was simplified in a block that dominates the current one + if dom.dominates(*cached_block, block) { + return resolve_cache(dfg, dom, cache, *cached_value, block); + } } + + 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(dfg, dom, constraint_simplification_mapping, value_id, block) + }) } /// Pushes a new [`Instruction`] into the [`DataFlowGraph`] which applies any optimizations @@ -378,7 +403,10 @@ impl<'brillig> Context<'brillig> { if let Instruction::Constrain(lhs, rhs, _) = instruction { // These `ValueId`s should be fully resolved now. if let Some((complex, simple)) = simplify(dfg, lhs, rhs) { - self.get_constraint_map(side_effects_enabled_var).insert(complex, simple); + self.get_constraint_map(side_effects_enabled_var) + .entry(complex) + .or_default() + .push((block, simple)); } } } @@ -402,7 +430,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() } @@ -1356,7 +1384,7 @@ mod test { jmp b2() b2(): v6 = eq v1, Field 0 - constrain v1 == Field 0 + constrain v1 == Field 0 // This `v1` here was incorrectly replaced with `Field 1` return } "; From c7ff48b8d71d2addf7b1a3a0246a62eeafb697d6 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Tue, 26 Nov 2024 14:27:35 -0300 Subject: [PATCH 4/5] Add a regression test for instruction hoisting --- .../src/ssa/opt/constant_folding.rs | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index b812854bb3..892829c1ed 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -1392,4 +1392,28 @@ mod test { 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); + } } From 3306f490bf100899cf46f366be2afe3732054f35 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Tue, 26 Nov 2024 15:00:27 -0300 Subject: [PATCH 5/5] Don't hoist instructions that have side effects --- .../src/ssa/opt/constant_folding.rs | 26 +++++++++++++++---- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index 892829c1ed..fac6e52104 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -270,11 +270,13 @@ impl<'brillig> Context<'brillig> { return; } CacheResult::NeedToHoistToCommonBlock(dominator, _cached) => { - // 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 - // that the previous instance can be deduplicated to this instance. - block = dominator; + if instruction_can_be_hoisted(&instruction, dfg, self.use_constraint_info) { + // 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 + // that the previous instance can be deduplicated to this instance. + block = dominator; + } } } } @@ -621,6 +623,20 @@ impl<'brillig> Context<'brillig> { } } +fn instruction_can_be_hoisted( + instruction: &Instruction, + dfg: &mut DataFlowGraph, + deduplicate_with_predicate: bool, +) -> bool { + // These two can never be hoisted as they have a side-effect + // (though it's fine to de-duplicate them, just not fine to hoist them) + if matches!(instruction, Instruction::Constrain(..) | Instruction::RangeCheck { .. }) { + return false; + } + + instruction.can_be_deduplicated(dfg, deduplicate_with_predicate) +} + impl ResultCache { /// Records that an `Instruction` in block `block` produced the result values `results`. fn cache(&mut self, block: BasicBlockId, results: Vec) {