From d570b7e61fc9e96229b74b51cc95d768cf87a399 Mon Sep 17 00:00:00 2001 From: Jake Fecher Date: Mon, 11 Nov 2024 07:54:49 -0600 Subject: [PATCH 1/5] Initial test --- .../src/ssa/opt/constant_folding.rs | 55 +++++++++---------- 1 file changed, 25 insertions(+), 30 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index 3b86ded4a87..191c8fb75a7 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -87,6 +87,18 @@ struct Context { /// Maps pre-folded ValueIds to the new ValueIds obtained by re-inserting the instruction. visited_blocks: HashSet, block_queue: Vec, + + // 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. + constraint_simplification_mappings: HashMap>, + + // Cache of instructions without any side-effects along with their outputs. + cached_instruction_results: InstructionResultCache, } /// HashMap from (Instruction, side_effects_enabled_var) to the results of the instruction. @@ -97,18 +109,6 @@ impl Context { fn fold_constants_in_block(&mut self, function: &mut Function, block: BasicBlockId) { let instructions = function.dfg[block].take_instructions(); - // Cache of instructions without any side-effects along with their outputs. - let mut cached_instruction_results = 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()); @@ -117,8 +117,6 @@ impl Context { &mut function.dfg, block, instruction_id, - &mut cached_instruction_results, - &mut constraint_simplification_mappings, &mut side_effects_enabled_var, ); } @@ -126,22 +124,19 @@ impl Context { } fn fold_constants_into_instruction( - &self, + &mut self, dfg: &mut DataFlowGraph, block: BasicBlockId, id: InstructionId, - instruction_result_cache: &mut InstructionResultCache, - 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 constraint_simplification_mapping = self.get_constraint_map(*side_effects_enabled_var); 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. if let Some(cached_results) = - Self::get_cached(dfg, instruction_result_cache, &instruction, *side_effects_enabled_var) + Self::get_cached(dfg, &mut self.cached_instruction_results, &instruction, *side_effects_enabled_var) { Self::replace_result_ids(dfg, &old_results, cached_results); return; @@ -156,8 +151,6 @@ impl Context { instruction.clone(), new_results, dfg, - instruction_result_cache, - constraint_simplification_mapping, *side_effects_enabled_var, ); @@ -229,12 +222,10 @@ impl Context { } fn cache_instruction( - &self, + &mut self, instruction: Instruction, instruction_results: Vec, dfg: &DataFlowGraph, - instruction_result_cache: &mut InstructionResultCache, - constraint_simplification_mapping: &mut HashMap, side_effects_enabled_var: ValueId, ) { if self.use_constraint_info { @@ -248,18 +239,18 @@ impl Context { // Prefer replacing with constants where possible. (Value::NumericConstant { .. }, _) => { - constraint_simplification_mapping.insert(rhs, lhs); + self.get_constraint_map(side_effects_enabled_var).insert(rhs, lhs); } (_, Value::NumericConstant { .. }) => { - constraint_simplification_mapping.insert(lhs, rhs); + 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 { .. }) => { - constraint_simplification_mapping.insert(rhs, lhs); + self.get_constraint_map(side_effects_enabled_var).insert(rhs, lhs); } (Value::Instruction { .. }, Value::Param { .. }) => { - constraint_simplification_mapping.insert(lhs, rhs); + self.get_constraint_map(side_effects_enabled_var).insert(lhs, rhs); } (_, _) => (), } @@ -273,13 +264,17 @@ impl Context { self.use_constraint_info && instruction.requires_acir_gen_predicate(dfg); let predicate = use_predicate.then_some(side_effects_enabled_var); - instruction_result_cache + self.cached_instruction_results .entry(instruction) .or_default() .insert(predicate, instruction_results); } } + fn get_constraint_map(&mut self, side_effects_enabled_var: ValueId) -> &mut HashMap { + self.constraint_simplification_mappings.entry(side_effects_enabled_var).or_default() + } + /// Replaces a set of [`ValueId`]s inside the [`DataFlowGraph`] with another. fn replace_result_ids( dfg: &mut DataFlowGraph, From 288e42564a94d9e7046ae2b31cf3a825f18cf49e Mon Sep 17 00:00:00 2001 From: Jake Fecher Date: Mon, 11 Nov 2024 09:01:12 -0600 Subject: [PATCH 2/5] Deduplicate more --- .../src/ssa/opt/constant_folding.rs | 79 +++++++++++++++---- 1 file changed, 63 insertions(+), 16 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index 191c8fb75a7..64d4667c36d 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -28,6 +28,7 @@ use crate::ssa::{ ir::{ basic_block::BasicBlockId, dfg::{DataFlowGraph, InsertInstructionResult}, + dom::DominatorTree, function::Function, instruction::{Instruction, InstructionId}, types::Type, @@ -67,7 +68,7 @@ impl Function { /// The structure of this pass is simple: /// Go through each block and re-insert all instructions. pub(crate) fn constant_fold(&mut self, use_constraint_info: bool) { - let mut context = Context { use_constraint_info, ..Default::default() }; + let mut context = Context::new(self, use_constraint_info); context.block_queue.push(self.entry_block()); while let Some(block) = context.block_queue.pop() { @@ -81,7 +82,6 @@ impl Function { } } -#[derive(Default)] struct Context { use_constraint_info: bool, /// Maps pre-folded ValueIds to the new ValueIds obtained by re-inserting the instruction. @@ -99,13 +99,34 @@ struct Context { // Cache of instructions without any side-effects along with their outputs. cached_instruction_results: InstructionResultCache, + + dom: DominatorTree, } /// 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. -type InstructionResultCache = HashMap, Vec>>; +/// +/// In addition to each result, the original BasicBlockId is stored as well. This allows us +/// to deduplicate instructions across blocks as long as the new block dominates the original. +type InstructionResultCache = HashMap, ResultCache>>; + +#[derive(Default)] +struct ResultCache { + results: Vec<(BasicBlockId, Vec)>, +} impl Context { + fn new(function: &Function, use_constraint_info: bool) -> Self { + Self { + use_constraint_info, + visited_blocks: Default::default(), + 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(); @@ -135,9 +156,15 @@ impl Context { 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(cached_results) = - Self::get_cached(dfg, &mut self.cached_instruction_results, &instruction, *side_effects_enabled_var) - { + if let Some(cached_results) = Self::get_cached( + dfg, + &mut self.cached_instruction_results, + &instruction, + *side_effects_enabled_var, + self.use_constraint_info, + block, + &mut self.dom, + ) { Self::replace_result_ids(dfg, &old_results, cached_results); return; } @@ -152,6 +179,7 @@ impl Context { new_results, dfg, *side_effects_enabled_var, + block, ); // If we just inserted an `Instruction::EnableSideEffectsIf`, we need to update `side_effects_enabled_var` @@ -227,6 +255,7 @@ impl Context { instruction_results: Vec, dfg: &DataFlowGraph, side_effects_enabled_var: ValueId, + block: BasicBlockId, ) { if self.use_constraint_info { // If the instruction was a constraint, then create a link between the two `ValueId`s @@ -267,11 +296,16 @@ impl Context { self.cached_instruction_results .entry(instruction) .or_default() - .insert(predicate, instruction_results); + .entry(predicate) + .or_default() + .cache(block, instruction_results); } } - fn get_constraint_map(&mut self, side_effects_enabled_var: ValueId) -> &mut HashMap { + fn get_constraint_map( + &mut self, + side_effects_enabled_var: ValueId, + ) -> &mut HashMap { self.constraint_simplification_mappings.entry(side_effects_enabled_var).or_default() } @@ -291,18 +325,31 @@ impl Context { instruction_result_cache: &'a mut InstructionResultCache, instruction: &Instruction, side_effects_enabled_var: ValueId, + use_constraint_info: bool, + block: BasicBlockId, + dom_tree: &mut DominatorTree, ) -> Option<&'a Vec> { - let results_for_instruction = instruction_result_cache.get(instruction); + let results_for_instruction = instruction_result_cache.get(instruction)?; - // See if there's a cached version with no predicate first - if let Some(results) = results_for_instruction.and_then(|map| map.get(&None)) { - return Some(results); - } + let predicate = use_constraint_info && instruction.requires_acir_gen_predicate(dfg); + let predicate = predicate.then_some(side_effects_enabled_var); - let predicate = - instruction.requires_acir_gen_predicate(dfg).then_some(side_effects_enabled_var); + results_for_instruction.get(&predicate)?.get(block, dom_tree) + } +} + +impl ResultCache { + fn cache(&mut self, block: BasicBlockId, results: Vec) { + self.results.push((block, results)); + } - results_for_instruction.and_then(|map| map.get(&predicate)) + fn get(&self, block: BasicBlockId, dom: &mut DominatorTree) -> Option<&Vec> { + for (origin_block, results) in &self.results { + if dom.dominates(block, *origin_block) { + return Some(results); + } + } + None } } From e8b54c1c9b781d1e1f284ebb4bd3e97e48dcc86a Mon Sep 17 00:00:00 2001 From: Jake Fecher Date: Mon, 11 Nov 2024 09:06:56 -0600 Subject: [PATCH 3/5] Use self --- .../src/ssa/opt/constant_folding.rs | 22 ++++++------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs index 64d4667c36d..d7c81e26252 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -156,15 +156,9 @@ impl Context { 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(cached_results) = Self::get_cached( - dfg, - &mut self.cached_instruction_results, - &instruction, - *side_effects_enabled_var, - self.use_constraint_info, - block, - &mut self.dom, - ) { + if let Some(cached_results) = + self.get_cached(dfg, &instruction, *side_effects_enabled_var, block) + { Self::replace_result_ids(dfg, &old_results, cached_results); return; } @@ -321,20 +315,18 @@ impl Context { } fn get_cached<'a>( + &'a mut self, dfg: &DataFlowGraph, - instruction_result_cache: &'a mut InstructionResultCache, instruction: &Instruction, side_effects_enabled_var: ValueId, - use_constraint_info: bool, block: BasicBlockId, - dom_tree: &mut DominatorTree, ) -> Option<&'a Vec> { - let results_for_instruction = instruction_result_cache.get(instruction)?; + let results_for_instruction = self.cached_instruction_results.get(instruction)?; - let predicate = use_constraint_info && instruction.requires_acir_gen_predicate(dfg); + 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, dom_tree) + results_for_instruction.get(&predicate)?.get(block, &mut self.dom) } } From 1b6e15702627401dbf240f03bffd5fce2842d73b Mon Sep 17 00:00:00 2001 From: Jake Fecher Date: Mon, 11 Nov 2024 09:23:50 -0600 Subject: [PATCH 4/5] Fix and add test --- .../src/ssa/opt/constant_folding.rs | 54 +++++++++++++++++-- 1 file changed, 49 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 d7c81e26252..bea0a9e40e8 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -19,7 +19,7 @@ //! //! This is the only pass which removes duplicated pure [`Instruction`]s however and so is needed when //! different blocks are merged, i.e. after the [`flatten_cfg`][super::flatten_cfg] pass. -use std::collections::HashSet; +use std::collections::{HashSet, VecDeque}; use acvm::{acir::AcirField, FieldElement}; use iter_extended::vecmap; @@ -69,9 +69,9 @@ impl Function { /// Go through each block and re-insert all instructions. pub(crate) fn constant_fold(&mut self, use_constraint_info: bool) { let mut context = Context::new(self, use_constraint_info); - context.block_queue.push(self.entry_block()); + context.block_queue.push_back(self.entry_block()); - while let Some(block) = context.block_queue.pop() { + while let Some(block) = context.block_queue.pop_front() { if context.visited_blocks.contains(&block) { continue; } @@ -86,7 +86,7 @@ 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, + block_queue: VecDeque, // Contains sets of values which are constrained to be equivalent to each other. // @@ -337,7 +337,7 @@ impl ResultCache { fn get(&self, block: BasicBlockId, dom: &mut DominatorTree) -> Option<&Vec> { for (origin_block, results) in &self.results { - if dom.dominates(block, *origin_block) { + if dom.dominates(*origin_block, block) { return Some(results); } } @@ -930,4 +930,48 @@ mod test { let ending_instruction_count = instructions.len(); assert_eq!(ending_instruction_count, 1); } + + #[test] + fn deduplicate_across_blocks() { + // fn main f0 { + // b0(v0: u1): + // v1 = not v0 + // jmp b1() + // b1(): + // v2 = not v0 + // return v2 + // } + let main_id = Id::test_new(0); + + // Compiling main + let mut builder = FunctionBuilder::new("main".into(), main_id); + let b1 = builder.insert_block(); + + let v0 = builder.add_parameter(Type::bool()); + let _v1 = builder.insert_not(v0); + builder.terminate_with_jmp(b1, Vec::new()); + + builder.switch_to_block(b1); + let v2 = builder.insert_not(v0); + builder.terminate_with_return(vec![v2]); + + let ssa = builder.finish(); + let main = ssa.main(); + assert_eq!(main.dfg[main.entry_block()].instructions().len(), 1); + assert_eq!(main.dfg[b1].instructions().len(), 1); + + // Expected output: + // + // fn main f0 { + // b0(v0: u1): + // v1 = not v0 + // jmp b1() + // b1(): + // return v1 + // } + let ssa = ssa.fold_constants_using_constraints(); + let main = ssa.main(); + assert_eq!(main.dfg[main.entry_block()].instructions().len(), 1); + assert_eq!(main.dfg[b1].instructions().len(), 0); + } } From 666c1bcfadb25068dd9c77b7367d42f0cbc8410a Mon Sep 17 00:00:00 2001 From: Tom French Date: Mon, 11 Nov 2024 16:51:27 +0000 Subject: [PATCH 5/5] chore: nits and documentation --- .../src/ssa/opt/constant_folding.rs | 28 +++++++++++++------ 1 file changed, 19 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 bea0a9e40e8..96ac1e8a6ac 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs @@ -88,13 +88,13 @@ struct Context { visited_blocks: HashSet, block_queue: VecDeque, - // 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. + /// 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. constraint_simplification_mappings: HashMap>, // Cache of instructions without any side-effects along with their outputs. @@ -110,6 +110,9 @@ struct Context { /// to deduplicate instructions across blocks as long as the new block dominates the original. type InstructionResultCache = HashMap, ResultCache>>; +/// Records the results of all duplicate [`Instruction`]s along with the blocks in which they sit. +/// +/// For more information see [`InstructionResultCache`]. #[derive(Default)] struct ResultCache { results: Vec<(BasicBlockId, Vec)>, @@ -320,7 +323,7 @@ impl Context { instruction: &Instruction, side_effects_enabled_var: ValueId, block: BasicBlockId, - ) -> Option<&'a Vec> { + ) -> Option<&'a [ValueId]> { let results_for_instruction = self.cached_instruction_results.get(instruction)?; let predicate = self.use_constraint_info && instruction.requires_acir_gen_predicate(dfg); @@ -331,11 +334,18 @@ impl Context { } impl ResultCache { + /// Records that an `Instruction` in block `block` produced the result values `results`. fn cache(&mut self, block: BasicBlockId, results: Vec) { self.results.push((block, results)); } - fn get(&self, block: BasicBlockId, dom: &mut DominatorTree) -> Option<&Vec> { + /// Returns a set of [`ValueId`]s produced from a copy of this [`Instruction`] which sits + /// within a block which dominates `block`. + /// + /// 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<&[ValueId]> { for (origin_block, results) in &self.results { if dom.dominates(*origin_block, block) { return Some(results);