From 4c3a1c1b36a819a7129ae18b420163ad34dde9c6 Mon Sep 17 00:00:00 2001 From: "Michael J. Klein" Date: Wed, 22 Oct 2025 14:52:51 -0400 Subject: [PATCH 1/3] chore: avoid unrolling loop headers twice --- .../noirc_evaluator/src/ssa/opt/unrolling.rs | 55 +++++++++++++++++-- 1 file changed, 49 insertions(+), 6 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs index f4faffcb7ae..9089bec959b 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs @@ -567,8 +567,14 @@ impl Loop { let mut unroll_into = self.get_pre_header(function, cfg)?; let mut jump_value = get_induction_variable(&function.dfg, unroll_into)?; - while let Some(context) = self.unroll_header(function, unroll_into, jump_value)? { - (unroll_into, jump_value) = context.unroll_loop_iteration(); + // TODO: WIP + while let (Some(context), loop_header_id) = + self.unroll_header(function, unroll_into, jump_value)? + { + // TODO: WIP + // context.visited_blocks.insert(loop_header_id); + + (unroll_into, jump_value) = context.unroll_loop_iteration(loop_header_id); } Ok(()) @@ -601,19 +607,24 @@ impl Loop { /// Unrolls the header block of the loop. This is the block that dominates all other blocks in the /// loop and contains the jmpif instruction that lets us know if we should continue looping. /// Returns Some(iteration context) if we should perform another iteration. + /// // TODO: WIP + /// and the loop header's BasicBlockId fn unroll_header<'a>( &'a self, function: &'a mut Function, unroll_into: BasicBlockId, induction_value: ValueId, - ) -> Result>, CallStack> { + // TODO: WIP + ) -> Result<(Option>, BasicBlockId), CallStack> { // We insert into a fresh block first and move instructions into the unroll_into block later // only once we verify the jmpif instruction has a constant condition. If it does not, we can // just discard this fresh block and leave the loop unmodified. let fresh_block = function.dfg.make_block(); let mut context = LoopIteration::new(function, self, fresh_block, self.header); - let source_block = &context.dfg()[context.source_block]; + // TODO: WIP + let loop_header_id = context.source_block; + let source_block = &context.dfg()[loop_header_id]; assert_eq!(source_block.parameters().len(), 1, "Expected only 1 argument in loop header"); // Insert the current value of the loop induction variable into our context. @@ -621,6 +632,8 @@ impl Loop { context.inserter.try_map_value(first_param, induction_value); // Copy over all instructions and a fresh terminator. context.inline_instructions_from_block(); + context.visited_blocks.insert(loop_header_id); + // Mutate the terminator if possible so that it points at the iteration block. match context.dfg()[fresh_block].unwrap_terminator() { TerminatorInstruction::JmpIf { @@ -651,7 +664,11 @@ impl Loop { // have no more loops to unroll, because that block was not part of the loop itself, // ie. it wasn't between `loop_header` and `loop_body`. Otherwise we have the `loop_body` // in `source_block` and can unroll that into the destination. - Ok(self.blocks.contains(&context.source_block).then_some(context)) + // TODO: WIP + Ok(( + self.blocks.contains(&context.source_block).then_some(context), + loop_header_id, + )) } else { // If this case is reached the loop either uses non-constant indices or we need // another pass, such as mem2reg to resolve them to constants. @@ -967,6 +984,11 @@ struct LoopIteration<'f> { original_blocks: HashMap, visited_blocks: HashSet, + // TODO: WIP + // #[cfg(test)] + // loop_header_id: Option, + encountered_loop_header: bool, + insert_block: BasicBlockId, source_block: BasicBlockId, @@ -992,6 +1014,9 @@ impl<'f> LoopIteration<'f> { blocks: HashMap::default(), original_blocks: HashMap::default(), visited_blocks: HashSet::default(), + // TODO: WIP + encountered_loop_header: false, + induction_value: None, } } @@ -1002,7 +1027,12 @@ impl<'f> LoopIteration<'f> { /// It is expected the terminator instructions are set up to branch into an empty block /// for further unrolling. When the loop is finished this will need to be mutated to /// jump to the end of the loop instead. - fn unroll_loop_iteration(mut self) -> (BasicBlockId, ValueId) { + // TODO: WIP + fn unroll_loop_iteration(mut self, loop_header: BasicBlockId) -> (BasicBlockId, ValueId) { + // TODO: WIP + // Mark the loop header as visited + self.visited_blocks.insert(loop_header); + // Kick off the unrolling from the initial source block. let mut next_blocks = self.unroll_loop_block(); @@ -1010,6 +1040,11 @@ impl<'f> LoopIteration<'f> { self.insert_block = block; self.source_block = self.get_original_block(block); + // TODO: WIP + // if cfg!(test) { + self.encountered_loop_header |= self.visited_blocks.contains(&self.source_block); + // } + if !self.visited_blocks.contains(&self.source_block) { let mut blocks = self.unroll_loop_block(); next_blocks.append(&mut blocks); @@ -1021,6 +1056,14 @@ impl<'f> LoopIteration<'f> { .induction_value .expect("Expected to find the induction variable by end of loop iteration"); + // TODO: WIP + // if cfg!(test) { + assert!( + self.encountered_loop_header, + "expected to encounter loop header when visiting blocks" + ); + // } + (end_block, induction_value) } From cb2b0b0cdc47f9770fa4b00d3188bb3fa41bb559 Mon Sep 17 00:00:00 2001 From: "Michael J. Klein" Date: Thu, 23 Oct 2025 14:14:20 -0400 Subject: [PATCH 2/3] . --- .../noirc_evaluator/src/ssa/opt/unrolling.rs | 45 +++++-------------- 1 file changed, 11 insertions(+), 34 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs index 9089bec959b..dcaca7c3fce 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs @@ -567,13 +567,9 @@ impl Loop { let mut unroll_into = self.get_pre_header(function, cfg)?; let mut jump_value = get_induction_variable(&function.dfg, unroll_into)?; - // TODO: WIP - while let (Some(context), loop_header_id) = + while let Some((context, loop_header_id)) = self.unroll_header(function, unroll_into, jump_value)? { - // TODO: WIP - // context.visited_blocks.insert(loop_header_id); - (unroll_into, jump_value) = context.unroll_loop_iteration(loop_header_id); } @@ -606,23 +602,19 @@ impl Loop { /// Unrolls the header block of the loop. This is the block that dominates all other blocks in the /// loop and contains the jmpif instruction that lets us know if we should continue looping. - /// Returns Some(iteration context) if we should perform another iteration. - /// // TODO: WIP - /// and the loop header's BasicBlockId + /// Returns Some((iteration context, loop_header_id)) if we should perform another iteration. fn unroll_header<'a>( &'a self, function: &'a mut Function, unroll_into: BasicBlockId, induction_value: ValueId, - // TODO: WIP - ) -> Result<(Option>, BasicBlockId), CallStack> { + ) -> Result, BasicBlockId)>, CallStack> { // We insert into a fresh block first and move instructions into the unroll_into block later // only once we verify the jmpif instruction has a constant condition. If it does not, we can // just discard this fresh block and leave the loop unmodified. let fresh_block = function.dfg.make_block(); let mut context = LoopIteration::new(function, self, fresh_block, self.header); - // TODO: WIP let loop_header_id = context.source_block; let source_block = &context.dfg()[loop_header_id]; assert_eq!(source_block.parameters().len(), 1, "Expected only 1 argument in loop header"); @@ -664,11 +656,11 @@ impl Loop { // have no more loops to unroll, because that block was not part of the loop itself, // ie. it wasn't between `loop_header` and `loop_body`. Otherwise we have the `loop_body` // in `source_block` and can unroll that into the destination. - // TODO: WIP - Ok(( - self.blocks.contains(&context.source_block).then_some(context), - loop_header_id, - )) + Ok(self + .blocks + .contains(&context.source_block) + .then_some(context) + .map(|iteration_context| (iteration_context, loop_header_id))) } else { // If this case is reached the loop either uses non-constant indices or we need // another pass, such as mem2reg to resolve them to constants. @@ -984,9 +976,7 @@ struct LoopIteration<'f> { original_blocks: HashMap, visited_blocks: HashSet, - // TODO: WIP - // #[cfg(test)] - // loop_header_id: Option, + /// Has `unroll_loop_iteration` reached the `loop_header_id` encountered_loop_header: bool, insert_block: BasicBlockId, @@ -1014,7 +1004,6 @@ impl<'f> LoopIteration<'f> { blocks: HashMap::default(), original_blocks: HashMap::default(), visited_blocks: HashSet::default(), - // TODO: WIP encountered_loop_header: false, induction_value: None, @@ -1027,23 +1016,14 @@ impl<'f> LoopIteration<'f> { /// It is expected the terminator instructions are set up to branch into an empty block /// for further unrolling. When the loop is finished this will need to be mutated to /// jump to the end of the loop instead. - // TODO: WIP - fn unroll_loop_iteration(mut self, loop_header: BasicBlockId) -> (BasicBlockId, ValueId) { - // TODO: WIP - // Mark the loop header as visited - self.visited_blocks.insert(loop_header); - + fn unroll_loop_iteration(mut self, loop_header_id: BasicBlockId) -> (BasicBlockId, ValueId) { // Kick off the unrolling from the initial source block. let mut next_blocks = self.unroll_loop_block(); while let Some(block) = next_blocks.pop() { self.insert_block = block; self.source_block = self.get_original_block(block); - - // TODO: WIP - // if cfg!(test) { - self.encountered_loop_header |= self.visited_blocks.contains(&self.source_block); - // } + self.encountered_loop_header |= loop_header_id == self.source_block; if !self.visited_blocks.contains(&self.source_block) { let mut blocks = self.unroll_loop_block(); @@ -1056,13 +1036,10 @@ impl<'f> LoopIteration<'f> { .induction_value .expect("Expected to find the induction variable by end of loop iteration"); - // TODO: WIP - // if cfg!(test) { assert!( self.encountered_loop_header, "expected to encounter loop header when visiting blocks" ); - // } (end_block, induction_value) } From fd0e8cb516047ec09ef2fe6e61599ab3ee81a077 Mon Sep 17 00:00:00 2001 From: "Michael J. Klein" Date: Thu, 23 Oct 2025 14:42:30 -0400 Subject: [PATCH 3/3] . --- compiler/noirc_evaluator/src/ssa/opt/unrolling.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs index dcaca7c3fce..0e097a012ec 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs @@ -976,7 +976,7 @@ struct LoopIteration<'f> { original_blocks: HashMap, visited_blocks: HashSet, - /// Has `unroll_loop_iteration` reached the `loop_header_id` + /// Has `unroll_loop_iteration` reached the `loop_header_id`? encountered_loop_header: bool, insert_block: BasicBlockId,