From 13a01c727eeac258a953971148cedfffd49fc29e Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Wed, 25 Sep 2024 14:51:51 +0000 Subject: [PATCH 1/7] handle array set optimization across blocks --- .../noirc_evaluator/src/ssa/opt/array_set.rs | 61 +++++++++++++------ 1 file changed, 42 insertions(+), 19 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs index 491a17adb66..5f56432a149 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs @@ -4,10 +4,11 @@ use crate::ssa::{ dfg::DataFlowGraph, instruction::{Instruction, InstructionId}, types::Type::{Array, Slice}, + value::ValueId, }, ssa_gen::Ssa, }; -use fxhash::{FxHashMap as HashMap, FxHashSet}; +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; impl Ssa { /// Map arrays with the last instruction that uses it @@ -16,28 +17,41 @@ impl Ssa { #[tracing::instrument(level = "trace", skip(self))] pub(crate) fn array_set_optimization(mut self) -> Self { for func in self.functions.values_mut() { - let mut reachable_blocks = func.reachable_blocks(); - let block = if !func.runtime().is_entry_point() { - assert_eq!(reachable_blocks.len(), 1, "Expected there to be 1 block remaining in Acir function for array_set optimization"); - reachable_blocks.pop_first().unwrap() - } else { - // We only apply the array set optimization in the return block of Brillig functions - func.find_last_block() - }; + let reachable_blocks = func.reachable_blocks(); - let instructions_to_update = analyze_last_uses(&func.dfg, block); - make_mutable(&mut func.dfg, block, instructions_to_update); + if !func.runtime().is_entry_point() { + assert_eq!(reachable_blocks.len(), 1, "Expected there to be 1 block remaining in Acir function for array_set optimization"); + } + let mut array_to_last_use = HashMap::default(); + let mut instructions_to_update = HashSet::default(); + let mut arrays_from_load = HashSet::default(); + for block in reachable_blocks.iter() { + analyze_last_uses_new( + &func.dfg, + *block, + &mut array_to_last_use, + &mut instructions_to_update, + &mut arrays_from_load, + ); + } + for block in reachable_blocks { + make_mutable(&mut func.dfg, block, &instructions_to_update); + } } self } } -/// Returns the set of ArraySet instructions that can be made mutable +/// Builds the set of ArraySet instructions that can be made mutable /// because their input value is unused elsewhere afterward. -fn analyze_last_uses(dfg: &DataFlowGraph, block_id: BasicBlockId) -> FxHashSet { +fn analyze_last_uses_new( + dfg: &DataFlowGraph, + block_id: BasicBlockId, + array_to_last_use: &mut HashMap, + instructions_that_can_be_made_mutable: &mut HashSet, + arrays_from_load: &mut HashSet, +) { let block = &dfg[block_id]; - let mut array_to_last_use = HashMap::default(); - let mut instructions_that_can_be_made_mutable = FxHashSet::default(); for instruction_id in block.instructions() { match &dfg[*instruction_id] { @@ -54,7 +68,12 @@ fn analyze_last_uses(dfg: &DataFlowGraph, block_id: BasicBlockId) -> FxHashSet { for argument in arguments { @@ -68,18 +87,22 @@ fn analyze_last_uses(dfg: &DataFlowGraph, block_id: BasicBlockId) -> FxHashSet { + let result = dfg.instruction_results(*instruction_id)[0]; + if matches!(dfg.type_of_value(result), Array { .. } | Slice { .. }) { + arrays_from_load.insert(result); + } + } _ => (), } } - - instructions_that_can_be_made_mutable } /// Make each ArraySet instruction in `instructions_to_update` mutable. fn make_mutable( dfg: &mut DataFlowGraph, block_id: BasicBlockId, - instructions_to_update: FxHashSet, + instructions_to_update: &HashSet, ) { if instructions_to_update.is_empty() { return; From ede8c2a37f386869194acf33429523e28b89d9a9 Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Wed, 25 Sep 2024 17:50:08 +0000 Subject: [PATCH 2/7] do not check whether an array is from a load when in a return block --- compiler/noirc_evaluator/src/ssa/opt/array_set.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs index 5f56432a149..d5034608e60 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs @@ -2,7 +2,7 @@ use crate::ssa::{ ir::{ basic_block::BasicBlockId, dfg::DataFlowGraph, - instruction::{Instruction, InstructionId}, + instruction::{Instruction, InstructionId, TerminatorInstruction}, types::Type::{Array, Slice}, value::ValueId, }, @@ -71,7 +71,11 @@ fn analyze_last_uses_new( // If the array we are setting does not come from a load we can safely mark it mutable. // If the array comes from a load we may potentially being mutating an array at a reference // that is loaded from by other values. - if !arrays_from_load.contains(&array) { + let is_return_block = matches!( + dfg[block_id].unwrap_terminator(), + TerminatorInstruction::Return { .. } + ); + if !arrays_from_load.contains(&array) || is_return_block { instructions_that_can_be_made_mutable.insert(*instruction_id); } } From 532eb1235ff9525621d834f0d8fcf2dc4341e973 Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Wed, 25 Sep 2024 18:06:32 +0000 Subject: [PATCH 3/7] rename func --- compiler/noirc_evaluator/src/ssa/opt/array_set.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs index d5034608e60..306e04dfc0d 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs @@ -26,7 +26,7 @@ impl Ssa { let mut instructions_to_update = HashSet::default(); let mut arrays_from_load = HashSet::default(); for block in reachable_blocks.iter() { - analyze_last_uses_new( + analyze_last_uses( &func.dfg, *block, &mut array_to_last_use, @@ -44,7 +44,7 @@ impl Ssa { /// Builds the set of ArraySet instructions that can be made mutable /// because their input value is unused elsewhere afterward. -fn analyze_last_uses_new( +fn analyze_last_uses( dfg: &DataFlowGraph, block_id: BasicBlockId, array_to_last_use: &mut HashMap, @@ -52,6 +52,7 @@ fn analyze_last_uses_new( arrays_from_load: &mut HashSet, ) { let block = &dfg[block_id]; + let block_params = block.parameters(); for instruction_id in block.instructions() { match &dfg[*instruction_id] { @@ -71,11 +72,14 @@ fn analyze_last_uses_new( // If the array we are setting does not come from a load we can safely mark it mutable. // If the array comes from a load we may potentially being mutating an array at a reference // that is loaded from by other values. + // If we are in a return block we are not concerned about the array potentially being mutated again. let is_return_block = matches!( dfg[block_id].unwrap_terminator(), TerminatorInstruction::Return { .. } ); - if !arrays_from_load.contains(&array) || is_return_block { + if (!arrays_from_load.contains(&array) || is_return_block) + && !block_params.contains(&array) + { instructions_that_can_be_made_mutable.insert(*instruction_id); } } From ce59f2a38202842d95c6e0bde0d453f6b2ae0ca5 Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Wed, 25 Sep 2024 18:17:18 +0000 Subject: [PATCH 4/7] do not check array in block param, check array in terminator --- .../noirc_evaluator/src/ssa/opt/array_set.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs index 306e04dfc0d..ef901e0f0b2 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs @@ -25,6 +25,7 @@ impl Ssa { let mut array_to_last_use = HashMap::default(); let mut instructions_to_update = HashSet::default(); let mut arrays_from_load = HashSet::default(); + for block in reachable_blocks.iter() { analyze_last_uses( &func.dfg, @@ -52,7 +53,6 @@ fn analyze_last_uses( arrays_from_load: &mut HashSet, ) { let block = &dfg[block_id]; - let block_params = block.parameters(); for instruction_id in block.instructions() { match &dfg[*instruction_id] { @@ -73,13 +73,15 @@ fn analyze_last_uses( // If the array comes from a load we may potentially being mutating an array at a reference // that is loaded from by other values. // If we are in a return block we are not concerned about the array potentially being mutated again. - let is_return_block = matches!( - dfg[block_id].unwrap_terminator(), - TerminatorInstruction::Return { .. } - ); - if (!arrays_from_load.contains(&array) || is_return_block) - && !block_params.contains(&array) - { + let terminator = dfg[block_id].unwrap_terminator(); + let is_return_block = matches!(terminator, TerminatorInstruction::Return { .. }); + let mut array_in_terminator = false; + terminator.for_each_value(|value| { + if value == array { + array_in_terminator = true; + } + }); + if (!arrays_from_load.contains(&array) || is_return_block) && !array_in_terminator { instructions_that_can_be_made_mutable.insert(*instruction_id); } } From 9eeddf41519b2282ae24d8185ea4e44255413033 Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Wed, 25 Sep 2024 20:08:43 +0000 Subject: [PATCH 5/7] add unit test to array sets opt --- .../noirc_evaluator/src/ssa/opt/array_set.rs | 129 +++++++++++++++++- 1 file changed, 128 insertions(+), 1 deletion(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs index ef901e0f0b2..055bc101446 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs @@ -72,9 +72,10 @@ fn analyze_last_uses( // If the array we are setting does not come from a load we can safely mark it mutable. // If the array comes from a load we may potentially being mutating an array at a reference // that is loaded from by other values. - // If we are in a return block we are not concerned about the array potentially being mutated again. let terminator = dfg[block_id].unwrap_terminator(); + // If we are in a return block we are not concerned about the array potentially being mutated again. let is_return_block = matches!(terminator, TerminatorInstruction::Return { .. }); + // We also want to check that the array is not part of the terminator arguments, as this means it is used again. let mut array_in_terminator = false; terminator.for_each_value(|value| { if value == array { @@ -138,3 +139,129 @@ fn make_mutable( *dfg[block_id].instructions_mut() = instructions; } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use im::vector; + + use crate::ssa::{ + function_builder::FunctionBuilder, + ir::{ + function::RuntimeType, + instruction::{self, BinaryOp, Instruction}, + map::Id, + types::Type, + }, + }; + + #[test] + fn array_set_in_loop_with_conditional_clone() { + // We want to make sure that we do not mark a single array set mutable which is loaded + // from and cloned in a loop. If the array is inadvertently marked mutable, and is cloned in a previous iteration + // of the loop, its clone will also be altered. + // + // acir(inline) fn main f0 { + // b0(): + // v2 = allocate + // store [Field 0, Field 0, Field 0, Field 0, Field 0] at v2 + // v3 = allocate + // store [Field 0, Field 0, Field 0, Field 0, Field 0] at v3 + // jmp b1(u32 0) + // b1(v5: u32): + // v7 = lt v5, u32 5 + // jmpif v7 then: b3, else: b2 + // b3(): + // v8 = eq v5, u32 5 + // jmpif v8 then: b4, else: b5 + // b4(): + // v9 = load v2 + // store v9 at v3 + // jmp b5() + // b5(): + // v10 = load v2 + // v12 = array_set v10, index v5, value Field 20 + // store v12 at v2 + // v14 = add v5, u32 1 + // jmp b1(v14) + // b2(): + // return + // } + let main_id = Id::test_new(0); + let mut builder = FunctionBuilder::new("main".into(), main_id); + builder.set_runtime(RuntimeType::Brillig); + + let array_type = Type::Array(Arc::new(vec![Type::field()]), 5); + let zero = builder.field_constant(0u128); + let array_constant = + builder.array_constant(vector![zero, zero, zero, zero, zero], array_type.clone()); + + let v2 = builder.insert_allocate(array_type.clone()); + + builder.insert_store(v2, array_constant); + + let v3 = builder.insert_allocate(array_type.clone()); + builder.insert_store(v3, array_constant); + + let b1 = builder.insert_block(); + let zero_u32 = builder.numeric_constant(0u128, Type::unsigned(32)); + builder.terminate_with_jmp(b1, vec![zero_u32]); + + // Loop header + builder.switch_to_block(b1); + let v5 = builder.add_block_parameter(b1, Type::unsigned(32)); + let five = builder.numeric_constant(5u128, Type::unsigned(32)); + let v7 = builder.insert_binary(v5, BinaryOp::Lt, five); + + let b2 = builder.insert_block(); + let b3 = builder.insert_block(); + let b4 = builder.insert_block(); + let b5 = builder.insert_block(); + builder.terminate_with_jmpif(v7, b3, b2); + + // Loop body + // b3 is the if statement conditional + builder.switch_to_block(b3); + let two = builder.numeric_constant(5u128, Type::unsigned(32)); + let v8 = builder.insert_binary(v5, BinaryOp::Eq, two); + builder.terminate_with_jmpif(v8, b4, b5); + + // b4 is the rest of the loop after the if statement + builder.switch_to_block(b4); + let v9 = builder.insert_load(v2, array_type.clone()); + builder.insert_store(v3, v9); + builder.terminate_with_jmp(b5, vec![]); + + builder.switch_to_block(b5); + let v10 = builder.insert_load(v2, array_type.clone()); + let twenty = builder.field_constant(20u128); + let v12 = builder.insert_array_set(v10, v5, twenty); + builder.insert_store(v2, v12); + let one = builder.numeric_constant(1u128, Type::unsigned(32)); + let v14 = builder.insert_binary(v5, BinaryOp::Add, one); + builder.terminate_with_jmp(b1, vec![v14]); + + builder.switch_to_block(b2); + builder.terminate_with_return(vec![]); + + let ssa = builder.finish(); + // We expect the same result as above + let ssa = ssa.array_set_optimization(); + + let main = ssa.main(); + assert_eq!(main.reachable_blocks().len(), 6); + + let array_set_instructions = main.dfg[b5] + .instructions() + .into_iter() + .filter(|instruction| matches!(&main.dfg[**instruction], Instruction::ArraySet { .. })) + .collect::>(); + + assert_eq!(array_set_instructions.len(), 1); + if let Instruction::ArraySet { mutable, .. } = &main.dfg[*array_set_instructions[0]] { + // The single array set should not be marked mutable + assert_eq!(*mutable, false) + } + } +} From 5c5783b686b2d3bdcfc323b5afa1e9322b832b99 Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Wed, 25 Sep 2024 20:16:22 +0000 Subject: [PATCH 6/7] clippy --- compiler/noirc_evaluator/src/ssa/opt/array_set.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs index 055bc101446..15550b21855 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs @@ -254,14 +254,14 @@ mod tests { let array_set_instructions = main.dfg[b5] .instructions() - .into_iter() + .iter() .filter(|instruction| matches!(&main.dfg[**instruction], Instruction::ArraySet { .. })) .collect::>(); assert_eq!(array_set_instructions.len(), 1); if let Instruction::ArraySet { mutable, .. } = &main.dfg[*array_set_instructions[0]] { // The single array set should not be marked mutable - assert_eq!(*mutable, false) + assert!(!mutable); } } } From b7d707403ec0980e90aeb2f7d936a738cf04267f Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Wed, 25 Sep 2024 20:24:43 +0000 Subject: [PATCH 7/7] clippy once more --- compiler/noirc_evaluator/src/ssa/opt/array_set.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs index 15550b21855..6d48b8c0d67 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs @@ -150,7 +150,7 @@ mod tests { function_builder::FunctionBuilder, ir::{ function::RuntimeType, - instruction::{self, BinaryOp, Instruction}, + instruction::{BinaryOp, Instruction}, map::Id, types::Type, },