diff --git a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs index 661109c1786..267dc6a3c20 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs @@ -40,54 +40,47 @@ use fxhash::FxHashMap as HashMap; impl Ssa { /// Loop unrolling can return errors, since ACIR functions need to be fully unrolled. /// This meta-pass will keep trying to unroll loops and simplifying the SSA until no more errors are found. + #[tracing::instrument(level = "trace", skip(ssa))] pub(crate) fn unroll_loops_iteratively(mut ssa: Ssa) -> Result { - // Try to unroll loops first: - let mut unroll_errors; - (ssa, unroll_errors) = ssa.try_to_unroll_loops(); - - // Keep unrolling until no more errors are found - while !unroll_errors.is_empty() { - let prev_unroll_err_count = unroll_errors.len(); - - // Simplify the SSA before retrying - - // Do a mem2reg after the last unroll to aid simplify_cfg - ssa = ssa.mem2reg(); - ssa = ssa.simplify_cfg(); - // Do another mem2reg after simplify_cfg to aid the next unroll - ssa = ssa.mem2reg(); - - // Unroll again - (ssa, unroll_errors) = ssa.try_to_unroll_loops(); - // If we didn't manage to unroll any more loops, exit - if unroll_errors.len() >= prev_unroll_err_count { - return Err(unroll_errors.swap_remove(0)); + let acir_functions = ssa.functions.iter_mut().filter(|(_, func)| { + // Loop unrolling in brillig can lead to a code explosion currently. This can + // also be true for ACIR, but we have no alternative to unrolling in ACIR. + // Brillig also generally prefers smaller code rather than faster code. + !matches!(func.runtime(), RuntimeType::Brillig(_)) + }); + + for (_, function) in acir_functions { + // Try to unroll loops first: + let mut unroll_errors = function.try_to_unroll_loops(); + + // Keep unrolling until no more errors are found + while !unroll_errors.is_empty() { + let prev_unroll_err_count = unroll_errors.len(); + + // Simplify the SSA before retrying + + // Do a mem2reg after the last unroll to aid simplify_cfg + function.mem2reg(); + function.simplify_function(); + // Do another mem2reg after simplify_cfg to aid the next unroll + function.mem2reg(); + + // Unroll again + unroll_errors = function.try_to_unroll_loops(); + // If we didn't manage to unroll any more loops, exit + if unroll_errors.len() >= prev_unroll_err_count { + return Err(unroll_errors.swap_remove(0)); + } } } - Ok(ssa) - } - /// Tries to unroll all loops in each SSA function. - /// If any loop cannot be unrolled, it is left as-is or in a partially unrolled state. - /// Returns the ssa along with all unrolling errors encountered - #[tracing::instrument(level = "trace", skip(self))] - pub(crate) fn try_to_unroll_loops(mut self) -> (Ssa, Vec) { - let mut errors = vec![]; - for function in self.functions.values_mut() { - function.try_to_unroll_loops(&mut errors); - } - (self, errors) + Ok(ssa) } } impl Function { - pub(crate) fn try_to_unroll_loops(&mut self, errors: &mut Vec) { - // Loop unrolling in brillig can lead to a code explosion currently. This can - // also be true for ACIR, but we have no alternative to unrolling in ACIR. - // Brillig also generally prefers smaller code rather than faster code. - if !matches!(self.runtime(), RuntimeType::Brillig(_)) { - errors.extend(find_all_loops(self).unroll_each_loop(self)); - } + fn try_to_unroll_loops(&mut self) -> Vec { + find_all_loops(self).unroll_each_loop(self) } } @@ -507,11 +500,26 @@ impl<'f> LoopIteration<'f> { #[cfg(test)] mod tests { - use crate::ssa::{ - function_builder::FunctionBuilder, - ir::{instruction::BinaryOp, map::Id, types::Type}, + use crate::{ + errors::RuntimeError, + ssa::{ + function_builder::FunctionBuilder, + ir::{instruction::BinaryOp, map::Id, types::Type}, + }, }; + use super::Ssa; + + /// Tries to unroll all loops in each SSA function. + /// If any loop cannot be unrolled, it is left as-is or in a partially unrolled state. + fn try_to_unroll_loops(mut ssa: Ssa) -> (Ssa, Vec) { + let mut errors = vec![]; + for function in ssa.functions.values_mut() { + errors.extend(function.try_to_unroll_loops()); + } + (ssa, errors) + } + #[test] fn unroll_nested_loops() { // fn main() { @@ -630,7 +638,7 @@ mod tests { // } // The final block count is not 1 because unrolling creates some unnecessary jmps. // If a simplify cfg pass is ran afterward, the expected block count will be 1. - let (ssa, errors) = ssa.try_to_unroll_loops(); + let (ssa, errors) = try_to_unroll_loops(ssa); assert_eq!(errors.len(), 0, "All loops should be unrolled"); assert_eq!(ssa.main().reachable_blocks().len(), 5); } @@ -680,7 +688,7 @@ mod tests { assert_eq!(ssa.main().reachable_blocks().len(), 4); // Expected that we failed to unroll the loop - let (_, errors) = ssa.try_to_unroll_loops(); + let (_, errors) = try_to_unroll_loops(ssa); assert_eq!(errors.len(), 1, "Expected to fail to unroll loop"); } }