feat(performance): brillig array set optimization

compiler/noirc_evaluator/src/acir/arrays.rs

@@ -84,7 +84,7 @@ impl Context<'_> {
         // Pass the instruction between array methods rather than the internal fields themselves
         let (array, index, store_value) = match dfg[instruction] {
             Instruction::ArrayGet { array, index, offset: _ } => (array, index, None),
-            Instruction::ArraySet { array, index, value, mutable } => {
+            Instruction::ArraySet { array, index, value, mutable, offset: _ } => {
                 mutable_array_set = mutable;
                 (array, index, Some(value))
             }

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs

@@ -7,7 +7,7 @@ use crate::brillig::brillig_ir::registers::RegisterAllocator;
 use crate::brillig::brillig_ir::{
     BRILLIG_MEMORY_ADDRESSING_BIT_SIZE, BrilligBinaryOp, BrilligContext, ReservedRegisters,
 };
-use crate::ssa::ir::instruction::{ArrayGetOffset, ConstrainError, Hint};
+use crate::ssa::ir::instruction::{ArrayOffset, ConstrainError, Hint};
 use crate::ssa::ir::{
     basic_block::BasicBlockId,
     dfg::DataFlowGraph,
@@ -802,30 +802,24 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> {
                 );
 
                 let array_variable = self.convert_ssa_value(*array, dfg);
-
                 let index_variable = self.convert_ssa_single_addr_value(*index, dfg);
 
-                if dfg.is_constant(*index) {
-                    // For constant indices it must be the case that they have been offseted during SSA
-                    assert!(*offset != ArrayGetOffset::None);
-                    self.brillig_context.codegen_load_with_offset(
-                        array_variable.extract_register(),
-                        index_variable,
-                        destination_variable.extract_register(),
-                    );
+                let has_offset = if dfg.is_constant(*index) {
+                    // For constant indices it must be the case that they have been offset during SSA
+                    assert!(*offset != ArrayOffset::None);
+                    true
                 } else {
-                    let items_pointer = self
-                        .brillig_context
-                        .codegen_make_array_or_vector_items_pointer(array_variable);
-                    self.brillig_context.codegen_load_with_offset(
-                        items_pointer,
-                        index_variable,
-                        destination_variable.extract_register(),
-                    );
-                    self.brillig_context.deallocate_register(items_pointer);
-                }
+                    false
+                };
+
+                self.convert_ssa_array_get(
+                    array_variable,
+                    index_variable,
+                    destination_variable,
+                    has_offset,
+                );
             }
-            Instruction::ArraySet { array, index, value, mutable } => {
+            Instruction::ArraySet { array, index, value, mutable, offset } => {
                 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);
@@ -838,12 +832,21 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> {
                     dfg,
                 );
 
+                let has_offset = if dfg.is_constant(*index) {
+                    // For constant indices it must be the case that they have been offset during SSA
+                    assert!(*offset != ArrayOffset::None);
+                    true
+                } else {
+                    false
+                };
+
                 self.convert_ssa_array_set(
                     source_variable,
                     destination_variable,
                     index_register,
                     value_variable,
                     *mutable,
+                    has_offset,
                 );
             }
             Instruction::RangeCheck { value, max_bit_size, assert_message } => {
@@ -1110,6 +1113,30 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> {
         }
     }
 
+    fn convert_ssa_array_get(
+        &mut self,
+        array_variable: BrilligVariable,
+        index_variable: SingleAddrVariable,
+        destination_variable: BrilligVariable,
+        has_offset: bool,
+    ) {
+        let items_pointer = if has_offset {
+            array_variable.extract_register()
+        } else {
+            self.brillig_context.codegen_make_array_or_vector_items_pointer(array_variable)
+        };
+
+        self.brillig_context.codegen_load_with_offset(
+            items_pointer,
+            index_variable,
+            destination_variable.extract_register(),
+        );
+
+        if !has_offset {
+            self.brillig_context.deallocate_register(items_pointer);
+        }
+    }
+
     /// Array set operation in SSA returns a new array or slice that is a copy of the parameter array or slice
     /// With a specific value changed.
     ///
@@ -1122,6 +1149,7 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> {
         index_register: SingleAddrVariable,
         value_variable: BrilligVariable,
         mutable: bool,
+        has_offset: bool,
     ) {
         assert!(index_register.bit_size == BRILLIG_MEMORY_ADDRESSING_BIT_SIZE);
         match (source_variable, destination_variable) {
@@ -1146,9 +1174,13 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> {
         }
 
         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_for_store);
+        let items_pointer = if has_offset {
+            destination_for_store.extract_register()
+        } else {
+            self.brillig_context.codegen_make_array_or_vector_items_pointer(destination_for_store)
+        };
 
         self.brillig_context.codegen_store_with_offset(
             items_pointer,
@@ -1164,7 +1196,9 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> {
             );
         }
 
-        self.brillig_context.deallocate_register(items_pointer);
+        if !has_offset {
+            self.brillig_context.deallocate_register(items_pointer);
+        }
     }
 
     /// Convert the SSA slice operations to brillig slice operations

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_globals.rs

@@ -442,7 +442,7 @@ mod tests {
 
         let ssa = Ssa::from_str(src).unwrap();
         // Need to run SSA pass that sets up Brillig array gets
-        let ssa = ssa.brillig_array_gets();
+        let ssa = ssa.brillig_array_get_and_set();
         // Need to run DIE to generate the used globals map, which is necessary for Brillig globals generation.
         let mut ssa = ssa.dead_instruction_elimination();
 

compiler/noirc_evaluator/src/ssa.rs

@@ -210,7 +210,7 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec<SsaPass> {
         // We can safely place the pass before DIE as that pass only removes instructions.
         // We also need DIE's tracking of used globals in case the array get transformations
         // end up using an existing constant from the globals space.
-        SsaPass::new(Ssa::brillig_array_gets, "Brillig Array Get Optimizations"),
+        SsaPass::new(Ssa::brillig_array_get_and_set, "Brillig Array Get and Set Optimizations"),
         SsaPass::new(Ssa::dead_instruction_elimination, "Dead Instruction Elimination"),
         // A function can be potentially unreachable post-DIE if all calls to that function were removed.
         SsaPass::new(Ssa::remove_unreachable_functions, "Removing Unreachable Functions"),
@@ -248,7 +248,7 @@ pub fn minimal_passes() -> Vec<SsaPass<'static>> {
         // We need a DIE pass to populate `used_globals`, otherwise it will panic later.
         SsaPass::new(Ssa::dead_instruction_elimination, "Dead Instruction Elimination"),
         // We need to add an offset to constant array indices in Brillig.
-        SsaPass::new(Ssa::brillig_array_gets, "Brillig Array Get Optimizations"),
+        SsaPass::new(Ssa::brillig_array_get_and_set, "Brillig Array Get and Set Optimizations"),
     ]
 }
 

compiler/noirc_evaluator/src/ssa/function_builder/data_bus.rs

@@ -2,6 +2,7 @@ use std::{collections::BTreeMap, sync::Arc};
 
 use crate::ssa::ir::{
     function::RuntimeType,
+    instruction::ArrayOffset,
     types::{NumericType, Type},
     value::{ValueId, ValueMapping},
 };
@@ -161,7 +162,9 @@ impl FunctionBuilder {
                         if let Type::Array(_, 0) = subitem_typ {
                             continue;
                         }
-                        let element = self.insert_array_get(value, index_var, subitem_typ.clone());
+                        let offset = ArrayOffset::None;
+                        let element =
+                            self.insert_array_get(value, index_var, offset, subitem_typ.clone());
                         index += match subitem_typ {
                             Type::Array(_, _) | Type::Slice(_) => subitem_typ.element_size(),
                             Type::Numeric(_) => 1,

compiler/noirc_evaluator/src/ssa/function_builder/mod.rs

@@ -23,7 +23,7 @@ use super::{
         basic_block::BasicBlock,
         dfg::{GlobalsGraph, InsertInstructionResult},
         function::RuntimeType,
-        instruction::{ArrayGetOffset, ConstrainError, InstructionId, Intrinsic},
+        instruction::{ArrayOffset, ConstrainError, InstructionId, Intrinsic},
         types::NumericType,
     },
     opt::pure::FunctionPurities,
@@ -351,17 +351,7 @@ impl FunctionBuilder {
         &mut self,
         array: ValueId,
         index: ValueId,
-        element_type: Type,
-    ) -> ValueId {
-        self.insert_array_get_with_offset(array, index, ArrayGetOffset::None, element_type)
-    }
-
-    /// Insert an instruction to extract an element from an array
-    pub fn insert_array_get_with_offset(
-        &mut self,
-        array: ValueId,
-        index: ValueId,
-        offset: ArrayGetOffset,
+        offset: ArrayOffset,
         element_type: Type,
     ) -> ValueId {
         let element_type = Some(vec![element_type]);
@@ -370,20 +360,16 @@ impl FunctionBuilder {
     }
 
     /// Insert an instruction to create a new array with the given index replaced with a new value
-    pub fn insert_array_set(&mut self, array: ValueId, index: ValueId, value: ValueId) -> ValueId {
-        self.insert_instruction(Instruction::ArraySet { array, index, value, mutable: false }, None)
-            .first()
-    }
-
-    #[cfg(test)]
-    pub fn insert_mutable_array_set(
+    pub fn insert_array_set(
         &mut self,
         array: ValueId,
         index: ValueId,
         value: ValueId,
+        mutable: bool,
+        offset: ArrayOffset,
     ) -> ValueId {
-        self.insert_instruction(Instruction::ArraySet { array, index, value, mutable: true }, None)
-            .first()
+        let instruction = Instruction::ArraySet { array, index, value, mutable, offset };
+        self.insert_instruction(instruction, None).first()
     }
 
     /// Insert an instruction to increment an array's reference count. This only has an effect

compiler/noirc_evaluator/src/ssa/interpreter/mod.rs

@@ -5,7 +5,7 @@
     ir::{
         dfg::DataFlowGraph,
         function::{Function, FunctionId, RuntimeType},
-        instruction::{ArrayGetOffset, Binary, BinaryOp, Instruction, TerminatorInstruction},
+        instruction::{ArrayOffset, Binary, BinaryOp, Instruction, TerminatorInstruction},
         types::Type,
         value::ValueId,
     },
@@ -446,8 +446,8 @@
             Instruction::ArrayGet { array, index, offset } => {
                 self.interpret_array_get(*array, *index, *offset, results[0])
             }
-            Instruction::ArraySet { array, index, value, mutable } => {
-                self.interpret_array_set(*array, *index, *value, *mutable, results[0])
+            Instruction::ArraySet { array, index, value, mutable, offset } => {
+                self.interpret_array_set(*array, *index, *value, *mutable, *offset, results[0])
             }
             Instruction::IncrementRc { value } => self.interpret_inc_rc(*value),
             Instruction::DecrementRc { value } => self.interpret_dec_rc(*value),
@@ -736,7 +736,7 @@
         &mut self,
         array: ValueId,
         index: ValueId,
-        offset: ArrayGetOffset,
+        offset: ArrayOffset,
         result: ValueId,
     ) -> IResult<()> {
         let element = if self.side_effects_enabled() {
@@ -758,12 +758,14 @@
         index: ValueId,
         value: ValueId,
         mutable: bool,
+        offset: ArrayOffset,
         result: ValueId,
     ) -> IResult<()> {
         let array = self.lookup_array_or_slice(array, "array set")?;
 
         let result_array = if self.side_effects_enabled() {
             let index = self.lookup_u32(index, "array set index")?;
+            let index = index - offset.to_u32();
             let value = self.lookup(value)?;
 
             let should_mutate =
@@ -875,7 +877,7 @@
     }
 }
 
 macro_rules! apply_int_binop {
     ($lhs:expr, $rhs:expr, $binary:expr, $f:expr) => {{
         use value::NumericValue::*;
         match ($lhs, $rhs) {
@@ -906,7 +908,7 @@
     }};
 }
 
 macro_rules! apply_int_binop_opt {
     ($dfg:expr, $lhs:expr, $rhs:expr, $binary:expr, $f:expr) => {{
         use value::NumericValue::*;
 
@@ -1006,7 +1008,7 @@
             }));
         }
 
         // Disable this instruction if it is side-effectful and side effects are disabled.
         if !self.side_effects_enabled() && binary.requires_acir_gen_predicate(self.dfg()) {
             let zero = NumericValue::zero(lhs.get_type());
             return Ok(Value::Numeric(zero));
@@ -1023,13 +1025,13 @@
         let dfg = self.dfg();
         let result = match binary.operator {
             BinaryOp::Add { unchecked: false } => {
                 apply_int_binop_opt!(dfg, lhs, rhs, binary, num_traits::CheckedAdd::checked_add)
             }
             BinaryOp::Add { unchecked: true } => {
                 apply_int_binop!(lhs, rhs, binary, num_traits::WrappingAdd::wrapping_add)
             }
             BinaryOp::Sub { unchecked: false } => {
                 apply_int_binop_opt!(dfg, lhs, rhs, binary, num_traits::CheckedSub::checked_sub)
             }
             BinaryOp::Sub { unchecked: true } => {
                 apply_int_binop!(lhs, rhs, binary, num_traits::WrappingSub::wrapping_sub)

compiler/noirc_evaluator/src/ssa/interpreter/tests/instructions.rs

@@ -725,6 +725,36 @@ fn array_set_disabled_by_enable_side_effects() {
     assert_eq!(*v2.elements.borrow(), expected);
 }
 
+#[test]
+fn array_set_with_offset() {
+    let values = expect_values(
+        "
+        acir(inline) fn main f0 {
+          b0():
+            v0 = make_array [Field 1, Field 2] : [Field; 2]
+            v1 = array_set v0, index u32 4 minus 3, value Field 5
+            return v0, v1
+        }
+    ",
+    );
+
+    let v0 = values[0].as_array_or_slice().unwrap();
+    let v1 = values[1].as_array_or_slice().unwrap();
+
+    // acir function, so all rcs are 1
+    assert_eq!(*v0.rc.borrow(), 1);
+    assert_eq!(*v1.rc.borrow(), 1);
+
+    let one = from_constant(1u32.into(), NumericType::NativeField);
+    let two = from_constant(2u32.into(), NumericType::NativeField);
+    let five = from_constant(5u32.into(), NumericType::NativeField);
+
+    // v0 was not mutated
+    assert_eq!(*v0.elements.borrow(), vec![one.clone(), two.clone()]);
+    // v1 was mutated
+    assert_eq!(*v1.elements.borrow(), vec![one, five]);
+}
+
 #[test]
 fn increment_rc() {
     let value = expect_value(

compiler/noirc_evaluator/src/ssa/ir/dfg/simplify.rs

@@ -2,7 +2,7 @@ use crate::ssa::{
     ir::{
         basic_block::BasicBlockId,
         instruction::{
-            ArrayGetOffset, Binary, BinaryOp, Instruction,
+            ArrayOffset, Binary, BinaryOp, Instruction,
             binary::{truncate, truncate_field},
         },
         types::Type,
@@ -376,7 +376,7 @@ fn try_optimize_array_set_from_previous_get(
 ) -> SimplifyResult {
     let array_from_get = match dfg.get_local_or_global_instruction(target_value) {
         Some(Instruction::ArrayGet { array, index, offset }) => {
-            if *offset == ArrayGetOffset::None && *array == array_id && *index == target_index {
+            if *offset == ArrayOffset::None && *array == array_id && *index == target_index {
                 // If array and index match from the value, we can immediately simplify
                 return SimplifyResult::SimplifiedTo(array_id);
             } else if *index == target_index {

compiler/noirc_evaluator/src/ssa/ir/dfg/simplify/call.rs

@@ -11,7 +11,7 @@ use crate::ssa::{
     ir::{
         basic_block::BasicBlockId,
         dfg::DataFlowGraph,
-        instruction::{ArrayGetOffset, Binary, BinaryOp, Endian, Hint, Instruction, Intrinsic},
+        instruction::{ArrayOffset, Binary, BinaryOp, Endian, Hint, Instruction, Intrinsic},
         types::{NumericType, Type},
         value::{Value, ValueId},
     },
@@ -506,6 +506,7 @@ fn simplify_slice_push_back(
         index: arguments[0],
         value: arguments[2],
         mutable: false,
+        offset: ArrayOffset::None,
     };
 
     let set_last_slice_value = dfg
@@ -557,7 +558,7 @@ fn simplify_slice_pop_back(
         let get_last_elem_instr = Instruction::ArrayGet {
             array: arguments[1],
             index: flattened_len,
-            offset: ArrayGetOffset::None,
+            offset: ArrayOffset::None,
         };
 
         let element_type = Some(vec![element_type.clone()]);