diff --git a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs index 58ba498e68b..0a8f05c709c 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/array_set.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/array_set.rs @@ -2,6 +2,9 @@ //! set a known index on a previous `make_array` instruction, as long as both instructions //! are under the same side-effects predicate or the side effects var for the `array_set` is `true`. //! +//! If the predicates are different, but the element type is numeric, it will use the `ValueMerger` +//! to merge the items. +//! //! For example, this: //! //! ```text @@ -20,9 +23,36 @@ //! //! ```text //! enable_side_effects v0 -//! v0 = make_array [Field 2, Field 3] : [Field; 2] -//! enable_side_effects v1 -//! v1 = array_set v0, index u32 0, value Field 4 +//! v1 = make_array [Field 2, Field 3] : [Field; 2] +//! enable_side_effects v2 +//! v3 = array_set v0, index u32 0, value Field 4 +//! ``` +//! +//! will change into this: +//! +//! ```text +//! enable_side_effects v0 +//! v1 = make_array [Field 2, Field 3] : [Field; 2] +//! enable_side_effects v2 +//! v3 = not v2 +//! v4 = cast v2 as Field +//! v5 = cast v3 as Field +//! v6 = mul v4, Field 4 +//! v7 = mul v5, Field 2 +//! v8 = add v6, v7 +//! v9 = make_array [v8, Field 3] : [Field; 2] +//! ``` +//! +//! and this: +//! +//! ```text +//! enable_side_effects v0 +//! v1 = make_array [Field 1, Field 2] : [Field; 2] +//! v2 = make_array [Field 3, Field 4] : [Field; 2] +//! v3 = make_array [v1, v2] : [[Field; 2]; 2] +//! enable_side_effects v4 +//! v5 = make_array [Field 5, Field 6] : [Field; 2] +//! v6 = array_set v3, index u32 0, value v5 //! ``` //! //! will remain unchanged. @@ -30,10 +60,12 @@ use std::collections::HashMap; use acvm::AcirField; use im::Vector; +use noirc_errors::call_stack::CallStackId; use crate::ssa::{ ir::{ - dfg::DataFlowGraph, + basic_block::BasicBlockId, + dfg::{DataFlowGraph, simplify::value_merger::ValueMerger}, function::Function, instruction::{Instruction, InstructionId}, value::{Value, ValueId}, @@ -92,8 +124,10 @@ impl Function { let array = *array; let value = *value; - if let Some(elements) = fold_array_set_into_make_array( + if let Some((elements, insert_predicate)) = fold_array_set_into_make_array( context.dfg, + context.block_id, + context.call_stack_id, array, value, index, @@ -111,10 +145,12 @@ impl Function { context.replace_value(result, new_result); // Keep track of the predicate of the newly inserted make_array instruction - let Value::Instruction { instruction: new_instruction_id, .. } = context.dfg[new_result] else { - unreachable!("Expected the previous make_array insertion to be an instruction value"); - }; - make_array_predicates.insert(new_instruction_id, context.enable_side_effects); + if insert_predicate { + let Value::Instruction { instruction: new_instruction_id, .. } = context.dfg[new_result] else { + unreachable!("Expected the last make_array insertion to be an instruction value"); + }; + make_array_predicates.insert(new_instruction_id, context.enable_side_effects); + } } } _ => {} @@ -123,14 +159,24 @@ impl Function { } } +/// Decide whether we can turn an `array_set` into a `make_array`, returning: +/// * `None` to keep the `array_set` +/// * `Some(elements, insert_predicate)` to replace it with a `make_array` with the updated `elements`; +/// +/// If `insert_predicate` is true then we should keep tracking the side effect variable of the new `make_array`. +/// Otherwise the `elements` are a result of merging the items at the `index` under different predicates, +/// and the items in the `make_array` can be a mix of various side effects, and tracking must be stopped for it. +#[allow(clippy::too_many_arguments)] fn fold_array_set_into_make_array( - dfg: &DataFlowGraph, + dfg: &mut DataFlowGraph, + block_id: BasicBlockId, + call_stack_id: CallStackId, array_id: ValueId, value: ValueId, index: u32, side_effects_var: ValueId, make_array_predicates: &HashMap, -) -> Option> { +) -> Option<(Vector, bool)> { let index = index as usize; let (instruction, instruction_id) = dfg.get_local_or_global_instruction_with_id(array_id)?; @@ -144,23 +190,57 @@ fn fold_array_set_into_make_array( } let side_effects_var_value = dfg.get_numeric_constant(side_effects_var); + let always_executes = side_effects_var_value.is_some_and(|var| var.is_one()); + let never_executes = side_effects_var_value.is_some_and(|var| var.is_zero()); - // If the current side effects var is `u1 0`, the `array_set` will never execute. In that case we - // can make its return value be the original `make_array` it's (not) modifying. - if side_effects_var_value.is_some_and(|var| var.is_zero()) { - return Some(elements.clone()); + // If the current side effects var is `u1 0`, the `array_set` will never execute. + // In that case we can make its return value be the original `make_array` it's (not) modifying. + if never_executes { + return Some((elements.clone(), true)); } - let can_fold = side_effects_var_value.is_some_and(|var| var.is_one()) - || make_array_predicates[&instruction_id] == side_effects_var; - if !can_fold { - // The array_set and make_array are under different predicates, and the array_set predicate is not `true`, - // so we can't fold them together. + // If the current side effects var is `u1 1`, then the `array_set` will always execute, and we can safely replace the element. + // This is also true if the `make_array` was under the same predicate as the `array_set`, however we might not have this + // information, if the `make_array` is the result of merging an earlier `make_array` with an `array_set` below; + // in that case different items may be under different predicates, and we must keep using merge. + let can_fold = always_executes + || make_array_predicates.get(&instruction_id).is_some_and(|p| *p == side_effects_var); + + if can_fold { + return Some((elements.update(index, value), true)); + } + + // If we are dealing with a simple numeric value, we can merge it, so the new value will be: + // elements[index] = side_effects * value + (1 - side_effects) * elements[index] + if !dfg.type_of_value(value).is_numeric() { return None; } - let elements = elements.update(index, value); - Some(elements) + // The array_set and make_array are under different predicates, and the array_set predicate is not `true`, + // so we can't fold them together. We can either keep the array_set, or merge the items at that index. + let mut elements = elements.clone(); + + let negated_side_effects_var = dfg + .insert_instruction_and_results( + Instruction::Not(side_effects_var), + block_id, + None, + call_stack_id, + ) + .first(); + + let merged_value = ValueMerger::merge_numeric_values( + dfg, + block_id, + side_effects_var, + negated_side_effects_var, + value, + elements[index], + ); + + elements[index] = merged_value; + + Some((elements, false)) } #[cfg(test)] @@ -219,23 +299,68 @@ mod tests { "); } - /// ArraySet on a constant array must not be folded into MakeArray when the + /// ArraySet on a constant array must use merging with the MakeArray when the /// side-effects predicate is different for both instructions, because the array_set may /// not actually execute. #[test] - fn does_not_fold_array_set_when_side_effects_predicate_is_unknown() { + fn merge_folds_array_set_chain_when_side_effects_predicate_is_unknown() { let src = " acir(inline) fn main f0 { b0(v0: u1): v1 = make_array [Field 10, Field 11] : [Field; 2] enable_side_effects v0 v2 = array_set v1, index u32 0, value Field 99 + v3 = array_set v2, index u32 0, value Field 100 enable_side_effects u1 1 - v3 = array_get v2, index u32 0 -> Field - return v3 + v4 = array_get v3, index u32 0 -> Field + return v4 } "; + let ssa = Ssa::from_str(src).unwrap(); + let ssa = ssa.array_set_optimization(); + + assert_ssa_snapshot!(ssa, @r" + acir(inline) fn main f0 { + b0(v0: u1): + v3 = make_array [Field 10, Field 11] : [Field; 2] + enable_side_effects v0 + v4 = not v0 + v5 = cast v0 as Field + v6 = cast v4 as Field + v8 = mul v5, Field 99 + v9 = mul v6, Field 10 + v10 = add v8, v9 + v11 = make_array [v10, Field 11] : [Field; 2] + v12 = not v0 + v13 = cast v0 as Field + v14 = cast v12 as Field + v16 = mul v13, Field 100 + v17 = mul v14, v10 + v18 = add v16, v17 + v19 = make_array [v18, Field 11] : [Field; 2] + enable_side_effects u1 1 + return v18 + } + "); + } + + #[test] + fn does_not_fold_array_set_on_complex_array() { + let src = r#" + acir(inline) fn main f0 { + b0(v0: u1, v4: u1): + enable_side_effects v0 + v1 = make_array [Field 1, Field 2] : [Field; 2] + v2 = make_array [Field 3, Field 4] : [Field; 2] + v3 = make_array [v1, v2] : [[Field; 2]; 2] + enable_side_effects v4 + v5 = make_array [Field 5, Field 6] : [Field; 2] + v6 = array_set v3, index u32 0, value v5 + return v6 + } + "#; + assert_ssa_does_not_change(src, Ssa::array_set_optimization); }