From fac2a3ee1af88d9488bf7273f1c0075e7bc314ba Mon Sep 17 00:00:00 2001 From: Tom French Date: Tue, 16 Jan 2024 23:40:35 +0000 Subject: [PATCH 1/4] feat: replace constrained values with simpler equivalents --- .../src/ssa/opt/constant_folding.rs | 93 +++++++++++++++++-- 1 file changed, 84 insertions(+), 9 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index addaee3ba8..5f1bf9c470 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, }; @@ -78,6 +80,10 @@ impl Context { // Cache of instructions without any side-effects along with their outputs. let mut cached_instruction_results: HashMap> = HashMap::default(); + let mut constrained_values: 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( @@ -85,6 +91,8 @@ impl Context { block, instruction_id, &mut cached_instruction_results, + &mut constrained_values, + &mut side_effects_enabled_var, ); } self.block_queue.extend(function.dfg[block].successors()); @@ -95,8 +103,14 @@ impl Context { block: BasicBlockId, id: InstructionId, instruction_result_cache: &mut HashMap>, + constrained_values: &mut HashMap>, + side_effects_enabled_var: &mut ValueId, ) { - let instruction = Self::resolve_instruction(id, dfg); + let instruction = Self::resolve_instruction( + id, + dfg, + constrained_values.entry(*side_effects_enabled_var).or_default(), + ); 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 +124,48 @@ 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, + constrained_values.entry(*side_effects_enabled_var).or_default(), + ); + + // 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_cache: &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_cache, value_id)) } /// Pushes a new [`Instruction`] into the [`DataFlowGraph`] which applies any optimizations @@ -156,7 +203,35 @@ impl Context { instruction_results: Vec, dfg: &DataFlowGraph, instruction_result_cache: &mut HashMap>, + constraint_cache: &mut HashMap, ) { + // 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_cache.insert(rhs, lhs); + } + (_, Value::NumericConstant { .. }) => { + constraint_cache.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_cache.insert(rhs, lhs); + } + (Value::Instruction { .. }, Value::Param { .. }) => { + constraint_cache.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 +411,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 +429,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 } ); } From c81f9d308aa4d881ef3bc881d0415cbad4308e5c Mon Sep 17 00:00:00 2001 From: Tom French Date: Wed, 7 Feb 2024 11:01:08 +0000 Subject: [PATCH 2/4] feat: add a constant folding pass which is unaware of constraints --- compiler/noirc_evaluator/src/ssa.rs | 4 + .../src/ssa/opt/constant_folding.rs | 74 ++++++++++++------- 2 files changed, 50 insertions(+), 28 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa.rs b/compiler/noirc_evaluator/src/ssa.rs index 694351006c..2ddee1765b 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 5f1bf9c470..4bf49741a6 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -45,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 } @@ -53,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() { @@ -69,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, @@ -86,7 +100,7 @@ impl Context { 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, @@ -99,6 +113,7 @@ impl Context { } fn fold_constants_into_instruction( + &self, dfg: &mut DataFlowGraph, block: BasicBlockId, id: InstructionId, @@ -124,7 +139,7 @@ impl Context { Self::replace_result_ids(dfg, &old_results, &new_results); - Self::cache_instruction( + self.cache_instruction( instruction.clone(), new_results, dfg, @@ -199,36 +214,39 @@ impl Context { } fn cache_instruction( + &self, instruction: Instruction, instruction_results: Vec, dfg: &DataFlowGraph, instruction_result_cache: &mut HashMap>, constraint_cache: &mut HashMap, ) { - // 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_cache.insert(rhs, lhs); - } - (_, Value::NumericConstant { .. }) => { - constraint_cache.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_cache.insert(rhs, lhs); - } - (Value::Instruction { .. }, Value::Param { .. }) => { - constraint_cache.insert(lhs, rhs); + 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_cache.insert(rhs, lhs); + } + (_, Value::NumericConstant { .. }) => { + constraint_cache.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_cache.insert(rhs, lhs); + } + (Value::Instruction { .. }, Value::Param { .. }) => { + constraint_cache.insert(lhs, rhs); + } + (_, _) => (), } - (_, _) => (), } } From a7796932fca8ec77a8c63bfb9576af8569b2ffa5 Mon Sep 17 00:00:00 2001 From: TomAFrench Date: Wed, 7 Feb 2024 17:05:49 +0000 Subject: [PATCH 3/4] chore: documentation --- .../src/ssa/opt/constant_folding.rs | 33 ++++++++++++------- 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index 4bf49741a6..8849582d5f 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -94,7 +94,15 @@ impl Context { // Cache of instructions without any side-effects along with their outputs. let mut cached_instruction_results: HashMap> = HashMap::default(); - let mut constrained_values: HashMap> = + + // 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()); @@ -105,7 +113,7 @@ impl Context { block, instruction_id, &mut cached_instruction_results, - &mut constrained_values, + &mut constraint_simplification_mappings, &mut side_effects_enabled_var, ); } @@ -118,13 +126,14 @@ impl Context { block: BasicBlockId, id: InstructionId, instruction_result_cache: &mut HashMap>, - constrained_values: &mut HashMap>, + constraint_simplification_mappings: &mut HashMap>, side_effects_enabled_var: &mut ValueId, ) { + let constraint_simplification_mapping = constraint_simplification_mappings.entry(*side_effects_enabled_var).or_default(); let instruction = Self::resolve_instruction( id, dfg, - constrained_values.entry(*side_effects_enabled_var).or_default(), + constraint_simplification_mapping ); let old_results = dfg.instruction_results(id).to_vec(); @@ -144,7 +153,7 @@ impl Context { new_results, dfg, instruction_result_cache, - constrained_values.entry(*side_effects_enabled_var).or_default(), + constraint_simplification_mapping ); // If we just inserted an `Instruction::EnableSideEffects`, we need to update `side_effects_enabled_var` @@ -158,7 +167,7 @@ impl Context { fn resolve_instruction( instruction_id: InstructionId, dfg: &DataFlowGraph, - constraint_cache: &HashMap, + constraint_simplification_mapping: &HashMap, ) -> Instruction { let instruction = dfg[instruction_id].clone(); @@ -180,7 +189,7 @@ impl Context { } // Resolve any inputs to ensure that we're comparing like-for-like instructions. - instruction.map_values(|value_id| resolve_cache(dfg, constraint_cache, 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 @@ -219,7 +228,7 @@ impl Context { instruction_results: Vec, dfg: &DataFlowGraph, instruction_result_cache: &mut HashMap>, - constraint_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 @@ -232,18 +241,18 @@ impl Context { // Prefer replacing with constants where possible. (Value::NumericConstant { .. }, _) => { - constraint_cache.insert(rhs, lhs); + constraint_simplification_mapping.insert(rhs, lhs); } (_, Value::NumericConstant { .. }) => { - constraint_cache.insert(lhs, rhs); + 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_cache.insert(rhs, lhs); + constraint_simplification_mapping.insert(rhs, lhs); } (Value::Instruction { .. }, Value::Param { .. }) => { - constraint_cache.insert(lhs, rhs); + constraint_simplification_mapping.insert(lhs, rhs); } (_, _) => (), } From 9a2bf9b6976ec03beecc1ac5c3c3c8c558f224b9 Mon Sep 17 00:00:00 2001 From: TomAFrench Date: Wed, 7 Feb 2024 17:11:25 +0000 Subject: [PATCH 4/4] chore: cargo fmt --- .../src/ssa/opt/constant_folding.rs | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index 8849582d5f..06ae4bf520 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -94,9 +94,9 @@ 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 @@ -129,12 +129,9 @@ impl Context { constraint_simplification_mappings: &mut HashMap>, side_effects_enabled_var: &mut ValueId, ) { - 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 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. @@ -153,7 +150,7 @@ impl Context { new_results, dfg, instruction_result_cache, - constraint_simplification_mapping + constraint_simplification_mapping, ); // If we just inserted an `Instruction::EnableSideEffects`, we need to update `side_effects_enabled_var` @@ -189,7 +186,8 @@ impl Context { } // 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, constraint_simplification_mapping, value_id)) } /// Pushes a new [`Instruction`] into the [`DataFlowGraph`] which applies any optimizations