From 26d235198f9a2fedbe438b3f7b39184554c5e1c1 Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Mon, 2 Dec 2024 11:53:34 -0500 Subject: [PATCH] feat(ssa): Hoisting of array get using known induction variable maximum (#6639) --- .../src/ssa/opt/loop_invariant.rs | 201 +++++++++++++++++- .../noirc_evaluator/src/ssa/opt/unrolling.rs | 4 +- .../loop_invariant_regression/src/main.nr | 13 ++ 3 files changed, 208 insertions(+), 10 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs b/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs index 14233ca73e5..f1dd511402e 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs @@ -7,14 +7,16 @@ //! - Already marked as loop invariants //! //! We also check that we are not hoisting instructions with side effects. -use fxhash::FxHashSet as HashSet; +use acvm::{acir::AcirField, FieldElement}; +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; use crate::ssa::{ ir::{ basic_block::BasicBlockId, function::{Function, RuntimeType}, function_inserter::FunctionInserter, - instruction::InstructionId, + instruction::{Instruction, InstructionId}, + types::Type, value::ValueId, }, Ssa, @@ -45,25 +47,51 @@ impl Function { } impl Loops { - fn hoist_loop_invariants(self, function: &mut Function) { + fn hoist_loop_invariants(mut self, function: &mut Function) { let mut context = LoopInvariantContext::new(function); - for loop_ in self.yet_to_unroll.iter() { + // The loops should be sorted by the number of blocks. + // We want to access outer nested loops first, which we do by popping + // from the top of the list. + while let Some(loop_) = self.yet_to_unroll.pop() { let Ok(pre_header) = loop_.get_pre_header(context.inserter.function, &self.cfg) else { // If the loop does not have a preheader we skip hoisting loop invariants for this loop continue; }; - context.hoist_loop_invariants(loop_, pre_header); + + context.hoist_loop_invariants(&loop_, pre_header); } context.map_dependent_instructions(); } } +impl Loop { + /// Find the value that controls whether to perform a loop iteration. + /// This is going to be the block parameter of the loop header. + /// + /// Consider the following example of a `for i in 0..4` loop: + /// ```text + /// brillig(inline) fn main f0 { + /// b0(v0: u32): + /// ... + /// jmp b1(u32 0) + /// b1(v1: u32): // Loop header + /// v5 = lt v1, u32 4 // Upper bound + /// jmpif v5 then: b3, else: b2 + /// ``` + /// In the example above, `v1` is the induction variable + fn get_induction_variable(&self, function: &Function) -> ValueId { + function.dfg.block_parameters(self.header)[0] + } +} + struct LoopInvariantContext<'f> { inserter: FunctionInserter<'f>, defined_in_loop: HashSet, loop_invariants: HashSet, + // Maps induction variable -> fixed upper loop bound + outer_induction_variables: HashMap, } impl<'f> LoopInvariantContext<'f> { @@ -72,6 +100,7 @@ impl<'f> LoopInvariantContext<'f> { inserter: FunctionInserter::new(function), defined_in_loop: HashSet::default(), loop_invariants: HashSet::default(), + outer_induction_variables: HashMap::default(), } } @@ -88,13 +117,29 @@ impl<'f> LoopInvariantContext<'f> { self.inserter.push_instruction(instruction_id, *block); } - self.update_values_defined_in_loop_and_invariants(instruction_id, hoist_invariant); + 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); + } } /// Gather the variables declared within the loop fn set_values_defined_in_loop(&mut self, loop_: &Loop) { + // Clear any values that may be defined in previous loops, as the context is per function. + self.defined_in_loop.clear(); + // These are safe to keep per function, but we want to be clear that these values + // are used per loop. + self.loop_invariants.clear(); + for block in loop_.blocks.iter() { let params = self.inserter.function.dfg.block_parameters(*block); self.defined_in_loop.extend(params); @@ -107,7 +152,7 @@ impl<'f> LoopInvariantContext<'f> { /// Update any values defined in the loop and loop invariants after a /// analyzing and re-inserting a loop's instruction. - fn update_values_defined_in_loop_and_invariants( + fn extend_values_defined_in_loop_and_invariants( &mut self, instruction_id: InstructionId, hoist_invariant: bool, @@ -143,9 +188,45 @@ impl<'f> LoopInvariantContext<'f> { is_loop_invariant &= !self.defined_in_loop.contains(&value) || self.loop_invariants.contains(&value); }); - is_loop_invariant && instruction.can_be_deduplicated(&self.inserter.function.dfg, false) + + let can_be_deduplicated = instruction + .can_be_deduplicated(&self.inserter.function.dfg, false) + || self.can_be_deduplicated_from_upper_bound(&instruction); + + is_loop_invariant && can_be_deduplicated + } + + /// 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 + /// index that is below the length of the array. + /// Checking an array get where the index is the loop's induction variable on its own + /// would determine that the instruction is not safe for hoisting. + /// However, if we know that the induction variable's upper bound will always be in bounds of the array + /// we can safely hoist the array access. + fn can_be_deduplicated_from_upper_bound(&self, instruction: &Instruction) -> bool { + match instruction { + Instruction::ArrayGet { array, index } => { + let array_typ = self.inserter.function.dfg.type_of_value(*array); + let upper_bound = self.outer_induction_variables.get(index); + if let (Type::Array(_, len), Some(upper_bound)) = (array_typ, upper_bound) { + upper_bound.to_u128() as usize <= len + } else { + false + } + } + _ => false, + } } + /// 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. + /// However, there may be instructions which are not within loops that are + /// still reliant upon the instruction results altered during the pass. + /// This method re-inserts all instructions so that all instructions have + /// correct new value IDs based upon the `FunctionInserter` internal map. + /// Leaving out this mapping could lead to instructions with values that do not exist. fn map_dependent_instructions(&mut self) { let blocks = self.inserter.function.reachable_blocks(); for block in blocks { @@ -375,4 +456,108 @@ mod test { // The code should be unchanged assert_normalized_ssa_equals(ssa, src); } + + #[test] + fn hoist_array_gets_using_induction_variable_with_const_bound() { + // SSA for the following program: + // + // fn triple_loop(x: u32) { + // let arr = [2; 5]; + // for i in 0..4 { + // for j in 0..4 { + // for _ in 0..4 { + // assert_eq(arr[i], x); + // assert_eq(arr[j], x); + // } + // } + // } + // } + // + // `arr[i]` and `arr[j]` are safe to hoist as we know the maximum possible index + // to be used for both array accesses. + // We want to make sure `arr[i]` is hoisted to the outermost loop body and that + // `arr[j]` is hoisted to the second outermost loop body. + let src = " + brillig(inline) fn main f0 { + b0(v0: u32, v1: u32): + v6 = make_array [u32 2, u32 2, u32 2, u32 2, u32 2] : [u32; 5] + inc_rc v6 + jmp b1(u32 0) + b1(v2: u32): + v9 = lt v2, u32 4 + jmpif v9 then: b3, else: b2 + b3(): + jmp b4(u32 0) + b4(v3: u32): + v10 = lt v3, u32 4 + jmpif v10 then: b6, else: b5 + b6(): + jmp b7(u32 0) + b7(v4: u32): + v13 = lt v4, u32 4 + jmpif v13 then: b9, else: b8 + b9(): + v15 = array_get v6, index v2 -> u32 + v16 = eq v15, v0 + constrain v15 == v0 + v17 = array_get v6, index v3 -> u32 + v18 = eq v17, v0 + constrain v17 == v0 + v19 = add v4, u32 1 + jmp b7(v19) + b8(): + v14 = add v3, u32 1 + jmp b4(v14) + b5(): + v12 = add v2, u32 1 + jmp b1(v12) + b2(): + return + } + "; + + let ssa = Ssa::from_str(src).unwrap(); + + let expected = " + brillig(inline) fn main f0 { + b0(v0: u32, v1: u32): + v6 = make_array [u32 2, u32 2, u32 2, u32 2, u32 2] : [u32; 5] + inc_rc v6 + jmp b1(u32 0) + b1(v2: u32): + v9 = lt v2, u32 4 + jmpif v9 then: b3, else: b2 + b3(): + v10 = array_get v6, index v2 -> u32 + v11 = eq v10, v0 + jmp b4(u32 0) + b4(v3: u32): + v12 = lt v3, u32 4 + jmpif v12 then: b6, else: b5 + b6(): + v15 = array_get v6, index v3 -> u32 + v16 = eq v15, v0 + jmp b7(u32 0) + b7(v4: u32): + v17 = lt v4, u32 4 + jmpif v17 then: b9, else: b8 + b9(): + constrain v10 == v0 + constrain v15 == v0 + v19 = add v4, u32 1 + jmp b7(v19) + b8(): + v18 = add v3, u32 1 + jmp b4(v18) + b5(): + v14 = add v2, u32 1 + jmp b1(v14) + b2(): + return + } + "; + + let ssa = ssa.loop_invariant_code_motion(); + assert_normalized_ssa_equals(ssa, expected); + } } diff --git a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs index 5883ce25936..1a13acc5435 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs @@ -109,7 +109,7 @@ impl Function { pub(super) struct Loop { /// The header block of a loop is the block which dominates all the /// other blocks in the loop. - header: BasicBlockId, + pub(super) header: BasicBlockId, /// The start of the back_edge n -> d is the block n at the end of /// the loop that jumps back to the header block d which restarts the loop. @@ -299,7 +299,7 @@ impl Loop { /// v5 = lt v1, u32 4 // Upper bound /// jmpif v5 then: b3, else: b2 /// ``` - fn get_const_upper_bound(&self, function: &Function) -> Option { + pub(super) fn get_const_upper_bound(&self, function: &Function) -> Option { let block = &function.dfg[self.header]; let instructions = block.instructions(); assert_eq!( diff --git a/test_programs/execution_success/loop_invariant_regression/src/main.nr b/test_programs/execution_success/loop_invariant_regression/src/main.nr index 25f6e92f868..c28ce063116 100644 --- a/test_programs/execution_success/loop_invariant_regression/src/main.nr +++ b/test_programs/execution_success/loop_invariant_regression/src/main.nr @@ -2,6 +2,7 @@ // to be hoisted to the loop's pre-header block. fn main(x: u32, y: u32) { loop(4, x, y); + array_read_loop(4, x); } fn loop(upper_bound: u32, x: u32, y: u32) { @@ -11,3 +12,15 @@ fn loop(upper_bound: u32, x: u32, y: u32) { assert_eq(z, 12); } } + +fn array_read_loop(upper_bound: u32, x: u32) { + let arr = [2; 5]; + for i in 0..upper_bound { + for j in 0..upper_bound { + for _ in 0..upper_bound { + assert_eq(arr[i], x); + assert_eq(arr[j], x); + } + } + } +}