From 5598059576c6cbc72474aff4b18bc5e4bb9f08e1 Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Fri, 20 Sep 2024 15:58:19 -0400 Subject: [PATCH] feat(perf): Allow array set last uses optimization in return block of Brillig functions (#6119) # Description ## Problem\* Resolves Part of general effort to reduce Brillig bytecode sizes. ## Summary\* We currently only perform the array set last uses optimization on ACIR functions. This makes sense as the pass only looks only last uses in a single block. We can start to expand the last uses opt to encompass more than just ACIR functions and make it aware of the last array set of an array across blocks. This is quite complex though so this PR currently just attempts to look at the return block of Brillig functions and see the effects. ## Additional Context ## Documentation\* Check one: - [X] No documentation needed. - [ ] Documentation included in this PR. - [ ] **[For Experimental Features]** Documentation to be submitted in a separate PR. # PR Checklist\* - [X] I have tested the changes locally. - [X] I have formatted the changes with [Prettier](https://prettier.io/) and/or `cargo fmt` on default settings. --- .../src/brillig/brillig_gen/brillig_block.rs | 24 +++++++++++++++---- .../noirc_evaluator/src/ssa/ir/function.rs | 11 +++++++++ .../noirc_evaluator/src/ssa/opt/array_set.rs | 15 +++++++----- compiler/noirc_evaluator/src/ssa/opt/rc.rs | 19 ++------------- 4 files changed, 42 insertions(+), 27 deletions(-) diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs index c4f0e944099..8929122f515 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs @@ -666,7 +666,7 @@ impl<'block> BrilligBlock<'block> { self.brillig_context.deallocate_register(items_pointer); } - Instruction::ArraySet { array, index, value, mutable: _ } => { + Instruction::ArraySet { array, index, value, mutable } => { let source_variable = self.convert_ssa_value(*array, dfg); let index_register = self.convert_ssa_single_addr_value(*index, dfg); let value_variable = self.convert_ssa_value(*value, dfg); @@ -688,6 +688,7 @@ impl<'block> BrilligBlock<'block> { destination_variable, index_register, value_variable, + *mutable, ); } Instruction::RangeCheck { value, max_bit_size, assert_message } => { @@ -868,6 +869,7 @@ impl<'block> BrilligBlock<'block> { destination_variable: BrilligVariable, index_register: SingleAddrVariable, value_variable: BrilligVariable, + mutable: bool, ) { assert!(index_register.bit_size == BRILLIG_MEMORY_ADDRESSING_BIT_SIZE); match (source_variable, destination_variable) { @@ -875,20 +877,26 @@ impl<'block> BrilligBlock<'block> { BrilligVariable::BrilligArray(source_array), BrilligVariable::BrilligArray(destination_array), ) => { - self.brillig_context.call_array_copy_procedure(source_array, destination_array); + if !mutable { + self.brillig_context.call_array_copy_procedure(source_array, destination_array); + } } ( BrilligVariable::BrilligVector(source_vector), BrilligVariable::BrilligVector(destination_vector), ) => { - self.brillig_context.call_vector_copy_procedure(source_vector, destination_vector); + if !mutable { + self.brillig_context + .call_vector_copy_procedure(source_vector, destination_vector); + } } _ => unreachable!("ICE: array set on non-array"), } + let destination_for_store = if mutable { source_variable } else { destination_variable }; // Then set the value in the newly created array let items_pointer = - self.brillig_context.codegen_make_array_or_vector_items_pointer(destination_variable); + self.brillig_context.codegen_make_array_or_vector_items_pointer(destination_for_store); self.brillig_context.codegen_store_with_offset( items_pointer, @@ -896,6 +904,14 @@ impl<'block> BrilligBlock<'block> { value_variable.extract_register(), ); + // If we mutated the source array we want instructions that use the destination array to point to the source array + if mutable { + self.brillig_context.mov_instruction( + destination_variable.extract_register(), + source_variable.extract_register(), + ); + } + self.brillig_context.deallocate_register(items_pointer); } diff --git a/compiler/noirc_evaluator/src/ssa/ir/function.rs b/compiler/noirc_evaluator/src/ssa/ir/function.rs index 65a616ef612..1466f2e5d44 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/function.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/function.rs @@ -160,6 +160,17 @@ impl Function { let returns = vecmap(self.returns(), |ret| self.dfg.type_of_value(*ret)); Signature { params, returns } } + + /// Finds the block of the function with the Return instruction + pub(crate) fn find_last_block(&self) -> BasicBlockId { + for block in self.reachable_blocks() { + if matches!(self.dfg[block].terminator(), Some(TerminatorInstruction::Return { .. })) { + return block; + } + } + + unreachable!("SSA Function {} has no reachable return instruction!", self.id()) + } } impl std::fmt::Display for RuntimeType { diff --git a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs index cf61d7fd73f..491a17adb66 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs @@ -16,14 +16,17 @@ impl Ssa { #[tracing::instrument(level = "trace", skip(self))] pub(crate) fn array_set_optimization(mut self) -> Self { for func in self.functions.values_mut() { - if !func.runtime().is_entry_point() { - let mut reachable_blocks = func.reachable_blocks(); + 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 block = reachable_blocks.pop_first().unwrap(); - let instructions_to_update = analyze_last_uses(&func.dfg, block); - make_mutable(&mut func.dfg, block, instructions_to_update); - } + let instructions_to_update = analyze_last_uses(&func.dfg, block); + make_mutable(&mut func.dfg, block, instructions_to_update); } self } diff --git a/compiler/noirc_evaluator/src/ssa/opt/rc.rs b/compiler/noirc_evaluator/src/ssa/opt/rc.rs index 4f109a27874..1750f2d80a5 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/rc.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/rc.rs @@ -2,9 +2,8 @@ use std::collections::{HashMap, HashSet}; use crate::ssa::{ ir::{ - basic_block::BasicBlockId, function::Function, - instruction::{Instruction, InstructionId, TerminatorInstruction}, + instruction::{Instruction, InstructionId}, types::Type, value::ValueId, }, @@ -107,7 +106,7 @@ impl Context { /// Find each dec_rc instruction and if the most recent inc_rc instruction for the same value /// is not possibly mutated, then we can remove them both. Returns each such pair. fn find_rcs_to_remove(&mut self, function: &Function) -> HashSet { - let last_block = Self::find_last_block(function); + let last_block = function.find_last_block(); let mut to_remove = HashSet::new(); for instruction in function.dfg[last_block].instructions() { @@ -124,20 +123,6 @@ impl Context { to_remove } - /// Finds the block of the function with the Return instruction - fn find_last_block(function: &Function) -> BasicBlockId { - for block in function.reachable_blocks() { - if matches!( - function.dfg[block].terminator(), - Some(TerminatorInstruction::Return { .. }) - ) { - return block; - } - } - - unreachable!("SSA Function {} has no reachable return instruction!", function.id()) - } - /// Finds and pops the IncRc for the given array value if possible. fn pop_rc_for(&mut self, value: ValueId, function: &Function) -> Option { let typ = function.dfg.type_of_value(value);