feat: Optimize array sets in if conditions (alternate version) (#4716)

# Description ## Problem\* Alternate version of #4695 Resolves #4692 ## Summary\* - Flattening produces a new `IfElse` instruction to merge values instead of immediately merging them - `IfElse` instructions are removed by actually merging the values once flattening and mem2reg is done. This is done in a new remove_if_else pass. - Since slice capacity tracking was done in flattening before and required to merge values, I've had to move this to the remove_if_else pass instead. It is somewhat simpler since we don't have to track loads/stores. ## Additional Context Currently only failing the nested_array_in_slice and nested_array_dynamic tests. The array being asserted somehow has the value `[1, 2, 3]` despite none of these values being set in the program. There are currently no actual optimizations here, this is just structural changes so it should otherwise be equivalent. ## 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. --------- Co-authored-by: Tom French <[email protected]> Co-authored-by: vezenovm <[email protected]>
noir-lang · May 3, 2024 · a87c655 · a87c655
1 parent 102ff13
commit a87c655
Show file tree

Hide file tree

Showing 16 changed files with 601 additions and 109 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
@@ -631,7 +631,7 @@ impl<'block> BrilligBlock<'block> {
                     destination_variable,
                 );
             }
-            Instruction::ArraySet { array, index, value, .. } => {
+            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);
@@ -712,6 +712,9 @@ impl<'block> BrilligBlock<'block> {
             Instruction::EnableSideEffects { .. } => {
                 todo!("enable_side_effects not supported by brillig")
             }
+            Instruction::IfElse { .. } => {
+                unreachable!("IfElse instructions should not be possible in brillig")
+            }
         };
 
         let dead_variables = self

diff --git a/compiler/noirc_evaluator/src/ssa.rs b/compiler/noirc_evaluator/src/ssa.rs
@@ -59,6 +59,7 @@ pub(crate) fn optimize_into_acir(
         .run_pass(Ssa::remove_bit_shifts, "After Removing Bit Shifts:")
         // Run mem2reg once more with the flattened CFG to catch any remaining loads/stores
         .run_pass(Ssa::mem2reg, "After Mem2Reg:")
+        .run_pass(Ssa::remove_if_else, "After Remove IfElse:")
         // 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.
         .run_pass(Ssa::inline_functions_with_no_predicates, "After Inlining:")

diff --git a/compiler/noirc_evaluator/src/ssa/acir_gen/mod.rs b/compiler/noirc_evaluator/src/ssa/acir_gen/mod.rs
@@ -686,6 +686,9 @@ impl<'a> Context<'a> {
                     assert_message.clone(),
                 )?;
             }
+            Instruction::IfElse { .. } => {
+                unreachable!("IfElse instruction remaining in acir-gen")
+            }
         }
 
         self.acir_context.set_call_stack(CallStack::new());
@@ -1009,6 +1012,7 @@ impl<'a> Context<'a> {
                             });
                         }
                     };
+
                     if self.acir_context.is_constant_one(&self.current_side_effects_enabled_var) {
                         // Report the error if side effects are enabled.
                         if index >= array_size {

diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs
@@ -554,7 +554,10 @@ impl<'dfg> InsertInstructionResult<'dfg> {
         match self {
             InsertInstructionResult::SimplifiedTo(value) => *value,
             InsertInstructionResult::SimplifiedToMultiple(values) => values[0],
-            InsertInstructionResult::Results(_, results) => results[0],
+            InsertInstructionResult::Results(_, results) => {
+                assert_eq!(results.len(), 1);
+                results[0]
+            }
             InsertInstructionResult::InstructionRemoved => {
                 panic!("Instruction was removed, no results")
             }
@@ -583,6 +586,24 @@ impl<'dfg> InsertInstructionResult<'dfg> {
     }
 }
 
+impl<'dfg> std::ops::Index<usize> for InsertInstructionResult<'dfg> {
+    type Output = ValueId;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        match self {
+            InsertInstructionResult::Results(_, results) => &results[index],
+            InsertInstructionResult::SimplifiedTo(result) => {
+                assert_eq!(index, 0);
+                result
+            }
+            InsertInstructionResult::SimplifiedToMultiple(results) => &results[index],
+            InsertInstructionResult::InstructionRemoved => {
+                panic!("Cannot index into InsertInstructionResult::InstructionRemoved")
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::DataFlowGraph;

diff --git a/compiler/noirc_evaluator/src/ssa/ir/instruction.rs b/compiler/noirc_evaluator/src/ssa/ir/instruction.rs
@@ -1,6 +1,8 @@
 use acvm::{acir::BlackBoxFunc, FieldElement};
 use iter_extended::vecmap;
 
+use crate::ssa::opt::flatten_cfg::value_merger::ValueMerger;
+
 use super::{
     basic_block::BasicBlockId,
     dfg::{CallStack, DataFlowGraph},
@@ -206,6 +208,14 @@ pub(crate) enum Instruction {
     /// implemented via reference counting. In ACIR code this is done with im::Vector and these
     /// DecrementRc instructions are ignored.
     DecrementRc { value: ValueId },
+
+    /// Merge two values returned from opposite branches of a conditional into one.
+    IfElse {
+        then_condition: ValueId,
+        then_value: ValueId,
+        else_condition: ValueId,
+        else_value: ValueId,
+    },
 }
 
 impl Instruction {
@@ -219,10 +229,12 @@ impl Instruction {
         match self {
             Instruction::Binary(binary) => binary.result_type(),
             Instruction::Cast(_, typ) => InstructionResultType::Known(typ.clone()),
-            Instruction::Not(value) | Instruction::Truncate { value, .. } => {
+            Instruction::Not(value)
+            | Instruction::Truncate { value, .. }
+            | Instruction::ArraySet { array: value, .. }
+            | Instruction::IfElse { then_value: value, .. } => {
                 InstructionResultType::Operand(*value)
             }
-            Instruction::ArraySet { array, .. } => InstructionResultType::Operand(*array),
             Instruction::Constrain(..)
             | Instruction::Store { .. }
             | Instruction::IncrementRc { .. }
@@ -270,6 +282,7 @@ impl Instruction {
             | Cast(_, _)
             | Not(_)
             | Truncate { .. }
+            | IfElse { .. }
             | ArrayGet { .. }
             | ArraySet { .. } => !self.requires_acir_gen_predicate(dfg),
         }
@@ -295,6 +308,7 @@ impl Instruction {
             | Allocate
             | Load { .. }
             | ArrayGet { .. }
+            | IfElse { .. }
             | ArraySet { .. } => true,
 
             Constrain(..)
@@ -349,6 +363,7 @@ impl Instruction {
             | Instruction::Allocate
             | Instruction::Load { .. }
             | Instruction::Store { .. }
+            | Instruction::IfElse { .. }
             | Instruction::IncrementRc { .. }
             | Instruction::DecrementRc { .. } => false,
         }
@@ -416,6 +431,14 @@ impl Instruction {
                     assert_message: assert_message.clone(),
                 }
             }
+            Instruction::IfElse { then_condition, then_value, else_condition, else_value } => {
+                Instruction::IfElse {
+                    then_condition: f(*then_condition),
+                    then_value: f(*then_value),
+                    else_condition: f(*else_condition),
+                    else_value: f(*else_value),
+                }
+            }
         }
     }
 
@@ -473,6 +496,12 @@ impl Instruction {
             | Instruction::RangeCheck { value, .. } => {
                 f(*value);
             }
+            Instruction::IfElse { then_condition, then_value, else_condition, else_value } => {
+                f(*then_condition);
+                f(*then_value);
+                f(*else_condition);
+                f(*else_value);
+            }
         }
     }
 
@@ -624,6 +653,36 @@ impl Instruction {
                     None
                 }
             }
+            Instruction::IfElse { then_condition, then_value, else_condition, else_value } => {
+                let typ = dfg.type_of_value(*then_value);
+
+                if let Some(constant) = dfg.get_numeric_constant(*then_condition) {
+                    if constant.is_one() {
+                        return SimplifiedTo(*then_value);
+                    } else if constant.is_zero() {
+                        return SimplifiedTo(*else_value);
+                    }
+                }
+
+                if matches!(&typ, Type::Numeric(_)) {
+                    let then_condition = *then_condition;
+                    let then_value = *then_value;
+                    let else_condition = *else_condition;
+                    let else_value = *else_value;
+
+                    let result = ValueMerger::merge_numeric_values(
+                        dfg,
+                        block,
+                        then_condition,
+                        else_condition,
+                        then_value,
+                        else_value,
+                    );
+                    SimplifiedTo(result)
+                } else {
+                    None
+                }
+            }
         }
     }
 }

diff --git a/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs b/compiler/noirc_evaluator/src/ssa/ir/instruction/call.rs
@@ -354,7 +354,9 @@ fn simplify_slice_push_back(
     slice_sizes.insert(set_last_slice_value, slice_size / element_size);
     slice_sizes.insert(new_slice, slice_size / element_size);
 
-    let mut value_merger = ValueMerger::new(dfg, block, &mut slice_sizes);
+    let unknown = &mut HashMap::default();
+    let mut value_merger = ValueMerger::new(dfg, block, &mut slice_sizes, unknown, None);
+
     let new_slice = value_merger.merge_values(
         len_not_equals_capacity,
         len_equals_capacity,

diff --git a/compiler/noirc_evaluator/src/ssa/ir/printer.rs b/compiler/noirc_evaluator/src/ssa/ir/printer.rs
@@ -184,7 +184,7 @@ fn display_instruction_inner(
             let index = show(*index);
             let value = show(*value);
             let mutable = if *mutable { " mut" } else { "" };
-            writeln!(f, "array_set{mutable} {array}, index {index}, value {value}",)
+            writeln!(f, "array_set{mutable} {array}, index {index}, value {value}")
         }
         Instruction::IncrementRc { value } => {
             writeln!(f, "inc_rc {}", show(*value))
@@ -195,6 +195,16 @@ fn display_instruction_inner(
         Instruction::RangeCheck { value, max_bit_size, .. } => {
             writeln!(f, "range_check {} to {} bits", show(*value), *max_bit_size,)
         }
+        Instruction::IfElse { then_condition, then_value, else_condition, else_value } => {
+            let then_condition = show(*then_condition);
+            let then_value = show(*then_value);
+            let else_condition = show(*else_condition);
+            let else_value = show(*else_value);
+            writeln!(
+                f,
+                "if {then_condition} then {then_value} else if {else_condition} then {else_value}"
+            )
+        }
     }
 }