diff --git a/compiler/noirc_evaluator/src/ssa.rs b/compiler/noirc_evaluator/src/ssa.rs index 694351006c5..2ddee1765b0 100644 --- a/compiler/noirc_evaluator/src/ssa.rs +++ b/compiler/noirc_evaluator/src/ssa.rs @@ -63,6 +63,10 @@ pub(crate) fn optimize_into_acir( // Run mem2reg once more with the flattened CFG to catch any remaining loads/stores .run_pass(Ssa::mem2reg, "After Mem2Reg:") .run_pass(Ssa::fold_constants, "After Constant Folding:") + .run_pass( + Ssa::fold_constants_using_constraints, + "After Constant Folding With Constraint Info:", + ) .run_pass(Ssa::dead_instruction_elimination, "After Dead Instruction Elimination:") .finish(); diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index addaee3ba8d..06ae4bf5202 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -22,6 +22,7 @@ //! different blocks are merged, i.e. after the [`flatten_cfg`][super::flatten_cfg] pass. use std::collections::HashSet; +use acvm::FieldElement; use iter_extended::vecmap; use crate::ssa::{ @@ -30,7 +31,8 @@ use crate::ssa::{ dfg::{DataFlowGraph, InsertInstructionResult}, function::Function, instruction::{Instruction, InstructionId}, - value::ValueId, + types::Type, + value::{Value, ValueId}, }, ssa_gen::Ssa, }; @@ -43,7 +45,20 @@ impl Ssa { #[tracing::instrument(level = "trace", skip(self))] pub(crate) fn fold_constants(mut self) -> Ssa { for function in self.functions.values_mut() { - constant_fold(function); + constant_fold(function, false); + } + self + } + + /// Performs constant folding on each instruction. + /// + /// Also uses constraint information to inform more optimizations. + /// + /// See [`constant_folding`][self] module for more information. + #[tracing::instrument(level = "trace", skip(self))] + pub(crate) fn fold_constants_using_constraints(mut self) -> Ssa { + for function in self.functions.values_mut() { + constant_fold(function, true); } self } @@ -51,8 +66,8 @@ impl Ssa { /// The structure of this pass is simple: /// Go through each block and re-insert all instructions. -fn constant_fold(function: &mut Function) { - let mut context = Context::default(); +fn constant_fold(function: &mut Function, use_constraint_info: bool) { + let mut context = Context { use_constraint_info, ..Default::default() }; context.block_queue.push(function.entry_block()); while let Some(block) = context.block_queue.pop() { @@ -67,6 +82,7 @@ fn constant_fold(function: &mut Function) { #[derive(Default)] struct Context { + use_constraint_info: bool, /// Maps pre-folded ValueIds to the new ValueIds obtained by re-inserting the instruction. visited_blocks: HashSet, block_queue: Vec, @@ -79,24 +95,43 @@ impl Context { // Cache of instructions without any side-effects along with their outputs. let mut cached_instruction_results: HashMap> = HashMap::default(); + // 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)`. + // + // 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. + let mut constraint_simplification_mappings: HashMap> = + HashMap::default(); + let mut side_effects_enabled_var = + function.dfg.make_constant(FieldElement::one(), Type::bool()); + for instruction_id in instructions { - Self::fold_constants_into_instruction( + self.fold_constants_into_instruction( &mut function.dfg, block, instruction_id, &mut cached_instruction_results, + &mut constraint_simplification_mappings, + &mut side_effects_enabled_var, ); } self.block_queue.extend(function.dfg[block].successors()); } fn fold_constants_into_instruction( + &self, dfg: &mut DataFlowGraph, block: BasicBlockId, id: InstructionId, instruction_result_cache: &mut HashMap>, + constraint_simplification_mappings: &mut HashMap>, + side_effects_enabled_var: &mut ValueId, ) { - let instruction = Self::resolve_instruction(id, dfg); + let constraint_simplification_mapping = + constraint_simplification_mappings.entry(*side_effects_enabled_var).or_default(); + let instruction = Self::resolve_instruction(id, dfg, 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. @@ -110,15 +145,49 @@ impl Context { Self::replace_result_ids(dfg, &old_results, &new_results); - Self::cache_instruction(instruction, new_results, dfg, instruction_result_cache); + self.cache_instruction( + instruction.clone(), + new_results, + dfg, + instruction_result_cache, + constraint_simplification_mapping, + ); + + // If we just inserted an `Instruction::EnableSideEffects`, we need to update `side_effects_enabled_var` + // so that we use the correct set of constrained values in future. + if let Instruction::EnableSideEffects { condition } = instruction { + *side_effects_enabled_var = condition; + }; } /// Fetches an [`Instruction`] by its [`InstructionId`] and fully resolves its inputs. - fn resolve_instruction(instruction_id: InstructionId, dfg: &DataFlowGraph) -> Instruction { + fn resolve_instruction( + instruction_id: InstructionId, + dfg: &DataFlowGraph, + constraint_simplification_mapping: &HashMap, + ) -> Instruction { let instruction = dfg[instruction_id].clone(); + // Alternate between resolving `value_id` in the `dfg` and checking to see if the resolved value + // has been constrained to be equal to some simpler value in the current block. + // + // This allows us to reach a stable final `ValueId` for each instruction input as we add more + // constraints to the cache. + fn resolve_cache( + dfg: &DataFlowGraph, + 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), + None => resolved_id, + } + } + // Resolve any inputs to ensure that we're comparing like-for-like instructions. - instruction.map_values(|value_id| dfg.resolve(value_id)) + instruction + .map_values(|value_id| resolve_cache(dfg, constraint_simplification_mapping, value_id)) } /// Pushes a new [`Instruction`] into the [`DataFlowGraph`] which applies any optimizations @@ -152,11 +221,42 @@ impl Context { } fn cache_instruction( + &self, instruction: Instruction, instruction_results: Vec, dfg: &DataFlowGraph, instruction_result_cache: &mut HashMap>, + constraint_simplification_mapping: &mut HashMap, ) { + if self.use_constraint_info { + // If the instruction was a constraint, then create a link between the two `ValueId`s + // 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 { .. }, _) => { + constraint_simplification_mapping.insert(rhs, lhs); + } + (_, Value::NumericConstant { .. }) => { + constraint_simplification_mapping.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 { .. }) => { + constraint_simplification_mapping.insert(rhs, lhs); + } + (Value::Instruction { .. }, Value::Param { .. }) => { + constraint_simplification_mapping.insert(lhs, rhs); + } + (_, _) => (), + } + } + } + // If the instruction doesn't have side-effects, cache the results so we can reuse them if // the same instruction appears again later in the block. if instruction.is_pure(dfg) { @@ -336,9 +436,9 @@ mod test { // // fn main f0 { // b0(v0: u16, Field 255: Field): - // v5 = div v0, Field 255 - // v6 = truncate v5 to 8 bits, max_bit_size: 16 - // return v6 + // v6 = div v0, Field 255 + // v7 = truncate v6 to 8 bits, max_bit_size: 16 + // return v7 // } main.dfg.set_value_from_id(v1, constant); @@ -354,7 +454,7 @@ mod test { ); assert_eq!( &main.dfg[instructions[1]], - &Instruction::Truncate { value: ValueId::test_new(5), bit_size: 8, max_bit_size: 16 } + &Instruction::Truncate { value: ValueId::test_new(6), bit_size: 8, max_bit_size: 16 } ); }