fix: Last use analysis & make it an SSA pass (#4686)

# Description ## Problem\* Resolves #4685 ## Summary\* Fixes some issues with the last_use analysis and makes it into a dedicated SSA pass (after all other passes, only on 1 block acir functions). ## Additional Context One case I noticed where we don't make array sets mutable currently is in parallel if branches: ```rs if array[0] == 1 { array[1] = 10; } else { array[1] = 20; } ``` Will produce the IR: ```rs v137 = array_get v0, index u64 0 v138 = eq v137, Field 1 enable_side_effects v138 v139 = array_set v0, index u64 1, value Field 10 v140 = not v138 enable_side_effects v140 v141 = array_set mut v0, index u64 1, value Field 20 enable_side_effects u1 1 ``` Note that with this change we can see the second array_set is mutable and not the first. Fixing this isn't in the scope of this PR but this could be a good future optimization. ## 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.
noir-lang · Apr 3, 2024 · 941eb5e · 941eb5e
1 parent 8adee4f
commit 941eb5e
Show file tree

Hide file tree

Showing 9 changed files with 154 additions and 125 deletions.
diff --git a/compiler/noirc_evaluator/src/ssa.rs b/compiler/noirc_evaluator/src/ssa.rs
@@ -63,17 +63,14 @@ pub(crate) fn optimize_into_acir(
         .run_pass(Ssa::fold_constants, "After Constant Folding:")
         .run_pass(Ssa::fold_constants_using_constraints, "After Constraint Folding:")
         .run_pass(Ssa::dead_instruction_elimination, "After Dead Instruction Elimination:")
+        .run_pass(Ssa::array_set_optimization, "After Array Set Optimizations:")
         .finish();
 
     let brillig = time("SSA to Brillig", print_timings, || ssa.to_brillig(print_brillig_trace));
 
     drop(ssa_gen_span_guard);
 
-    let last_array_uses = ssa.find_last_array_uses();
-
-    time("SSA to ACIR", print_timings, || {
-        ssa.into_acir(&brillig, abi_distinctness, &last_array_uses)
-    })
+    time("SSA to ACIR", print_timings, || ssa.into_acir(&brillig, abi_distinctness))
 }
 
 // Helper to time SSA passes

diff --git a/compiler/noirc_evaluator/src/ssa/acir_gen/mod.rs b/compiler/noirc_evaluator/src/ssa/acir_gen/mod.rs
@@ -181,15 +181,12 @@ impl Ssa {
         self,
         brillig: &Brillig,
         abi_distinctness: Distinctness,
-        last_array_uses: &HashMap<ValueId, InstructionId>,
     ) -> Result<Vec<GeneratedAcir>, RuntimeError> {
         let mut acirs = Vec::new();
         // TODO: can we parallelise this?
         for function in self.functions.values() {
             let context = Context::new();
-            if let Some(generated_acir) =
-                context.convert_ssa_function(&self, function, brillig, last_array_uses)?
-            {
+            if let Some(generated_acir) = context.convert_ssa_function(&self, function, brillig)? {
                 acirs.push(generated_acir);
             }
         }
@@ -245,7 +242,6 @@ impl Context {
         ssa: &Ssa,
         function: &Function,
         brillig: &Brillig,
-        last_array_uses: &HashMap<ValueId, InstructionId>,
     ) -> Result<Option<GeneratedAcir>, RuntimeError> {
         match function.runtime() {
             RuntimeType::Acir(inline_type) => {
@@ -258,7 +254,7 @@ impl Context {
                     }
                 }
                 // We only want to convert entry point functions. This being `main` and those marked with `#[fold]`
-                Ok(Some(self.convert_acir_main(function, ssa, brillig, last_array_uses)?))
+                Ok(Some(self.convert_acir_main(function, ssa, brillig)?))
             }
             RuntimeType::Brillig => {
                 if function.id() == ssa.main_id {
@@ -275,7 +271,6 @@ impl Context {
         main_func: &Function,
         ssa: &Ssa,
         brillig: &Brillig,
-        last_array_uses: &HashMap<ValueId, InstructionId>,
     ) -> Result<GeneratedAcir, RuntimeError> {
         let dfg = &main_func.dfg;
         let entry_block = &dfg[main_func.entry_block()];
@@ -284,13 +279,7 @@ impl Context {
         self.data_bus = dfg.data_bus.to_owned();
         let mut warnings = Vec::new();
         for instruction_id in entry_block.instructions() {
-            warnings.extend(self.convert_ssa_instruction(
-                *instruction_id,
-                dfg,
-                ssa,
-                brillig,
-                last_array_uses,
-            )?);
+            warnings.extend(self.convert_ssa_instruction(*instruction_id, dfg, ssa, brillig)?);
         }
 
         warnings.extend(self.convert_ssa_return(entry_block.unwrap_terminator(), dfg)?);
@@ -463,7 +452,6 @@ impl Context {
         dfg: &DataFlowGraph,
         ssa: &Ssa,
         brillig: &Brillig,
-        last_array_uses: &HashMap<ValueId, InstructionId>,
     ) -> Result<Vec<SsaReport>, RuntimeError> {
         let instruction = &dfg[instruction_id];
         self.acir_context.set_call_stack(dfg.get_call_stack(instruction_id));
@@ -523,7 +511,7 @@ impl Context {
                 self.current_side_effects_enabled_var = acir_var;
             }
             Instruction::ArrayGet { .. } | Instruction::ArraySet { .. } => {
-                self.handle_array_operation(instruction_id, dfg, last_array_uses)?;
+                self.handle_array_operation(instruction_id, dfg)?;
             }
             Instruction::Allocate => {
                 unreachable!("Expected all allocate instructions to be removed before acir_gen")
@@ -721,12 +709,16 @@ impl Context {
         &mut self,
         instruction: InstructionId,
         dfg: &DataFlowGraph,
-        last_array_uses: &HashMap<ValueId, InstructionId>,
     ) -> Result<(), RuntimeError> {
+        let mut mutable_array_set = false;
+
         // Pass the instruction between array methods rather than the internal fields themselves
         let (array, index, store_value) = match dfg[instruction] {
             Instruction::ArrayGet { array, index } => (array, index, None),
-            Instruction::ArraySet { array, index, value, .. } => (array, index, Some(value)),
+            Instruction::ArraySet { array, index, value, mutable } => {
+                mutable_array_set = mutable;
+                (array, index, Some(value))
+            }
             _ => {
                 return Err(InternalError::Unexpected {
                     expected: "Instruction should be an ArrayGet or ArraySet".to_owned(),
@@ -744,11 +736,8 @@ impl Context {
         let (new_index, new_value) =
             self.convert_array_operation_inputs(array, dfg, index, store_value)?;
 
-        let resolved_array = dfg.resolve(array);
-        let map_array = last_array_uses.get(&resolved_array) == Some(&instruction);
-
         if let Some(new_value) = new_value {
-            self.array_set(instruction, new_index, new_value, dfg, map_array)?;
+            self.array_set(instruction, new_index, new_value, dfg, mutable_array_set)?;
         } else {
             self.array_get(instruction, array, new_index, dfg)?;
         }
@@ -1028,16 +1017,18 @@ impl Context {
         }
     }
 
-    /// Copy the array and generates a write opcode on the new array
-    ///
-    /// Note: Copying the array is inefficient and is not the way we want to do it in the end.
+    /// If `mutate_array` is:
+    /// - true: Mutate the array directly
+    /// - false: Copy the array and generates a write opcode on the new array. This is
+    ///          generally very inefficient and should be avoided if possible. Currently
+    ///          this is controlled by SSA's array set optimization pass.
     fn array_set(
         &mut self,
         instruction: InstructionId,
         mut var_index: AcirVar,
         store_value: AcirValue,
         dfg: &DataFlowGraph,
-        map_array: bool,
+        mutate_array: bool,
     ) -> Result<(), RuntimeError> {
         // Pass the instruction between array methods rather than the internal fields themselves
         let array = match dfg[instruction] {
@@ -1075,7 +1066,7 @@ impl Context {
             .first()
             .expect("Array set does not have one result");
         let result_block_id;
-        if map_array {
+        if mutate_array {
             self.memory_blocks.insert(*result_id, block_id);
             result_block_id = block_id;
         } else {
@@ -2401,14 +2392,13 @@ mod test {
         ssa::{
             function_builder::FunctionBuilder,
             ir::{
-                function::{FunctionId, InlineType, RuntimeType},
+                function::{FunctionId, InlineType},
                 instruction::BinaryOp,
                 map::Id,
                 types::Type,
             },
         },
     };
-    use fxhash::FxHashMap as HashMap;
 
     fn build_basic_foo_with_return(builder: &mut FunctionBuilder, foo_id: FunctionId) {
         // acir(fold) fn foo f1 {
@@ -2461,11 +2451,7 @@ mod test {
         let ssa = builder.finish();
 
         let acir_functions = ssa
-            .into_acir(
-                &Brillig::default(),
-                noirc_frontend::Distinctness::Distinct,
-                &HashMap::default(),
-            )
+            .into_acir(&Brillig::default(), noirc_frontend::Distinctness::Distinct)
             .expect("Should compile manually written SSA into ACIR");
         // Expected result:
         // main f0
@@ -2564,11 +2550,7 @@ mod test {
         let ssa = builder.finish();
 
         let acir_functions = ssa
-            .into_acir(
-                &Brillig::default(),
-                noirc_frontend::Distinctness::Distinct,
-                &HashMap::default(),
-            )
+            .into_acir(&Brillig::default(), noirc_frontend::Distinctness::Distinct)
             .expect("Should compile manually written SSA into ACIR");
         // The expected result should look very similar to the abvoe test expect that the input witnesses of the `Call`
         // opcodes will be different. The changes can discerned from the checks below.
@@ -2659,11 +2641,7 @@ mod test {
         let ssa = builder.finish();
 
         let acir_functions = ssa
-            .into_acir(
-                &Brillig::default(),
-                noirc_frontend::Distinctness::Distinct,
-                &HashMap::default(),
-            )
+            .into_acir(&Brillig::default(), noirc_frontend::Distinctness::Distinct)
             .expect("Should compile manually written SSA into ACIR");
 
         assert_eq!(acir_functions.len(), 3, "Should have three ACIR functions");

diff --git a/compiler/noirc_evaluator/src/ssa/function_builder/mod.rs b/compiler/noirc_evaluator/src/ssa/function_builder/mod.rs
@@ -310,7 +310,8 @@ impl FunctionBuilder {
         index: ValueId,
         value: ValueId,
     ) -> ValueId {
-        self.insert_instruction(Instruction::ArraySet { array, index, value }, None).first()
+        self.insert_instruction(Instruction::ArraySet { array, index, value, mutable: false }, None)
+            .first()
     }
 
     /// Insert an instruction to increment an array's reference count. This only has an effect
@@ -500,7 +501,6 @@ mod tests {
     use acvm::FieldElement;
 
     use crate::ssa::ir::{
-        function::{InlineType, RuntimeType},
         instruction::{Endian, Intrinsic},
         map::Id,
         types::Type,

diff --git a/compiler/noirc_evaluator/src/ssa/ir/instruction.rs b/compiler/noirc_evaluator/src/ssa/ir/instruction.rs
@@ -189,8 +189,9 @@ pub(crate) enum Instruction {
     ArrayGet { array: ValueId, index: ValueId },
 
     /// Creates a new array with the new value at the given index. All other elements are identical
-    /// to those in the given array. This will not modify the original array.
-    ArraySet { array: ValueId, index: ValueId, value: ValueId },
+    /// to those in the given array. This will not modify the original array unless `mutable` is
+    /// set. This flag is off by default and only enabled when optimizations determine it is safe.
+    ArraySet { array: ValueId, index: ValueId, value: ValueId, mutable: bool },
 
     /// An instruction to increment the reference count of a value.
     ///
@@ -363,9 +364,12 @@ impl Instruction {
             Instruction::ArrayGet { array, index } => {
                 Instruction::ArrayGet { array: f(*array), index: f(*index) }
             }
-            Instruction::ArraySet { array, index, value } => {
-                Instruction::ArraySet { array: f(*array), index: f(*index), value: f(*value) }
-            }
+            Instruction::ArraySet { array, index, value, mutable } => Instruction::ArraySet {
+                array: f(*array),
+                index: f(*index),
+                value: f(*value),
+                mutable: *mutable,
+            },
             Instruction::IncrementRc { value } => Instruction::IncrementRc { value: f(*value) },
             Instruction::DecrementRc { value } => Instruction::DecrementRc { value: f(*value) },
             Instruction::RangeCheck { value, max_bit_size, assert_message } => {
@@ -416,7 +420,7 @@ impl Instruction {
                 f(*array);
                 f(*index);
             }
-            Instruction::ArraySet { array, index, value } => {
+            Instruction::ArraySet { array, index, value, mutable: _ } => {
                 f(*array);
                 f(*index);
                 f(*value);

diff --git a/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs b/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs
@@ -339,8 +339,13 @@ fn simplify_slice_push_back(
     let element_size = element_type.element_size();
     let new_slice = dfg.make_array(slice, element_type);
 
-    let set_last_slice_value_instr =
-        Instruction::ArraySet { array: new_slice, index: arguments[0], value: arguments[2] };
+    let set_last_slice_value_instr = Instruction::ArraySet {
+        array: new_slice,
+        index: arguments[0],
+        value: arguments[2],
+        mutable: false,
+    };
+
     let set_last_slice_value = dfg
         .insert_instruction_and_results(set_last_slice_value_instr, block, None, call_stack)
         .first();

diff --git a/compiler/noirc_evaluator/src/ssa/ir/printer.rs b/compiler/noirc_evaluator/src/ssa/ir/printer.rs
@@ -176,14 +176,12 @@ fn display_instruction_inner(
         Instruction::ArrayGet { array, index } => {
             writeln!(f, "array_get {}, index {}", show(*array), show(*index))
         }
-        Instruction::ArraySet { array, index, value } => {
-            writeln!(
-                f,
-                "array_set {}, index {}, value {}",
-                show(*array),
-                show(*index),
-                show(*value)
-            )
+        Instruction::ArraySet { array, index, value, mutable } => {
+            let array = show(*array);
+            let index = show(*index);
+            let value = show(*value);
+            let mutable = if *mutable { " mut" } else { "" };
+            writeln!(f, "array_set{mutable} {array}, index {index}, value {value}",)
         }
         Instruction::IncrementRc { value } => {
             writeln!(f, "inc_rc {}", show(*value))