From 7480b0e3d35eb7d052a080ab8ccb91a028583daf Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Mon, 25 Nov 2024 23:03:44 +0000 Subject: [PATCH 1/4] bring back local allocations set in flattening --- .../src/ssa/opt/flatten_cfg.rs | 30 +++++++++++++++++-- .../noirc_evaluator/src/ssa/opt/mem2reg.rs | 13 +++----- 2 files changed, 32 insertions(+), 11 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs index 61a93aee58d..522ef297392 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs @@ -131,7 +131,7 @@ //! v11 = mul v4, Field 12 //! v12 = add v10, v11 //! store v12 at v5 (new store) -use fxhash::FxHashMap as HashMap; +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; use acvm::{acir::AcirField, acir::BlackBoxFunc, FieldElement}; use iter_extended::vecmap; @@ -201,6 +201,13 @@ struct Context<'f> { /// When processing a block, we pop this stack to get its arguments /// and at the end we push the arguments for his successor arguments_stack: Vec>, + + /// Stores all allocations local to the current branch. + /// Since these branches are local to the current branch (ie. only defined within one branch of + /// an if expression), they should not be merged with their previous value or stored value in + /// the other branch since there is no such value. The ValueId here is that which is returned + /// by the allocate instruction. + local_allocations: HashSet, } #[derive(Clone)] @@ -211,6 +218,8 @@ struct ConditionalBranch { old_condition: ValueId, // The condition of the branch condition: ValueId, + // The allocations accumulated when processing the branch + local_allocations: HashSet, } struct ConditionalContext { @@ -243,6 +252,7 @@ fn flatten_function_cfg(function: &mut Function, no_predicates: &HashMap Context<'f> { let old_condition = *condition; let then_condition = self.inserter.resolve(old_condition); + let old_allocations = std::mem::take(&mut self.local_allocations); let branch = ConditionalBranch { old_condition, condition: self.link_condition(then_condition), last_block: *then_destination, + local_allocations: old_allocations, }; let cond_context = ConditionalContext { condition: then_condition, @@ -435,11 +447,14 @@ impl<'f> Context<'f> { ); let else_condition = self.link_condition(else_condition); + let old_allocations = std::mem::take(&mut self.local_allocations); let else_branch = ConditionalBranch { old_condition: cond_context.then_branch.old_condition, condition: else_condition, last_block: *block, + local_allocations: old_allocations, }; + cond_context.then_branch.local_allocations.clear(); cond_context.else_branch = Some(else_branch); self.condition_stack.push(cond_context); @@ -461,6 +476,7 @@ impl<'f> Context<'f> { } let mut else_branch = cond_context.else_branch.unwrap(); + self.local_allocations = std::mem::take(&mut else_branch.local_allocations); else_branch.last_block = *block; cond_context.else_branch = Some(else_branch); @@ -604,10 +620,18 @@ impl<'f> Context<'f> { call_stack.clone(), *previous_allocate_result, ); + let is_allocate = matches!(instruction, Instruction::Allocate); let instruction_is_allocate = matches!(&instruction, Instruction::Allocate); let entry = self.inserter.function.entry_block(); let results = self.inserter.push_instruction_value(instruction, id, entry, call_stack); + + // Remember an allocate was created local to this branch so that we do not try to merge store + // values across branches for it later. + if is_allocate { + self.local_allocations.insert(results.first()); + } + *previous_allocate_result = instruction_is_allocate.then(|| results.first()); } @@ -652,7 +676,9 @@ impl<'f> Context<'f> { Instruction::Store { address, value } => { // If this instruction immediately follows an allocate, and stores to that // address there is no previous value to load and we don't need a merge anyway. - if Some(address) == previous_allocate_result { + if Some(address) == previous_allocate_result + || self.local_allocations.contains(&address) + { Instruction::Store { address, value } } else { // Instead of storing `value`, store `if condition { value } else { previous_value }` diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs index 77133d7d88d..7f30948f6ae 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs @@ -415,13 +415,11 @@ impl<'f> PerFunctionContext<'f> { let address = self.inserter.function.dfg.resolve(*address); let value = self.inserter.function.dfg.resolve(*value); - // FIXME: This causes errors in the sha256 tests - // // If there was another store to this instruction without any (unremoved) loads or // function calls in-between, we can remove the previous store. - // if let Some(last_store) = references.last_stores.get(&address) { - // self.instructions_to_remove.insert(*last_store); - // } + if let Some(last_store) = references.last_stores.get(&address) { + self.instructions_to_remove.insert(*last_store); + } if self.inserter.function.dfg.value_is_reference(value) { if let Some(expression) = references.expressions.get(&value) { @@ -876,7 +874,6 @@ mod tests { // acir fn main f0 { // b0(): // v9 = allocate - // store Field 0 at v9 // v10 = allocate // jmp b1() // b1(): @@ -901,9 +898,7 @@ mod tests { // in the same block, and the store is not needed before the later store. // The rest of the stores are also removed as no loads are done within any blocks // to the stored values. - // - // NOTE: This store is not removed due to the FIXME when handling Instruction::Store. - assert_eq!(count_stores(b1, &main.dfg), 1); + assert_eq!(count_stores(b1, &main.dfg), 0); let b1_instructions = main.dfg[b1].instructions(); From 321a777a87a237bb76637a7bfad56a745a771981 Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Mon, 25 Nov 2024 23:27:48 +0000 Subject: [PATCH 2/4] fix load_aliases_in_predecessor_block --- .../src/ssa/opt/flatten_cfg.rs | 27 +++++-------------- .../noirc_evaluator/src/ssa/opt/mem2reg.rs | 5 ++-- 2 files changed, 9 insertions(+), 23 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs index 522ef297392..fa0f383ac4e 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs @@ -327,7 +327,6 @@ impl<'f> Context<'f> { // If this is not a separate variable, clippy gets confused and says the to_vec is // unnecessary, when removing it actually causes an aliasing/mutability error. let instructions = self.inserter.function.dfg[block].instructions().to_vec(); - let mut previous_allocate_result = None; for instruction in instructions.iter() { if self.is_no_predicate(no_predicates, instruction) { @@ -342,10 +341,10 @@ impl<'f> Context<'f> { None, im::Vector::new(), ); - self.push_instruction(*instruction, &mut previous_allocate_result); + self.push_instruction(*instruction); self.insert_current_side_effects_enabled(); } else { - self.push_instruction(*instruction, &mut previous_allocate_result); + self.push_instruction(*instruction); } } } @@ -609,18 +608,9 @@ impl<'f> Context<'f> { /// `previous_allocate_result` should only be set to the result of an allocate instruction /// if that instruction was the instruction immediately previous to this one - if there are /// any instructions in between it should be None. - fn push_instruction( - &mut self, - id: InstructionId, - previous_allocate_result: &mut Option, - ) { + fn push_instruction(&mut self, id: InstructionId) { let (instruction, call_stack) = self.inserter.map_instruction(id); - let instruction = self.handle_instruction_side_effects( - instruction, - call_stack.clone(), - *previous_allocate_result, - ); - let is_allocate = matches!(instruction, Instruction::Allocate); + let instruction = self.handle_instruction_side_effects(instruction, call_stack.clone()); let instruction_is_allocate = matches!(&instruction, Instruction::Allocate); let entry = self.inserter.function.entry_block(); @@ -628,11 +618,9 @@ impl<'f> Context<'f> { // Remember an allocate was created local to this branch so that we do not try to merge store // values across branches for it later. - if is_allocate { + if instruction_is_allocate { self.local_allocations.insert(results.first()); } - - *previous_allocate_result = instruction_is_allocate.then(|| results.first()); } /// If we are currently in a branch, we need to modify constrain instructions @@ -645,7 +633,6 @@ impl<'f> Context<'f> { &mut self, instruction: Instruction, call_stack: CallStack, - previous_allocate_result: Option, ) -> Instruction { if let Some(condition) = self.get_last_condition() { match instruction { @@ -676,9 +663,7 @@ impl<'f> Context<'f> { Instruction::Store { address, value } => { // If this instruction immediately follows an allocate, and stores to that // address there is no previous value to load and we don't need a merge anyway. - if Some(address) == previous_allocate_result - || self.local_allocations.contains(&address) - { + if self.local_allocations.contains(&address) { Instruction::Store { address, value } } else { // Instead of storing `value`, store `if condition { value } else { previous_value }` diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs index 7f30948f6ae..9bbfe963d48 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs @@ -874,6 +874,7 @@ mod tests { // acir fn main f0 { // b0(): // v9 = allocate + // store Field 0 at v9 // v10 = allocate // jmp b1() // b1(): @@ -902,8 +903,8 @@ mod tests { let b1_instructions = main.dfg[b1].instructions(); - // We expect the last eq to be optimized out, only the store from above remains - assert_eq!(b1_instructions.len(), 1); + // We expect the last eq to be optimized out + assert_eq!(b1_instructions.len(), 0); } #[test] From 413869d4071b7364626f03d025794f2c95351c6a Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Mon, 25 Nov 2024 23:52:37 +0000 Subject: [PATCH 3/4] switch load_aliases_in_predecessor_block to ssa parser --- .../noirc_evaluator/src/ssa/opt/mem2reg.rs | 110 ++++++------------ 1 file changed, 38 insertions(+), 72 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs index 9bbfe963d48..a48c57cdb58 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/mem2reg.rs @@ -612,6 +612,8 @@ mod tests { map::Id, types::Type, }, + opt::assert_normalized_ssa_equals, + Ssa, }; #[test] @@ -822,89 +824,53 @@ mod tests { // is later stored in a successor block #[test] fn load_aliases_in_predecessor_block() { - // fn main { - // b0(): - // v0 = allocate - // store Field 0 at v0 - // v2 = allocate - // store v0 at v2 - // v3 = load v2 - // v4 = load v2 - // jmp b1() - // b1(): - // store Field 1 at v3 - // store Field 2 at v4 - // v7 = load v3 - // v8 = eq v7, Field 2 - // return - // } - let main_id = Id::test_new(0); - let mut builder = FunctionBuilder::new("main".into(), main_id); - - let v0 = builder.insert_allocate(Type::field()); - - let zero = builder.field_constant(0u128); - builder.insert_store(v0, zero); - - let v2 = builder.insert_allocate(Type::Reference(Arc::new(Type::field()))); - builder.insert_store(v2, v0); - - let v3 = builder.insert_load(v2, Type::field()); - let v4 = builder.insert_load(v2, Type::field()); - let b1 = builder.insert_block(); - builder.terminate_with_jmp(b1, vec![]); - - builder.switch_to_block(b1); - - let one = builder.field_constant(1u128); - builder.insert_store(v3, one); - - let two = builder.field_constant(2u128); - builder.insert_store(v4, two); - - let v8 = builder.insert_load(v3, Type::field()); - let _ = builder.insert_binary(v8, BinaryOp::Eq, two); - - builder.terminate_with_return(vec![]); - - let ssa = builder.finish(); - assert_eq!(ssa.main().reachable_blocks().len(), 2); + let src = " + acir(inline) fn main f0 { + b0(): + v0 = allocate -> &mut Field + store Field 0 at v0 + v2 = allocate -> &mut &mut Field + store v0 at v2 + v3 = load v2 -> &mut Field + v4 = load v2 -> &mut Field + jmp b1() + b1(): + store Field 1 at v3 + store Field 2 at v4 + v7 = load v3 -> Field + v8 = eq v7, Field 2 + return + } + "; - // Expected result: - // acir fn main f0 { - // b0(): - // v9 = allocate - // store Field 0 at v9 - // v10 = allocate - // jmp b1() - // b1(): - // return - // } - let ssa = ssa.mem2reg(); - println!("{}", ssa); + let mut ssa = Ssa::from_str(src).unwrap(); + let main = ssa.main_mut(); - let main = ssa.main(); - assert_eq!(main.reachable_blocks().len(), 2); + let instructions = main.dfg[main.entry_block()].instructions(); + assert_eq!(instructions.len(), 6); // The final return is not counted // All loads should be removed - assert_eq!(count_loads(main.entry_block(), &main.dfg), 0); - assert_eq!(count_loads(b1, &main.dfg), 0); - // The first store is not removed as it is used as a nested reference in another store. - // We would need to track whether the store where `v9` is the store value gets removed to know whether + // We would need to track whether the store where `v0` is the store value gets removed to know whether // to remove it. - assert_eq!(count_stores(main.entry_block(), &main.dfg), 1); - // The first store in b1 is removed since there is another store to the same reference // in the same block, and the store is not needed before the later store. // The rest of the stores are also removed as no loads are done within any blocks // to the stored values. - assert_eq!(count_stores(b1, &main.dfg), 0); - - let b1_instructions = main.dfg[b1].instructions(); + let expected = " + acir(inline) fn main f0 { + b0(): + v0 = allocate -> &mut Field + store Field 0 at v0 + v2 = allocate -> &mut &mut Field + jmp b1() + b1(): + return + } + "; - // We expect the last eq to be optimized out - assert_eq!(b1_instructions.len(), 0); + let ssa = ssa.mem2reg(); + assert_normalized_ssa_equals(ssa, expected); } #[test] From 2449d22f1f3a57792669dbc4e1507b55d9149134 Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Mon, 25 Nov 2024 18:54:26 -0500 Subject: [PATCH 4/4] Update compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs Co-authored-by: Tom French <15848336+TomAFrench@users.noreply.github.com> --- compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs index fa0f383ac4e..034c854837e 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs @@ -203,10 +203,12 @@ struct Context<'f> { arguments_stack: Vec>, /// Stores all allocations local to the current branch. - /// Since these branches are local to the current branch (ie. only defined within one branch of + /// + /// Since these branches are local to the current branch (i.e. only defined within one branch of /// an if expression), they should not be merged with their previous value or stored value in - /// the other branch since there is no such value. The ValueId here is that which is returned - /// by the allocate instruction. + /// the other branch since there is no such value. + /// + /// The `ValueId` here is that which is returned by the allocate instruction. local_allocations: HashSet, }