diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs index 007eb349055..eb145ee6dd5 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs @@ -812,17 +812,26 @@ impl DataFlowGraph { /// Uses value information to determine whether an instruction is from /// this function's DFG or the global space's DFG. pub(crate) fn get_local_or_global_instruction(&self, value: ValueId) -> Option<&Instruction> { + self.get_local_or_global_instruction_with_id(value).map(|(instruction, _id)| instruction) + } + + /// Uses value information to determine whether an instruction is from + /// this function's DFG or the global space's DFG, including the instruction ID. + pub(crate) fn get_local_or_global_instruction_with_id( + &self, + value: ValueId, + ) -> Option<(&Instruction, InstructionId)> { match &self[value] { - Value::Instruction { instruction, .. } => { + Value::Instruction { instruction: instruction_id, .. } => { let instruction = if self.is_global(value) { - let instruction = &self.globals[*instruction]; + let instruction = &self.globals[*instruction_id]; // We expect to only have MakeArray instructions in the global space assert!(matches!(instruction, Instruction::MakeArray { .. })); instruction } else { - &self[*instruction] + &self[*instruction_id] }; - Some(instruction) + Some((instruction, *instruction_id)) } _ => None, } diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify.rs index 615f1b50a97..1aec4130fc5 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify.rs @@ -1,12 +1,15 @@ -use crate::ssa::ir::{ - basic_block::BasicBlockId, - dfg::simplify::value_merger::ValueMerger, - instruction::{ - Binary, BinaryOp, ConstrainError, Instruction, - binary::{truncate, truncate_field}, +use crate::ssa::{ + ir::{ + basic_block::BasicBlockId, + dfg::simplify::value_merger::ValueMerger, + instruction::{ + Binary, BinaryOp, ConstrainError, Instruction, + binary::{truncate, truncate_field}, + }, + types::{NumericType, Type}, + value::{Value, ValueId}, }, - types::{NumericType, Type}, - value::{Value, ValueId}, + opt::ArrayGetOptimizationResult, }; use acvm::{ AcirField as _, FieldElement, @@ -115,7 +118,7 @@ pub(crate) fn simplify( Instruction::ConstrainNotEqual(..) => None, Instruction::ArrayGet { array, index } => { if let Some(index) = dfg.get_numeric_constant(*index) { - return try_optimize_array_get_from_previous_set(dfg, *array, index); + return try_optimize_array_get_from_previous_instructions(dfg, *array, index); } let array_or_vector_type = dfg.type_of_value(*array); @@ -357,7 +360,8 @@ fn optimize_length_one_array_read( ); dfg.insert_instruction_and_results(index_constraint, block, None, call_stack); - let result = try_optimize_array_get_from_previous_set(dfg, array, FieldElement::zero()); + let result = + try_optimize_array_get_from_previous_instructions(dfg, array, FieldElement::zero()); if let SimplifyResult::None = result { SimplifyResult::SimplifiedToInstruction(Instruction::ArrayGet { array, index: zero }) } else { @@ -365,73 +369,35 @@ fn optimize_length_one_array_read( } } -/// Given a chain of operations like: -/// v1 = array_set [10, 11, 12], index 1, value: 5 -/// v2 = array_set v1, index 2, value: 6 -/// v3 = array_set v2, index 2, value: 7 -/// v4 = array_get v3, index 1 -/// -/// We want to optimize `v4` to `11`. To do this we need to follow the array value -/// through several array sets. For each array set: -/// - If the index is non-constant we fail the optimization since any index may be changed -/// - If the index is constant and is our target index, we conservatively fail the optimization -/// in case the array_set is disabled from a previous `enable_side_effects_if` and the array get -/// was not. -/// - Otherwise, we check the array value of the array set. -/// - If the array value is constant, we use that array. -/// - If the array value is from a previous array-set, we recur. -/// - If the array value is from an array parameter, we use that array. -/// -/// That is, we have multiple `array_set` instructions setting various constant indexes -/// of the same array, returning a modified version. We want to go backwards until we -/// find the last `array_set` for the index we are interested in, and return the value set. -fn try_optimize_array_get_from_previous_set( +/// See [`crate::ssa::opt::try_optimize_array_get_from_previous_instructions`] for more information. +fn try_optimize_array_get_from_previous_instructions( dfg: &mut DataFlowGraph, - mut array_id: ValueId, + array_id: ValueId, target_index: FieldElement, ) -> SimplifyResult { - // The target index must be less than the maximum array length - let Some(target_index_u32) = target_index.try_to_u32() else { - return SimplifyResult::None; - }; - - // Arbitrary number of maximum tries just to prevent this optimization from taking too long. - let max_tries = 5; - for _ in 0..max_tries { - if let Some(instruction) = dfg.get_local_or_global_instruction(array_id) { - match instruction { - Instruction::ArraySet { array, index, value, .. } => { - if let Some(constant) = dfg.get_numeric_constant(*index) { - if constant == target_index { - return SimplifyResult::SimplifiedTo(*value); - } + let side_effects = None; + let result = crate::ssa::opt::try_optimize_array_get_from_previous_instructions( + array_id, + target_index, + dfg, + side_effects, + ); + match result { + Some(ArrayGetOptimizationResult::Value(value)) => SimplifyResult::SimplifiedTo(value), + Some(ArrayGetOptimizationResult::ArrayGet(new_array_id)) => { + assert_ne!( + new_array_id, array_id, + "ArrayGetOptimizationResult::ArrayGet returned the same array_id" + ); - array_id = *array; // recur - continue; - } - } - Instruction::MakeArray { elements: array, typ: _ } => { - let index = target_index_u32 as usize; - if index < array.len() { - return SimplifyResult::SimplifiedTo(array[index]); - } - } - _ => (), - } - } else if let Value::Param { typ: Type::Array(_, length), .. } = &dfg[array_id] - && target_index_u32 < length.0 - { let index = dfg.make_constant(target_index, NumericType::length_type()); - return SimplifyResult::SimplifiedToInstruction(Instruction::ArrayGet { - array: array_id, + SimplifyResult::SimplifiedToInstruction(Instruction::ArrayGet { + array: new_array_id, index, - }); + }) } - - break; + None => SimplifyResult::None, } - - SimplifyResult::None } /// If we have an array set whose value is from an array get on the same array at the same index, @@ -753,12 +719,17 @@ mod tests { "; let ssa = Ssa::from_str_simplifying(src).unwrap(); + // Only v3 was optimized because v2 reads from the same index + // as v1 and this basic optimization can't know if that array_set + // is under the same predicate as the array_get (there's a dedicated + // `array_get` optimization for that) assert_ssa_snapshot!(ssa, @r" acir(inline) predicate_pure fn main f0 { b0(v0: [Field; 2]): v3 = array_set mut v0, index u32 0, value Field 4 - v5 = array_get v0, index u32 1 -> Field - return Field 4, v5 + v4 = array_get v3, index u32 0 -> Field + v6 = array_get v0, index u32 1 -> Field + return v4, v6 } "); } diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify/call.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify/call.rs index c7f896060c5..206bcdcf512 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify/call.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify/call.rs @@ -551,7 +551,14 @@ fn simplify_vector_push_back( vector_sizes.insert(set_last_vector_value, vector_size / element_size); vector_sizes.insert(new_vector, vector_size / element_size); - let mut value_merger = ValueMerger::new(dfg, block, &vector_sizes, call_stack); + let array_get_optimization_side_effects = None; + let mut value_merger = ValueMerger::new( + dfg, + block, + &vector_sizes, + call_stack, + array_get_optimization_side_effects, + ); let Ok(new_vector) = value_merger.merge_values( len_not_equals_capacity, diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify/value_merger.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify/value_merger.rs index 19a45d8bb22..3bdfca4d631 100644 --- a/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify/value_merger.rs +++ b/compiler/noirc_evaluator/src/ssa/ir/dfg/simplify/value_merger.rs @@ -1,16 +1,25 @@ -use acvm::acir::brillig::lengths::{ElementTypesLength, SemanticLength, SemiFlattenedLength}; +use acvm::{ + FieldElement, + acir::brillig::lengths::{ElementTypesLength, SemanticLength, SemiFlattenedLength}, +}; use noirc_errors::{Location, call_stack::CallStackId}; use rustc_hash::FxHashMap as HashMap; use crate::{ brillig::assert_u32, errors::{RtResult, RuntimeError}, - ssa::ir::{ - basic_block::BasicBlockId, - dfg::DataFlowGraph, - instruction::{BinaryOp, Instruction}, - types::{NumericType, Type}, - value::ValueId, + ssa::{ + ir::{ + basic_block::BasicBlockId, + dfg::DataFlowGraph, + instruction::{BinaryOp, Instruction}, + types::{NumericType, Type}, + value::ValueId, + }, + opt::{ + ArrayGetOptimizationResult, ArrayGetOptimizationSideEffects, + try_optimize_array_get_from_previous_instructions, + }, }, }; @@ -23,6 +32,19 @@ pub(crate) struct ValueMerger<'a> { vector_sizes: &'a HashMap, call_stack: CallStackId, + + /// Optional information about the current side effects variable, and what + /// side effects variables are applied to each `array_set` instruction that happen + /// before the values being merged. + /// + /// When arrays or vectors are merged as a result of an `if ... else` instruction, + /// a new array will be built that does `array_get` on the arrays on the "then" and + /// "else" branches, and combine those values. These newly inserted `array_get` could + /// be optimized by reusing previously inserted instructions, such as an `array_set` at + /// the same index as the one in the `array_get`. However, this is only safe to do + /// if we know the side effects var of those two instructions is the same. Hence, that + /// information can optionally be specified here. + array_get_optimization_side_effects: Option>, } impl<'a> ValueMerger<'a> { @@ -31,8 +53,9 @@ impl<'a> ValueMerger<'a> { block: BasicBlockId, vector_sizes: &'a HashMap, call_stack: CallStackId, + array_get_optimization_side_effects: Option>, ) -> Self { - ValueMerger { dfg, block, vector_sizes, call_stack } + ValueMerger { dfg, block, vector_sizes, call_stack, array_get_optimization_side_effects } } /// Choose a call stack to return with the [RuntimeError]. @@ -147,7 +170,7 @@ impl<'a> ValueMerger<'a> { /// Given an if expression that returns an array: `if c { array1 } else { array2 }`, /// this function will recursively merge array1 and array2 into a single resulting array /// by creating a new array containing the result of `self.merge_values` for each element. - pub(crate) fn merge_array_values( + fn merge_array_values( &mut self, typ: Type, then_condition: ValueId, @@ -166,16 +189,13 @@ impl<'a> ValueMerger<'a> { for i in 0..len.0 { for (element_index, element_type) in element_types.iter().enumerate() { - let index = u128::from(i * element_count + element_index as u32).into(); - let index = self.dfg.make_constant(index, NumericType::length_type()); + let index_value = u128::from(i * element_count + element_index as u32).into(); + let index = self.dfg.make_constant(index_value, NumericType::length_type()); let typevars = Some(vec![element_type.clone()]); let mut get_element = |array, typevars| { - let get = Instruction::ArrayGet { array, index }; - self.dfg - .insert_instruction_and_results(get, self.block, typevars, self.call_stack) - .first() + self.maybe_optimized_array_get(array, index, index_value, typevars) }; let then_element = get_element(then_value, typevars.clone()); @@ -235,14 +255,8 @@ impl<'a> ValueMerger<'a> { let mut get_element = |array, typevars, len: SemiFlattenedLength| { assert!(index_u32 < len.0, "get_element invoked with an out of bounds index"); - let get = Instruction::ArrayGet { array, index }; - let results = self.dfg.insert_instruction_and_results( - get, - self.block, - typevars, - self.call_stack, - ); - results.first() + + self.maybe_optimized_array_get(array, index, index_value, typevars) }; // If it's out of bounds for the "then" vector, a value in the "else" *must* exist. @@ -280,4 +294,39 @@ impl<'a> ValueMerger<'a> { self.dfg.insert_instruction_and_results(instruction, self.block, None, call_stack); Ok(result.first()) } + + fn maybe_optimized_array_get( + &mut self, + array: ValueId, + index: ValueId, + index_value: FieldElement, + typevars: Option>, + ) -> ValueId { + let side_effects = self.array_get_optimization_side_effects.as_ref(); + match try_optimize_array_get_from_previous_instructions( + array, + index_value, + self.dfg, + side_effects, + ) { + Some(ArrayGetOptimizationResult::Value(value)) => value, + Some(ArrayGetOptimizationResult::ArrayGet(new_array)) => { + assert_ne!( + new_array, array, + "ArrayGetOptimizationResult::ArrayGet returned the same array_id" + ); + + let get = Instruction::ArrayGet { array: new_array, index }; + self.dfg + .insert_instruction_and_results(get, self.block, typevars, self.call_stack) + .first() + } + None => { + let get = Instruction::ArrayGet { array, index }; + self.dfg + .insert_instruction_and_results(get, self.block, typevars, self.call_stack) + .first() + } + } + } } diff --git a/compiler/noirc_evaluator/src/ssa/mod.rs b/compiler/noirc_evaluator/src/ssa/mod.rs index b8ef4e459b1..424c71806bc 100644 --- a/compiler/noirc_evaluator/src/ssa/mod.rs +++ b/compiler/noirc_evaluator/src/ssa/mod.rs @@ -130,12 +130,14 @@ pub struct ArtifactsAndWarnings(pub Artifacts, pub Vec); pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec> { vec![ SsaPass::new(Ssa::black_box_bypass, "black_box bypass"), + SsaPass::new(Ssa::array_get_optimization, "ArrayGet optimization"), SsaPass::new(Ssa::expand_signed_checks, "expand signed checks"), SsaPass::new(Ssa::remove_unreachable_functions, "Removing Unreachable Functions"), SsaPass::new(Ssa::defunctionalize, "Defunctionalization"), SsaPass::new_try(Ssa::inline_simple_functions, "Inlining simple functions") .and_then(Ssa::remove_unreachable_functions), SsaPass::new(Ssa::mem2reg, "Mem2Reg"), + SsaPass::new(Ssa::array_get_optimization, "ArrayGet optimization"), SsaPass::new(Ssa::purity_analysis, "Purity Analysis"), SsaPass::new_try( move |ssa| { @@ -146,6 +148,7 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec> { }, "Preprocessing Functions", ), + SsaPass::new(Ssa::array_get_optimization, "ArrayGet optimization"), SsaPass::new_try( move |ssa| { ssa.inline_functions( @@ -157,6 +160,7 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec> { ), // Run mem2reg with the CFG separated into blocks SsaPass::new(Ssa::mem2reg, "Mem2Reg"), + SsaPass::new(Ssa::array_get_optimization, "ArrayGet optimization"), // Running DIE here might remove some unused instructions mem2reg could not eliminate. SsaPass::new( Ssa::dead_instruction_elimination_pre_flattening, @@ -183,6 +187,7 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec> { SsaPass::new(Ssa::simplify_cfg, "Simplifying"), SsaPass::new(Ssa::mem2reg, "Mem2Reg"), SsaPass::new(Ssa::remove_bit_shifts, "Removing Bit Shifts"), + SsaPass::new(Ssa::array_get_optimization, "ArrayGet optimization"), // Expand signed lt/div/mod after "Removing Bit Shifts" because that pass might // introduce signed divisions. SsaPass::new(Ssa::expand_signed_math, "Expand signed math"), @@ -191,6 +196,7 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec> { // Run mem2reg once more with the flattened CFG to catch any remaining loads/stores, // then try to free memory before inlining, which involves copying a instructions. SsaPass::new(Ssa::mem2reg, "Mem2Reg").and_then(Ssa::remove_unused_instructions), + SsaPass::new(Ssa::array_get_optimization, "ArrayGet optimization"), // Run the inlining pass again to handle functions with `InlineType::NoPredicates`. // Before flattening is run, we treat functions marked with the `InlineType::NoPredicates` as an entry point. // This pass must come immediately following `mem2reg` as the succeeding passes @@ -243,6 +249,7 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec> { // We cannot run mem2reg after DIE, because it removes Store instructions. // We have to run it before, to give it a chance to turn Store+Load into known values. SsaPass::new(Ssa::mem2reg, "Mem2Reg"), + SsaPass::new(Ssa::array_get_optimization, "ArrayGet optimization"), SsaPass::new(Ssa::dead_instruction_elimination, "Dead Instruction Elimination"), SsaPass::new(Ssa::brillig_entry_point_analysis, "Brillig Entry Point Analysis") // Remove any potentially unnecessary duplication from the Brillig entry point analysis. diff --git a/compiler/noirc_evaluator/src/ssa/opt/array_get.rs b/compiler/noirc_evaluator/src/ssa/opt/array_get.rs new file mode 100644 index 00000000000..5f0f04c1510 --- /dev/null +++ b/compiler/noirc_evaluator/src/ssa/opt/array_get.rs @@ -0,0 +1,380 @@ +//! Replaces `array_get` with known indices with values from previous instructions +//! such as `array_set` or `make_array`. +//! +//! Given these two instructions: +//! +//! ```text +//! v1 = array_set v0, index 0, value: 42 +//! v2 = array_get v1, index 0 -> Field +//! ``` +//! +//! because we get from `v1` at `index 0`, but `v1` is the result of setting the value "42" +//! at `index 0`, we can conclude that `v2` will be "42", and so this SSA pass will do that. +//! However, this is only safe to do if the `array_set` happened under the same side effects +//! variable as the `array_get`. For example, in this case: +//! +//! ```text +//! enable_side_effects v100 +//! v1 = array_set v0, index 0, value: 42 +//! enable_side_effects v200 +//! v2 = array_get v1, index 0 -> Field +//! ``` +//! +//! it would be wrong to replace `v2` with "42" as the previous array_set might not have +//! been executed. +//! +//! In this case: +//! +//! ```text +//! v1 = array_set v0, index 0, value: 42 +//! v2 = array_set v1, index 1, value: 15 +//! v3 = array_get v2, index 0 -> Field +//! ``` +//! +//! for `v3` the optimization will try to find a previous `array_set` with the same index (`index 0`). +//! It will first find `v2`. Because it's an `array_set` of a different **known** index, it will +//! then find `v1` and apply the same optimization as before. Note that it's safe to skip `v2` and +//! look at `v1` even if `v2` was under a different side effects var because it doesn't affect +//! the index used in `v3`. +//! +//! In this case: +//! +//! ```text +//! v1 = array_set v0, index 0, value: 42 +//! v2 = array_set v1, index v88, value: 15 +//! v3 = array_get v2, index 0 -> Field +//! ``` +//! +//! for `v3` the optimization will find `v2`. Because the set index is unknown, and it might be zero, +//! the optimization can't deduce anything so it won't do anything. +//! +//! Another case where the optimization applies is when it finds a `make_array`: +//! +//! ```text +//! v1 = make_array [Field 10, Field 20] : [Field; 2] +//! v2 = array_get v1, index 0 -> Field +//! ``` +//! +//! In this case `v2` will be replaced with `Field 10`. A `make_array` could also be reached +//! after passing through other `array_set` with a different index, as previously shown. +//! +//! Finally, the optimization might also reach to params: +//! +//! ```text +//! b0(v1: [Field; 2]): +//! v2 = array_set v1, index 1, value: 42 +//! v3 = array_get v2, index 0 -> Field +//! ``` +//! +//! In this case `v3` will be replaced with `array_get v1, index 0`, directly getting from `v1` +//! instead of from `v2`, because `v2` is the same as `v1` except for what's in index 1, but +//! `v3` is getting from index 0. +//! +//! This module also provides a [`try_optimize_array_get_from_previous_instructions`] function +//! that is used in other SSA-related optimizations. For example, whenever an `array_get` is inserted +//! into a [`DFG`][crate::ssa::ir::dfg::DataFlowGraph]: in this case a previous `array_set` with the +//! same index as the `array_get` cannot be used because we don't know under which side effects var it +//! happens. However, `array_set` with a different known index can be skipped through to eventually +//! reach a `make_array` or param. +use std::collections::HashMap; + +use acvm::{AcirField, FieldElement}; + +use crate::ssa::{ + ir::{ + dfg::DataFlowGraph, + function::Function, + instruction::{Instruction, InstructionId}, + types::Type, + value::{Value, ValueId}, + }, + ssa_gen::Ssa, +}; + +impl Ssa { + /// Replaces `array_get` instructions with known indices with known values from + /// previous instructions. See the [`array_get`][self] module for more information. + #[tracing::instrument(level = "trace", skip(self))] + pub(crate) fn array_get_optimization(mut self) -> Self { + for func in self.functions.values_mut() { + func.array_get_optimization(); + } + self + } +} + +impl Function { + fn array_get_optimization(&mut self) { + // Keeps track of side effect vars associated to each `array_set` instruction. + let mut array_set_predicates = HashMap::new(); + + self.simple_optimization(|context| { + let instruction_id = context.instruction_id; + + match context.instruction() { + Instruction::ArraySet { .. } => { + array_set_predicates.insert(instruction_id, context.enable_side_effects); + } + Instruction::ArrayGet { array, index } => { + let Some(index_field) = context.dfg.get_numeric_constant(*index) else { + return; + }; + let array = *array; + let index = *index; + let side_effects = Some(&ArrayGetOptimizationSideEffects { + side_effects_var: context.enable_side_effects, + array_set_predicates: &array_set_predicates, + }); + let Some(result) = try_optimize_array_get_from_previous_instructions( + array, + index_field, + context.dfg, + side_effects, + ) else { + return; + }; + + context.remove_current_instruction(); + + match result { + ArrayGetOptimizationResult::Value(new_value) => { + let [result] = context.dfg.instruction_result(instruction_id); + context.replace_value(result, new_value); + } + ArrayGetOptimizationResult::ArrayGet(new_array) => { + assert_ne!( + new_array, array, + "ArrayGetOptimizationResult::ArrayGet returned the same array_id" + ); + + let array_get = Instruction::ArrayGet { array: new_array, index }; + let [result] = context.dfg.instruction_result(instruction_id); + let result_typ = context.dfg.type_of_value(result); + let ctrl_typevars = Some(vec![result_typ]); + let new_result = context.insert_instruction(array_get, ctrl_typevars); + let new_result = new_result.first(); + context.replace_value(result, new_result); + } + } + } + _ => {} + } + }); + } +} + +/// The result of the array_get optimization. +pub(crate) enum ArrayGetOptimizationResult { + /// The `array_get` can be replaced with the given value. + Value(ValueId), + /// The `array_get` can be replaced by fetching from the given array at the same index as + /// the `array_get`'s index. + ArrayGet(ValueId), +} + +/// Side effects information to be able to optimize `array_get` more efficiently. +pub(crate) struct ArrayGetOptimizationSideEffects<'a> { + /// The current value of the side effects var. + pub(crate) side_effects_var: ValueId, + /// The side effects var applied to each known `array_set` instruction. + pub(crate) array_set_predicates: &'a HashMap, +} + +/// Tries to replace an `array_get` instructions with values from previous instructions. +/// See the [`array_get`][self] module for more information. +pub(crate) fn try_optimize_array_get_from_previous_instructions( + mut array_id: ValueId, + target_index: FieldElement, + dfg: &DataFlowGraph, + side_effects: Option<&ArrayGetOptimizationSideEffects>, +) -> Option { + let original_array_id = array_id; + let target_index_u32 = target_index.try_to_u32()?; + + // Arbitrary number of maximum tries just to prevent this optimization from taking too long. + let max_tries = 5; + for _ in 0..max_tries { + if let Some((instruction, other_instruction_id)) = + dfg.get_local_or_global_instruction_with_id(array_id) + { + match instruction { + Instruction::ArraySet { array, index, value, .. } => { + if let Some(constant) = dfg.get_numeric_constant(*index) { + if constant == target_index { + match side_effects { + None => { + // If it's an array_set with the same index as the array_get, we don't + // use the value at that index. The reason is that the array_set might + // be under a different predicate than the array_get, so the set value + // might not be the correct one in the end. + return None; + } + Some(ArrayGetOptimizationSideEffects { + side_effects_var, + array_set_predicates, + }) => { + // If there's an array_set with the same index as the array_get, we + // can only apply this optimization if they are under the same predicate. + let array_set_predicate = array_set_predicates + .get(&other_instruction_id) + .expect("Expected to know the predicate of every array_set preceding this array_get"); + + if array_set_predicate != side_effects_var { + return None; + } + } + } + + return Some(ArrayGetOptimizationResult::Value(*value)); + } + + // If it's for a different known index, we can safely recur, because + // regardless of whether the array_set ends up being executed or not, it + // won't modify the value at the array_get index. + array_id = *array; + continue; + } + } + Instruction::MakeArray { elements: array, typ: _ } => { + let index = target_index_u32 as usize; + if index < array.len() { + return Some(ArrayGetOptimizationResult::Value(array[index])); + } + } + _ => (), + } + } else if let Value::Param { typ: Type::Array(_, length), .. } = &dfg[array_id] + && target_index_u32 < length.0 + { + // There's no optimization if we end up getting from the original array + if array_id == original_array_id { + return None; + } + + return Some(ArrayGetOptimizationResult::ArrayGet(array_id)); + } + + break; + } + + None +} + +#[cfg(test)] +mod tests { + use crate::{assert_ssa_snapshot, ssa::opt::assert_ssa_does_not_change}; + + use super::Ssa; + + #[test] + fn optimizes_array_get_from_array_set_to_set_value_under_default_predicate() { + let src = " + acir(inline) fn main f0 { + b0(v0: [Field; 3]): + v1 = array_set v0, index u32 0, value Field 1 + v2 = array_get v1, index u32 0 -> Field + return v2 + } + "; + let ssa = Ssa::from_str(src).unwrap(); + + let ssa = ssa.array_get_optimization(); + assert_ssa_snapshot!(ssa, @r" + acir(inline) fn main f0 { + b0(v0: [Field; 3]): + v3 = array_set v0, index u32 0, value Field 1 + return Field 1 + } + "); + } + + #[test] + fn optimizes_array_get_from_array_set_to_array_get_from_param() { + let src = " + acir(inline) fn main f0 { + b0(v0: [Field; 3]): + v1 = array_set v0, index u32 1, value Field 1 + v2 = array_get v1, index u32 0 -> Field + return v2 + } + "; + let ssa = Ssa::from_str(src).unwrap(); + + let ssa = ssa.array_get_optimization(); + assert_ssa_snapshot!(ssa, @r" + acir(inline) fn main f0 { + b0(v0: [Field; 3]): + v3 = array_set v0, index u32 1, value Field 1 + v5 = array_get v0, index u32 0 -> Field + return v5 + } + "); + } + + #[test] + fn optimizes_array_get_from_array_set_to_make_array_value() { + let src = " + acir(inline) fn main f0 { + b0(): + v0 = make_array [Field 2, Field 4, Field 8] : [Field; 3] + v2 = array_get v0, index u32 1 -> Field + return v2 + } + "; + let ssa = Ssa::from_str(src).unwrap(); + + let ssa = ssa.array_get_optimization(); + assert_ssa_snapshot!(ssa, @r" + acir(inline) fn main f0 { + b0(): + v3 = make_array [Field 2, Field 4, Field 8] : [Field; 3] + return Field 4 + } + "); + } + + #[test] + fn does_not_optimize_array_get_from_array_set_with_different_predicate() { + let src = " + acir(inline) fn main f0 { + b0(v0: [Field; 3], v10: u1, v11: u1): + enable_side_effects v10 + v1 = array_set v0, index u32 0, value Field 1 + enable_side_effects v11 + v2 = array_get v1, index u32 0 -> Field + return v2 + } + "; + assert_ssa_does_not_change(src, Ssa::array_get_optimization); + } + + #[test] + fn optimizes_array_get_from_array_set_to_set_value_under_predicate() { + let src = " + acir(inline) fn main f0 { + b0(v0: [Field; 3], v10: u1, v11: u1): + enable_side_effects v10 + v1 = array_set v0, index u32 0, value Field 1 + enable_side_effects v11 + v12 = not v10 + enable_side_effects v10 + v3 = array_get v1, index u32 0 -> Field + return v3 + } + "; + let ssa = Ssa::from_str(src).unwrap(); + + let ssa = ssa.array_get_optimization(); + assert_ssa_snapshot!(ssa, @r" + acir(inline) fn main f0 { + b0(v0: [Field; 3], v1: u1, v2: u1): + enable_side_effects v1 + v5 = array_set v0, index u32 0, value Field 1 + enable_side_effects v2 + v6 = not v1 + enable_side_effects v1 + return Field 1 + } + "); + } +} diff --git a/compiler/noirc_evaluator/src/ssa/opt/mod.rs b/compiler/noirc_evaluator/src/ssa/opt/mod.rs index bdb5bc07867..c18507216ef 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/mod.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/mod.rs @@ -4,6 +4,7 @@ //! simpler form until the IR only has a single function remaining with 1 block within it. //! Generally, these passes are also expected to minimize the final amount of instructions. +mod array_get; mod array_set; mod as_vector_length; mod basic_conditional; @@ -42,6 +43,10 @@ mod simple_optimization; mod simplify_cfg; mod unrolling; +pub(crate) use array_get::{ + ArrayGetOptimizationResult, ArrayGetOptimizationSideEffects, + try_optimize_array_get_from_previous_instructions, +}; pub use constant_folding::DEFAULT_MAX_ITER as CONSTANT_FOLDING_MAX_ITER; pub use inlining::MAX_SIMPLE_FUNCTION_WEIGHT as INLINING_MAX_INSTRUCTIONS; pub use unrolling::FORCE_UNROLL_THRESHOLD; diff --git a/compiler/noirc_evaluator/src/ssa/opt/remove_if_else.rs b/compiler/noirc_evaluator/src/ssa/opt/remove_if_else.rs index 8fd68aa439c..f31e285be3a 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/remove_if_else.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/remove_if_else.rs @@ -108,6 +108,7 @@ use crate::errors::RtResult; use crate::ssa::ir::dfg::simplify::value_merger::ValueMerger; use crate::ssa::ir::types::NumericType; +use crate::ssa::opt::ArrayGetOptimizationSideEffects; use crate::ssa::opt::simple_optimization::SimpleOptimizationContext; use crate::ssa::{ Ssa, @@ -188,6 +189,9 @@ impl Context { return Ok(()); } + // Keeps track of side effect vars associated to each `array_set` instruction. + let mut array_set_predicates = std::collections::HashMap::new(); + function.simple_optimization_result(|context| { let instruction_id = context.instruction_id; let instruction = context.instruction(); @@ -222,8 +226,17 @@ impl Context { } let call_stack = context.dfg.get_instruction_call_stack_id(instruction_id); - let mut value_merger = - ValueMerger::new(context.dfg, block, &self.vector_sizes, call_stack); + let array_get_optimization_data = Some(ArrayGetOptimizationSideEffects { + side_effects_var: context.enable_side_effects, + array_set_predicates: &array_set_predicates, + }); + let mut value_merger = ValueMerger::new( + context.dfg, + block, + &self.vector_sizes, + call_stack, + array_get_optimization_data, + ); let value = value_merger.merge_values( then_condition, @@ -260,6 +273,8 @@ impl Context { } // Track vector sizes through array set instructions Instruction::ArraySet { array, .. } => { + array_set_predicates.insert(instruction_id, context.enable_side_effects); + let [result] = context.dfg.instruction_result(instruction_id); self.set_capacity(context.dfg, *array, result, |c| c); } diff --git a/compiler/noirc_evaluator/src/ssa/opt/simple_optimization.rs b/compiler/noirc_evaluator/src/ssa/opt/simple_optimization.rs index 04ffe2c70f6..e2c305f2db6 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/simple_optimization.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/simple_optimization.rs @@ -64,10 +64,11 @@ impl Function { F: FnMut(&mut SimpleOptimizationContext<'_, '_>) -> RtResult<()>, { let mut values_to_replace = ValueMapping::default(); - let mut enable_side_effects = - self.dfg.make_constant(FieldElement::from(1_u128), NumericType::bool()); + let one = self.dfg.make_constant(FieldElement::from(1_u128), NumericType::bool()); let reverse_post_order = PostOrder::with_function(self).into_vec_reverse(); for block_id in reverse_post_order { + let mut enable_side_effects = one; + let instruction_ids = self.dfg[block_id].take_instructions(); self.dfg[block_id].instructions_mut().reserve(instruction_ids.len()); for instruction_id in &instruction_ids {