feat(performance): Use unchecked ops based upon known induction variables

compiler/noirc_evaluator/src/ssa/ir/instruction/binary.rs

@@ -39,9 +39,9 @@
     Or,
     /// Bitwise xor (^)
     Xor,
     /// Bitshift left (<<)
     Shl,
     /// Bitshift right (>>)
     Shr,
 }
 
@@ -107,7 +107,7 @@
         }
 
         let operator = if lhs_type == NumericType::NativeField {
             // Unchecked operations between fields or bools don't make sense, so we convert those to non-unchecked
             // to reduce noise and confusion in the generated SSA.
             match operator {
                 BinaryOp::Add { unchecked: true } => BinaryOp::Add { unchecked: false },
@@ -116,7 +116,7 @@
                 _ => operator,
             }
         } else if lhs_type == NumericType::bool() {
             // Unchecked mul between bools doesn't make sense, so we convert that to non-unchecked
             if let BinaryOp::Mul { unchecked: true } = operator {
                 BinaryOp::Mul { unchecked: false }
             } else {
@@ -289,7 +289,7 @@
                 }
                 if lhs_type == NumericType::bool() {
                     // Boolean AND is equivalent to multiplication, which is a cheaper operation.
                     // (mul unchecked because these are bools so it doesn't matter really)
                     let instruction =
                         Instruction::binary(BinaryOp::Mul { unchecked: true }, lhs, rhs);
                     return SimplifyResult::SimplifiedToInstruction(instruction);
@@ -563,6 +563,24 @@
             BinaryOp::Shr => |x, y| Some(x >> y),
         }
     }
+
+    pub(crate) fn into_unchecked(self) -> Self {
+        match self {
+            BinaryOp::Add { .. } => BinaryOp::Add { unchecked: true },
+            BinaryOp::Sub { .. } => BinaryOp::Sub { unchecked: true },
+            BinaryOp::Mul { .. } => BinaryOp::Mul { unchecked: true },
+            _ => self,
+        }
+    }
+
+    pub(crate) fn is_unchecked(self) -> bool {
+        match self {
+            BinaryOp::Add { unchecked }
+            | BinaryOp::Sub { unchecked }
+            | BinaryOp::Mul { unchecked } => unchecked,
+            _ => true,
+        }
+    }
 }
 
 #[cfg(test)]

compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs

@@ -15,7 +15,9 @@ use crate::ssa::{
         basic_block::BasicBlockId,
         function::Function,
         function_inserter::FunctionInserter,
-        instruction::{binary::eval_constant_binary_op, BinaryOp, Instruction, InstructionId},
+        instruction::{
+            binary::eval_constant_binary_op, Binary, BinaryOp, Instruction, InstructionId,
+        },
         post_order::PostOrder,
         types::Type,
         value::ValueId,
@@ -87,7 +89,9 @@ struct LoopInvariantContext<'f> {
     inserter: FunctionInserter<'f>,
     defined_in_loop: HashSet<ValueId>,
     loop_invariants: HashSet<ValueId>,
-    // Maps induction variable -> fixed upper loop bound
+    // Maps current loop induction variable -> fixed upper loop bound
+    current_induction_variables: HashMap<ValueId, FieldElement>,
+    // Maps outer loop induction variable -> fixed upper loop bound
     outer_induction_variables: HashMap<ValueId, FieldElement>,
 }
 
@@ -97,15 +101,19 @@ impl<'f> LoopInvariantContext<'f> {
             inserter: FunctionInserter::new(function),
             defined_in_loop: HashSet::default(),
             loop_invariants: HashSet::default(),
+            current_induction_variables: HashMap::default(),
             outer_induction_variables: HashMap::default(),
         }
     }
 
     fn hoist_loop_invariants(&mut self, loop_: &Loop, pre_header: BasicBlockId) {
         self.set_values_defined_in_loop(loop_);
+        self.set_induction_var_bounds(loop_, true);
 
         for block in loop_.blocks.iter() {
             for instruction_id in self.inserter.function.dfg[*block].take_instructions() {
+                self.transform_to_unchecked_from_upper_bound(instruction_id);
+
                 let hoist_invariant = self.can_hoist_invariant(instruction_id);
 
                 if hoist_invariant {
@@ -135,20 +143,11 @@ impl<'f> LoopInvariantContext<'f> {
                 } else {
                     self.inserter.push_instruction(instruction_id, *block);
                 }
-
                 self.extend_values_defined_in_loop_and_invariants(instruction_id, hoist_invariant);
             }
         }
 
-        // Keep track of a loop induction variable and respective upper bound.
-        // This will be used by later loops to determine whether they have operations
-        // reliant upon the maximum induction variable.
-        let upper_bound = loop_.get_const_upper_bound(self.inserter.function);
-        if let Some(upper_bound) = upper_bound {
-            let induction_variable = loop_.get_induction_variable(self.inserter.function);
-            let induction_variable = self.inserter.resolve(induction_variable);
-            self.outer_induction_variables.insert(induction_variable, upper_bound);
-        }
+        self.set_induction_var_bounds(loop_, false);
     }
 
     /// Gather the variables declared within the loop
@@ -216,6 +215,24 @@ impl<'f> LoopInvariantContext<'f> {
         is_loop_invariant && can_be_deduplicated
     }
 
+    /// Keep track of a loop induction variable and respective upper bound.
+    /// In the case of a nested loop, this will be used by later loops to determine
+    /// whether they have operations reliant upon the maximum induction variable.
+    /// When within the current loop, the known upper bound can be used to simplify instructions,
+    /// such as transforming a checked add to an unchecked add.
+    fn set_induction_var_bounds(&mut self, loop_: &Loop, current_loop: bool) {
+        let upper_bound = loop_.get_const_upper_bound(self.inserter.function);
+        if let Some(upper_bound) = upper_bound {
+            let induction_variable = loop_.get_induction_variable(self.inserter.function);
+            let induction_variable = self.inserter.resolve(induction_variable);
+            if current_loop {
+                self.current_induction_variables.insert(induction_variable, upper_bound);
+            } else {
+                self.outer_induction_variables.insert(induction_variable, upper_bound);
+            }
+        }
+    }
+
     /// Certain instructions can take advantage of that our induction variable has a fixed maximum.
     ///
     /// For example, an array access can usually only be safely deduplicated when we have a constant
@@ -236,34 +253,68 @@ impl<'f> LoopInvariantContext<'f> {
                 }
             }
             Instruction::Binary(binary) => {
-                if !matches!(binary.operator, BinaryOp::Add { .. } | BinaryOp::Mul { .. }) {
-                    return false;
-                }
-
-                let operand_type =
-                    self.inserter.function.dfg.type_of_value(binary.lhs).unwrap_numeric();
-
-                let lhs_const =
-                    self.inserter.function.dfg.get_numeric_constant_with_type(binary.lhs);
-                let rhs_const =
-                    self.inserter.function.dfg.get_numeric_constant_with_type(binary.rhs);
-                let (lhs, rhs) = match (
-                    lhs_const,
-                    rhs_const,
-                    self.outer_induction_variables.get(&binary.lhs),
-                    self.outer_induction_variables.get(&binary.rhs),
-                ) {
-                    (Some((lhs, _)), None, None, Some(upper_bound)) => (lhs, *upper_bound),
-                    (None, Some((rhs, _)), Some(upper_bound), None) => (*upper_bound, rhs),
-                    _ => return false,
-                };
-
-                eval_constant_binary_op(lhs, rhs, binary.operator, operand_type).is_some()
+                self.can_evaluate_binary_op(binary, &self.outer_induction_variables)
             }
             _ => false,
         }
     }
 
+    /// Binary operations can take advantage of that our induction variable has a fixed maximum,
+    /// to be transformed from a checked operation to an unchecked operation.
+    /// Checked operations require more bytecode and thus we aim to minimize their usage wherever possible.
+    ///
+    /// If one side of a binary operation is a constant and the other is an induction variable
+    /// with a known upper bound, we know whether that binary operation will ever overflow.
+    /// If we determine that an overflow is not possible we can convert the checked operation to unchecked.
+    fn transform_to_unchecked_from_upper_bound(&mut self, instruction_id: InstructionId) {
+        let Instruction::Binary(binary) = &self.inserter.function.dfg[instruction_id] else {
+            return;
+        };
+
+        if binary.operator.is_unchecked()
+            || !self.can_evaluate_binary_op(binary, &self.current_induction_variables)
+        {
+            return;
+        }
+
+        if let Instruction::Binary(binary) = &mut self.inserter.function.dfg[instruction_id] {
+            binary.operator = binary.operator.into_unchecked();
+        };
+    }
+
+    /// Checks whether a binary operation can be evaluated using the upper bound of the given loop induction variables.
+    /// If it cannot be evaluated, it means that we either have a dynamic loop bound or
+    /// that the operation can potentially overflow at the upper loop bound.
+    fn can_evaluate_binary_op(
+        &self,
+        binary: &Binary,
+        induction_vars: &HashMap<ValueId, FieldElement>,
+    ) -> bool {
+        if !matches!(binary.operator, BinaryOp::Add { .. } | BinaryOp::Mul { .. }) {
+            return false;
+        }
+
+        let operand_type = self.inserter.function.dfg.type_of_value(binary.lhs).unwrap_numeric();
+
+        let lhs_const = self.inserter.function.dfg.get_numeric_constant_with_type(binary.lhs);
+        let rhs_const = self.inserter.function.dfg.get_numeric_constant_with_type(binary.rhs);
+        let (lhs, rhs) = match (
+            lhs_const,
+            rhs_const,
+            induction_vars.get(&binary.lhs),
+            induction_vars.get(&binary.rhs),
+        ) {
+            (Some((lhs, _)), None, None, Some(upper_bound)) => (lhs, *upper_bound),
+            (None, Some((rhs, _)), Some(upper_bound), None) => (*upper_bound, rhs),
+            _ => return false,
+        };
+
+        // We evaluate this expression using the upper bounds of its inputs to check whether it will ever overflow.
+        // If so, this will cause `eval_constant_binary_op` to return `None`.
+        // Therefore a `Some` value shows that this operation is safe.
+        eval_constant_binary_op(lhs, rhs, binary.operator, operand_type).is_some()
+    }
+
     /// Loop invariant hoisting only operates over loop instructions.
     /// The `FunctionInserter` is used for mapping old values to new values after
     /// re-inserting loop invariant instructions.
@@ -304,7 +355,7 @@ mod test {
           b3():
               v6 = mul v0, v1
               constrain v6 == i32 6
-              v8 = add v2, i32 1
+              v8 = unchecked_add v2, i32 1
               jmp b1(v8)
         }
         ";
@@ -328,7 +379,7 @@ mod test {
             return
           b3():
             constrain v3 == i32 6
-            v9 = add v2, i32 1
+            v9 = unchecked_add v2, i32 1
             jmp b1(v9)
         }
         ";
@@ -356,12 +407,12 @@ mod test {
             v7 = lt v3, i32 4
             jmpif v7 then: b6, else: b5
           b5():
-            v9 = add v2, i32 1
+            v9 = unchecked_add v2, i32 1
             jmp b1(v9)
           b6():
             v10 = mul v0, v1
             constrain v10 == i32 6
-            v12 = add v3, i32 1
+            v12 = unchecked_add v3, i32 1
             jmp b4(v12)
         }
         ";
@@ -389,11 +440,11 @@ mod test {
             v8 = lt v3, i32 4
             jmpif v8 then: b6, else: b5
           b5():
-            v10 = add v2, i32 1
+            v10 = unchecked_add v2, i32 1
             jmp b1(v10)
           b6():
             constrain v4 == i32 6
-            v12 = add v3, i32 1
+            v12 = unchecked_add v3, i32 1
             jmp b4(v12)
         }
         ";
@@ -429,7 +480,7 @@ mod test {
             v7 = mul v6, v0
             v8 = eq v7, i32 12
             constrain v7 == i32 12
-            v9 = add v2, i32 1
+            v9 = unchecked_add v2, i32 1
             jmp b1(v9)
         }
         ";
@@ -454,7 +505,7 @@ mod test {
             return
           b3():
             constrain v4 == i32 12
-            v11 = add v2, i32 1
+            v11 = unchecked_add v2, i32 1
             jmp b1(v11)
         }
         ";
@@ -488,7 +539,7 @@ mod test {
             v12 = load v5 -> [u32; 5]
             v13 = array_set v12, index v0, value v1
             store v13 at v5
-            v15 = add v2, u32 1
+            v15 = unchecked_add v2, u32 1
             jmp b1(v15)
         }
         ";
@@ -541,15 +592,15 @@ mod test {
             v10 = lt v3, u32 4
             jmpif v10 then: b6, else: b5
           b5():
-            v12 = add v2, u32 1
+            v12 = unchecked_add v2, u32 1
             jmp b1(v12)
           b6():
             jmp b7(u32 0)
           b7(v4: u32):
             v13 = lt v4, u32 4
             jmpif v13 then: b9, else: b8
           b8():
-            v14 = add v3, u32 1
+            v14 = unchecked_add v3, u32 1
             jmp b4(v14)
           b9():
             v15 = array_get v6, index v2 -> u32
@@ -558,7 +609,7 @@ mod test {
             v17 = array_get v6, index v3 -> u32
             v18 = eq v17, v0
             constrain v17 == v0
-            v19 = add v4, u32 1
+            v19 = unchecked_add v4, u32 1
             jmp b7(v19)
         }
         ";
@@ -584,7 +635,7 @@ mod test {
             v12 = lt v3, u32 4
             jmpif v12 then: b6, else: b5
           b5():
-            v14 = add v2, u32 1
+            v14 = unchecked_add v2, u32 1
             jmp b1(v14)
           b6():
             v15 = array_get v6, index v3 -> u32
@@ -594,12 +645,12 @@ mod test {
             v17 = lt v4, u32 4
             jmpif v17 then: b9, else: b8
           b8():
-            v18 = add v3, u32 1
+            v18 = unchecked_add v3, u32 1
             jmp b4(v18)
           b9():
             constrain v10 == v0
             constrain v15 == v0
-            v19 = add v4, u32 1
+            v19 = unchecked_add v4, u32 1
             jmp b7(v19)
         }
         ";
@@ -649,7 +700,7 @@ mod test {
             v21 = add v1, v2
             v23 = array_set v19, index v21, value Field 128
             call f1(v23)
-            v24 = add v2, u32 1
+            v24 = unchecked_add v2, u32 1
             jmp b1(v24)
         }
         brillig(inline) fn foo f1 {
@@ -685,7 +736,7 @@ mod test {
             v21 = add v1, v2
             v23 = array_set v14, index v21, value Field 128
             call f1(v23)
-            v24 = add v2, u32 1
+            v24 = unchecked_add v2, u32 1
             jmp b1(v24)
         }
         brillig(inline) fn foo f1 {
@@ -697,4 +748,49 @@ mod test {
         let ssa = ssa.loop_invariant_code_motion();
         assert_normalized_ssa_equals(ssa, expected);
     }
+
+    #[test]
+    fn transform_safe_ops_to_unchecked_during_code_motion() {
+        // This test is identical to `simple_loop_invariant_code_motion`, except this test
+        // uses a checked add in `b3`.
+        let src = "
+        brillig(inline) fn main f0 {
+          b0(v0: i32, v1: i32):
+              jmp b1(i32 0)
+          b1(v2: i32):
+              v5 = lt v2, i32 4
+              jmpif v5 then: b3, else: b2
+          b2():
+              return
+          b3():
+              v6 = mul v0, v1
+              constrain v6 == i32 6
+              v8 = add v2, i32 1
+              jmp b1(v8)
+        }
+        ";
+
+        let ssa = Ssa::from_str(src).unwrap();
+
+        // `v8 = add v2, i32 1` in b3 should now be `v9 = unchecked_add v2, i32 1` in b3
+        let expected = "
+        brillig(inline) fn main f0 {
+          b0(v0: i32, v1: i32):
+            v3 = mul v0, v1
+            jmp b1(i32 0)
+          b1(v2: i32):
+            v6 = lt v2, i32 4
+            jmpif v6 then: b3, else: b2
+          b2():
+            return
+          b3():
+            constrain v3 == i32 6
+            v9 = unchecked_add v2, i32 1
+            jmp b1(v9)
+        }
+        ";
+
+        let ssa = ssa.loop_invariant_code_motion();
+        assert_normalized_ssa_equals(ssa, expected);
+    }
 }