From 0ba9fb8333842e1c05f6f4b1825fd838f507ceab Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Tue, 5 Nov 2024 17:03:57 -0300 Subject: [PATCH 01/16] Add a skeleton pass to inline const brillig calls --- compiler/noirc_evaluator/src/ssa.rs | 1 + .../src/ssa/opt/inline_const_brillig_calls.rs | 8 ++++++++ compiler/noirc_evaluator/src/ssa/opt/mod.rs | 1 + 3 files changed, 10 insertions(+) create mode 100644 compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs diff --git a/compiler/noirc_evaluator/src/ssa.rs b/compiler/noirc_evaluator/src/ssa.rs index 0c4e42f09ef..637aa1db8ff 100644 --- a/compiler/noirc_evaluator/src/ssa.rs +++ b/compiler/noirc_evaluator/src/ssa.rs @@ -98,6 +98,7 @@ pub(crate) fn optimize_into_acir( .run_pass(Ssa::separate_runtime, "After Runtime Separation:") .run_pass(Ssa::resolve_is_unconstrained, "After Resolving IsUnconstrained:") .run_pass(|ssa| ssa.inline_functions(options.inliner_aggressiveness), "After Inlining:") + .run_pass(Ssa::inline_const_brillig_calls, "After Inlining Const Brillig Calls:") // Run mem2reg with the CFG separated into blocks .run_pass(Ssa::mem2reg, "After Mem2Reg (1st):") .run_pass(Ssa::simplify_cfg, "After Simplifying (1st):") diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs new file mode 100644 index 00000000000..fc8769dcfc4 --- /dev/null +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -0,0 +1,8 @@ +use crate::ssa::Ssa; + +impl Ssa { + #[tracing::instrument(level = "trace", skip(self))] + pub(crate) fn inline_const_brillig_calls(self) -> Self { + self + } +} diff --git a/compiler/noirc_evaluator/src/ssa/opt/mod.rs b/compiler/noirc_evaluator/src/ssa/opt/mod.rs index f3dbd58fa69..c9aa5e17efe 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/mod.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/mod.rs @@ -10,6 +10,7 @@ mod constant_folding; mod defunctionalize; mod die; pub(crate) mod flatten_cfg; +mod inline_const_brillig_calls; mod inlining; mod mem2reg; mod normalize_value_ids; From 4cf74f65757d20c1ae01ca7f5c42c207fe68da80 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Wed, 6 Nov 2024 12:08:44 -0300 Subject: [PATCH 02/16] Try to inline brillig calls with all constant arguments --- .../src/ssa/opt/flatten_cfg.rs | 5 +- .../src/ssa/opt/inline_const_brillig_calls.rs | 164 +++++++++++++++++- 2 files changed, 166 insertions(+), 3 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs index 984f639df00..bc7b8f2d88a 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs @@ -250,7 +250,10 @@ struct ConditionalContext { call_stack: CallStack, } -fn flatten_function_cfg(function: &mut Function, no_predicates: &HashMap) { +pub(crate) fn flatten_function_cfg( + function: &mut Function, + no_predicates: &HashMap, +) { // This pass may run forever on a brillig function. // Analyze will check if the predecessors have been processed and push the block to the back of // the queue. This loops forever if there are still any loops present in the program. diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index fc8769dcfc4..59744c9d641 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -1,8 +1,168 @@ -use crate::ssa::Ssa; +use std::collections::BTreeMap; + +use fxhash::FxHashMap; +use noirc_frontend::monomorphization::ast::InlineType; + +use crate::{ + errors::RuntimeError, + ssa::{ + ir::{ + function::{Function, FunctionId, RuntimeType}, + instruction::{Instruction, TerminatorInstruction}, + value::Value, + }, + opt::flatten_cfg::flatten_function_cfg, + Ssa, + }, +}; impl Ssa { #[tracing::instrument(level = "trace", skip(self))] - pub(crate) fn inline_const_brillig_calls(self) -> Self { + pub(crate) fn inline_const_brillig_calls(mut self) -> Self { + // Collect all brillig functions so that later we can find them when processing a call instruction + let mut brillig_functions = BTreeMap::::new(); + for (func_id, func) in &self.functions { + if let RuntimeType::Brillig(..) = func.runtime() { + let cloned_function = Function::clone_with_id(*func_id, func); + brillig_functions.insert(*func_id, cloned_function); + } + } + + for func in self.functions.values_mut() { + func.inline_const_brillig_calls(&brillig_functions); + } self } } + +impl Function { + pub(crate) fn inline_const_brillig_calls( + &mut self, + brillig_functions: &BTreeMap, + ) { + let reachable_block_ids = self.reachable_blocks(); + + for block_id in reachable_block_ids { + let block = &mut self.dfg[block_id]; + let instruction_ids = block.take_instructions(); + + for instruction_id in instruction_ids { + let instruction = &self.dfg[instruction_id]; + let Instruction::Call { func: func_id, arguments } = instruction else { + self.dfg[block_id].instructions_mut().push(instruction_id); + continue; + }; + + let func_value = &self.dfg[*func_id]; + let Value::Function(func_id) = func_value else { + self.dfg[block_id].instructions_mut().push(instruction_id); + continue; + }; + + let Some(function) = brillig_functions.get(func_id) else { + self.dfg[block_id].instructions_mut().push(instruction_id); + continue; + }; + + if !arguments.iter().all(|argument| self.dfg.is_constant(*argument)) { + self.dfg[block_id].instructions_mut().push(instruction_id); + continue; + } + + // The function we have is already a copy of the original function, but we need to clone + // it again because there might be multiple calls to the same brillig function. + let mut function = Function::clone_with_id(*func_id, function); + + // Find the entry block and remove its parameters + let entry_block_id = function.entry_block(); + let entry_block = &mut function.dfg[entry_block_id]; + let entry_block_parameters = entry_block.take_parameters(); + + assert_eq!(arguments.len(), entry_block_parameters.len()); + + // Replace the ValueId of parameters with the ValueId of arguments + for (parameter_id, argument_id) in entry_block_parameters.iter().zip(arguments) { + // Lookup the argument in the current function and insert it in the function copy + let argument = &self.dfg[*argument_id]; + let new_argument_id = function.dfg.make_value(argument.clone()); + + function.dfg.set_value_from_id(*parameter_id, new_argument_id); + } + + // Try to fully optimize the function. If we can't, we can't inline it's constant value. + if optimize(&mut function).is_err() { + self.dfg[block_id].instructions_mut().push(instruction_id); + continue; + } + + let entry_block = &mut function.dfg[entry_block_id]; + + // If the entry block has instructions, we can't inline it (we need a terminator) + if !entry_block.instructions().is_empty() { + self.dfg[block_id].instructions_mut().push(instruction_id); + continue; + } + + let terminator = entry_block.take_terminator(); + let TerminatorInstruction::Return { return_values, call_stack: _ } = terminator + else { + self.dfg[block_id].instructions_mut().push(instruction_id); + continue; + }; + + // Sanity check: make sure all returned values are constant + if !return_values.iter().all(|value_id| function.dfg.is_constant(*value_id)) { + self.dfg[block_id].instructions_mut().push(instruction_id); + continue; + } + + // Replace the instruction results with the constant values we got + let current_results = self.dfg.instruction_results(instruction_id).to_vec(); + assert_eq!(return_values.len(), current_results.len()); + + for (current_result, return_value) in current_results.iter().zip(return_values) { + let return_value = &function.dfg[return_value]; + let new_result_id = self.dfg.make_value(return_value.clone()); + self.dfg.set_value_from_id(*current_result, new_result_id); + } + } + } + } +} + +/// Optimizes a function by running the same passes as `optimize_into_acir` +/// after the `inline_const_brillig_calls` pass. +/// The function is changed to be an ACIR function so the function can potentially +/// be optimized into a single return terminator. +fn optimize(function: &mut Function) -> Result<(), RuntimeError> { + function.set_runtime(RuntimeType::Acir(InlineType::InlineAlways)); + + function.mem2reg(); + function.simplify_function(); + function.as_slice_optimization(); + function.evaluate_static_assert_and_assert_constant()?; + + let mut errors = Vec::new(); + function.try_to_unroll_loops(&mut errors); + if !errors.is_empty() { + return Err(errors.swap_remove(0)); + } + + function.simplify_function(); + + let mut no_predicates = FxHashMap::default(); + no_predicates.insert(function.id(), function.is_no_predicates()); + flatten_function_cfg(function, &no_predicates); + + function.remove_bit_shifts(); + function.mem2reg(); + function.remove_if_else(); + function.constant_fold(false); + function.remove_enable_side_effects(); + function.constant_fold(true); + function.dead_instruction_elimination(true); + function.simplify_function(); + function.array_set_optimization(); + + Ok(()) +} From 726f3bd1cc2cf7e7e743bbb980fbf7471f4c0799 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Wed, 6 Nov 2024 13:39:53 -0300 Subject: [PATCH 03/16] Fully copy constants --- .../src/ssa/opt/inline_const_brillig_calls.rs | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index 59744c9d641..4717a528359 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -9,7 +9,7 @@ use crate::{ ir::{ function::{Function, FunctionId, RuntimeType}, instruction::{Instruction, TerminatorInstruction}, - value::Value, + value::{Value, ValueId}, }, opt::flatten_cfg::flatten_function_cfg, Ssa, @@ -83,9 +83,8 @@ impl Function { // Replace the ValueId of parameters with the ValueId of arguments for (parameter_id, argument_id) in entry_block_parameters.iter().zip(arguments) { // Lookup the argument in the current function and insert it in the function copy - let argument = &self.dfg[*argument_id]; - let new_argument_id = function.dfg.make_value(argument.clone()); - + let new_argument_id = + self.copy_constant_to_function(*argument_id, &mut function); function.dfg.set_value_from_id(*parameter_id, new_argument_id); } @@ -128,6 +127,20 @@ impl Function { } } } + + fn copy_constant_to_function(&self, argument_id: ValueId, function: &mut Function) -> ValueId { + if let Some((constant, typ)) = self.dfg.get_numeric_constant_with_type(argument_id) { + function.dfg.make_constant(constant, typ) + } else if let Some((constants, typ)) = self.dfg.get_array_constant(argument_id) { + let new_constants = constants + .iter() + .map(|constant_id| self.copy_constant_to_function(*constant_id, function)) + .collect(); + function.dfg.make_array(new_constants, typ) + } else { + unreachable!("A constant should be either a numeric constant or an array constant") + } + } } /// Optimizes a function by running the same passes as `optimize_into_acir` From c7dbac11f749b7e71225ab87336c19abdea96ec5 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Wed, 6 Nov 2024 21:38:08 -0300 Subject: [PATCH 04/16] Replace IDs at the end --- .../src/ssa/opt/inline_const_brillig_calls.rs | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index 4717a528359..efe7602d897 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -40,32 +40,30 @@ impl Function { &mut self, brillig_functions: &BTreeMap, ) { + let mut ids_to_replace = vec![]; + let reachable_block_ids = self.reachable_blocks(); for block_id in reachable_block_ids { - let block = &mut self.dfg[block_id]; - let instruction_ids = block.take_instructions(); - - for instruction_id in instruction_ids { + let block = &self.dfg[block_id]; + for instruction_id in block.instructions() { + let instruction_id = *instruction_id; let instruction = &self.dfg[instruction_id]; + let Instruction::Call { func: func_id, arguments } = instruction else { - self.dfg[block_id].instructions_mut().push(instruction_id); continue; }; let func_value = &self.dfg[*func_id]; let Value::Function(func_id) = func_value else { - self.dfg[block_id].instructions_mut().push(instruction_id); continue; }; let Some(function) = brillig_functions.get(func_id) else { - self.dfg[block_id].instructions_mut().push(instruction_id); continue; }; if !arguments.iter().all(|argument| self.dfg.is_constant(*argument)) { - self.dfg[block_id].instructions_mut().push(instruction_id); continue; } @@ -90,7 +88,6 @@ impl Function { // Try to fully optimize the function. If we can't, we can't inline it's constant value. if optimize(&mut function).is_err() { - self.dfg[block_id].instructions_mut().push(instruction_id); continue; } @@ -98,20 +95,17 @@ impl Function { // If the entry block has instructions, we can't inline it (we need a terminator) if !entry_block.instructions().is_empty() { - self.dfg[block_id].instructions_mut().push(instruction_id); continue; } let terminator = entry_block.take_terminator(); let TerminatorInstruction::Return { return_values, call_stack: _ } = terminator else { - self.dfg[block_id].instructions_mut().push(instruction_id); continue; }; // Sanity check: make sure all returned values are constant if !return_values.iter().all(|value_id| function.dfg.is_constant(*value_id)) { - self.dfg[block_id].instructions_mut().push(instruction_id); continue; } @@ -121,11 +115,15 @@ impl Function { for (current_result, return_value) in current_results.iter().zip(return_values) { let return_value = &function.dfg[return_value]; - let new_result_id = self.dfg.make_value(return_value.clone()); - self.dfg.set_value_from_id(*current_result, new_result_id); + ids_to_replace.push((*current_result, return_value.clone())); } } } + + for (current_result, return_value) in ids_to_replace { + let new_result_id = self.dfg.make_value(return_value); + self.dfg.set_value_from_id(current_result, new_result_id); + } } fn copy_constant_to_function(&self, argument_id: ValueId, function: &mut Function) -> ValueId { From 16dd32fc8a2590b44a2a9b67da92d67dd622f8c7 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Wed, 6 Nov 2024 21:59:16 -0300 Subject: [PATCH 05/16] Trying something... --- compiler/noirc_evaluator/src/ssa/opt/unrolling.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs index 661109c1786..97c45fd4ddd 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs @@ -218,7 +218,7 @@ fn unroll_loop( cfg: &ControlFlowGraph, loop_: &Loop, ) -> Result<(), CallStack> { - let mut unroll_into = get_pre_header(cfg, loop_); + let mut unroll_into = get_pre_header(cfg, loop_)?; let mut jump_value = get_induction_variable(function, unroll_into)?; while let Some(context) = unroll_loop_header(function, loop_, unroll_into, jump_value)? { @@ -233,14 +233,17 @@ fn unroll_loop( /// The loop pre-header is the block that comes before the loop begins. Generally a header block /// is expected to have 2 predecessors: the pre-header and the final block of the loop which jumps /// back to the beginning. -fn get_pre_header(cfg: &ControlFlowGraph, loop_: &Loop) -> BasicBlockId { +fn get_pre_header(cfg: &ControlFlowGraph, loop_: &Loop) -> Result { let mut pre_header = cfg .predecessors(loop_.header) .filter(|predecessor| *predecessor != loop_.back_edge_start) .collect::>(); - assert_eq!(pre_header.len(), 1); - pre_header.remove(0) + if pre_header.len() == 1 { + Ok(pre_header.remove(0)) + } else { + Err(CallStack::new()) + } } /// Return the induction value of the current iteration of the loop, from the given block's jmp arguments. From dfefc0a679f7c4d6ab6ede52133d108cb50a3b8c Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Thu, 7 Nov 2024 07:01:22 -0300 Subject: [PATCH 06/16] Found the bug where constants weren't being fully copied from the optimized function --- .../src/ssa/opt/inline_const_brillig_calls.rs | 63 ++++++++++--------- 1 file changed, 35 insertions(+), 28 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index efe7602d897..4a0eedabfe5 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -40,30 +40,31 @@ impl Function { &mut self, brillig_functions: &BTreeMap, ) { - let mut ids_to_replace = vec![]; + for block_id in self.reachable_blocks() { + let block = &mut self.dfg[block_id]; + let instruction_ids = block.take_instructions(); - let reachable_block_ids = self.reachable_blocks(); - - for block_id in reachable_block_ids { - let block = &self.dfg[block_id]; - for instruction_id in block.instructions() { - let instruction_id = *instruction_id; + for instruction_id in instruction_ids { let instruction = &self.dfg[instruction_id]; let Instruction::Call { func: func_id, arguments } = instruction else { + self.dfg[block_id].instructions_mut().push(instruction_id); continue; }; let func_value = &self.dfg[*func_id]; let Value::Function(func_id) = func_value else { + self.dfg[block_id].instructions_mut().push(instruction_id); continue; }; let Some(function) = brillig_functions.get(func_id) else { + self.dfg[block_id].instructions_mut().push(instruction_id); continue; }; if !arguments.iter().all(|argument| self.dfg.is_constant(*argument)) { + self.dfg[block_id].instructions_mut().push(instruction_id); continue; } @@ -82,12 +83,13 @@ impl Function { for (parameter_id, argument_id) in entry_block_parameters.iter().zip(arguments) { // Lookup the argument in the current function and insert it in the function copy let new_argument_id = - self.copy_constant_to_function(*argument_id, &mut function); + copy_constant_to_function(self, &mut function, *argument_id); function.dfg.set_value_from_id(*parameter_id, new_argument_id); } // Try to fully optimize the function. If we can't, we can't inline it's constant value. if optimize(&mut function).is_err() { + self.dfg[block_id].instructions_mut().push(instruction_id); continue; } @@ -95,17 +97,20 @@ impl Function { // If the entry block has instructions, we can't inline it (we need a terminator) if !entry_block.instructions().is_empty() { + self.dfg[block_id].instructions_mut().push(instruction_id); continue; } let terminator = entry_block.take_terminator(); let TerminatorInstruction::Return { return_values, call_stack: _ } = terminator else { + self.dfg[block_id].instructions_mut().push(instruction_id); continue; }; // Sanity check: make sure all returned values are constant if !return_values.iter().all(|value_id| function.dfg.is_constant(*value_id)) { + self.dfg[block_id].instructions_mut().push(instruction_id); continue; } @@ -113,31 +118,33 @@ impl Function { let current_results = self.dfg.instruction_results(instruction_id).to_vec(); assert_eq!(return_values.len(), current_results.len()); - for (current_result, return_value) in current_results.iter().zip(return_values) { - let return_value = &function.dfg[return_value]; - ids_to_replace.push((*current_result, return_value.clone())); + for (current_result_id, return_value_id) in + current_results.iter().zip(return_values) + { + let new_return_value_id = + copy_constant_to_function(&function, self, return_value_id); + self.dfg.set_value_from_id(*current_result_id, new_return_value_id); } } } - - for (current_result, return_value) in ids_to_replace { - let new_result_id = self.dfg.make_value(return_value); - self.dfg.set_value_from_id(current_result, new_result_id); - } } +} - fn copy_constant_to_function(&self, argument_id: ValueId, function: &mut Function) -> ValueId { - if let Some((constant, typ)) = self.dfg.get_numeric_constant_with_type(argument_id) { - function.dfg.make_constant(constant, typ) - } else if let Some((constants, typ)) = self.dfg.get_array_constant(argument_id) { - let new_constants = constants - .iter() - .map(|constant_id| self.copy_constant_to_function(*constant_id, function)) - .collect(); - function.dfg.make_array(new_constants, typ) - } else { - unreachable!("A constant should be either a numeric constant or an array constant") - } +fn copy_constant_to_function( + from_function: &Function, + to_function: &mut Function, + argument_id: ValueId, +) -> ValueId { + if let Some((constant, typ)) = from_function.dfg.get_numeric_constant_with_type(argument_id) { + to_function.dfg.make_constant(constant, typ) + } else if let Some((constants, typ)) = from_function.dfg.get_array_constant(argument_id) { + let new_constants = constants + .iter() + .map(|constant_id| copy_constant_to_function(from_function, to_function, *constant_id)) + .collect(); + to_function.dfg.make_array(new_constants, typ) + } else { + unreachable!("A constant should be either a numeric constant or an array constant") } } From 521d25486b427bdb3b5db5cc45127b809bbcd022 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Thu, 7 Nov 2024 07:16:59 -0300 Subject: [PATCH 07/16] Simplify, and add some comments --- .../src/ssa/opt/inline_const_brillig_calls.rs | 155 +++++++++--------- 1 file changed, 79 insertions(+), 76 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index 4a0eedabfe5..d92f7bf58c2 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -1,3 +1,4 @@ +//! This pass tries to inline calls to brillig functions that have all constant arguments. use std::collections::BTreeMap; use fxhash::FxHashMap; @@ -8,7 +9,7 @@ use crate::{ ssa::{ ir::{ function::{Function, FunctionId, RuntimeType}, - instruction::{Instruction, TerminatorInstruction}, + instruction::{Instruction, InstructionId, TerminatorInstruction}, value::{Value, ValueId}, }, opt::flatten_cfg::flatten_function_cfg, @@ -41,103 +42,105 @@ impl Function { brillig_functions: &BTreeMap, ) { for block_id in self.reachable_blocks() { - let block = &mut self.dfg[block_id]; - let instruction_ids = block.take_instructions(); - - for instruction_id in instruction_ids { - let instruction = &self.dfg[instruction_id]; - - let Instruction::Call { func: func_id, arguments } = instruction else { - self.dfg[block_id].instructions_mut().push(instruction_id); - continue; - }; - - let func_value = &self.dfg[*func_id]; - let Value::Function(func_id) = func_value else { - self.dfg[block_id].instructions_mut().push(instruction_id); - continue; - }; - - let Some(function) = brillig_functions.get(func_id) else { - self.dfg[block_id].instructions_mut().push(instruction_id); - continue; - }; - - if !arguments.iter().all(|argument| self.dfg.is_constant(*argument)) { + for instruction_id in self.dfg[block_id].take_instructions() { + if !self.optimize_const_brillig_call(instruction_id, brillig_functions) { self.dfg[block_id].instructions_mut().push(instruction_id); - continue; } + } + } + } - // The function we have is already a copy of the original function, but we need to clone - // it again because there might be multiple calls to the same brillig function. - let mut function = Function::clone_with_id(*func_id, function); + /// Tries to optimize an instruction if it's a call that points to a brillig function, + /// and all its arguments are constant. If the optimization is successful, returns true. + /// Otherwise returns false. The given instruction is not removed from the function. + fn optimize_const_brillig_call( + &mut self, + instruction_id: InstructionId, + brillig_functions: &BTreeMap, + ) -> bool { + let instruction = &self.dfg[instruction_id]; + let Instruction::Call { func: func_id, arguments } = instruction else { + return false; + }; + + let func_value = &self.dfg[*func_id]; + let Value::Function(func_id) = func_value else { + return false; + }; + + let Some(function) = brillig_functions.get(func_id) else { + return false; + }; + + if !arguments.iter().all(|argument| self.dfg.is_constant(*argument)) { + return false; + } - // Find the entry block and remove its parameters - let entry_block_id = function.entry_block(); - let entry_block = &mut function.dfg[entry_block_id]; - let entry_block_parameters = entry_block.take_parameters(); + // The function we have is already a copy of the original function, but we need to clone + // it again because there might be multiple calls to the same brillig function. + let mut function = Function::clone_with_id(*func_id, function); - assert_eq!(arguments.len(), entry_block_parameters.len()); + // Find the entry block and remove its parameters + let entry_block_id = function.entry_block(); + let entry_block = &mut function.dfg[entry_block_id]; + let entry_block_parameters = entry_block.take_parameters(); - // Replace the ValueId of parameters with the ValueId of arguments - for (parameter_id, argument_id) in entry_block_parameters.iter().zip(arguments) { - // Lookup the argument in the current function and insert it in the function copy - let new_argument_id = - copy_constant_to_function(self, &mut function, *argument_id); - function.dfg.set_value_from_id(*parameter_id, new_argument_id); - } + assert_eq!(arguments.len(), entry_block_parameters.len()); - // Try to fully optimize the function. If we can't, we can't inline it's constant value. - if optimize(&mut function).is_err() { - self.dfg[block_id].instructions_mut().push(instruction_id); - continue; - } + // Replace the ValueId of parameters with the ValueId of arguments + for (parameter_id, argument_id) in entry_block_parameters.iter().zip(arguments) { + // Lookup the argument in the current function and insert it in the function copy + let new_argument_id = copy_constant_to_function(self, &mut function, *argument_id); + function.dfg.set_value_from_id(*parameter_id, new_argument_id); + } - let entry_block = &mut function.dfg[entry_block_id]; + // Try to fully optimize the function. If we can't, we can't inline it's constant value. + if optimize(&mut function).is_err() { + return false; + } - // If the entry block has instructions, we can't inline it (we need a terminator) - if !entry_block.instructions().is_empty() { - self.dfg[block_id].instructions_mut().push(instruction_id); - continue; - } + let entry_block = &mut function.dfg[entry_block_id]; - let terminator = entry_block.take_terminator(); - let TerminatorInstruction::Return { return_values, call_stack: _ } = terminator - else { - self.dfg[block_id].instructions_mut().push(instruction_id); - continue; - }; + // If the entry block has instructions, we can't inline it (we need a terminator) + if !entry_block.instructions().is_empty() { + return false; + } - // Sanity check: make sure all returned values are constant - if !return_values.iter().all(|value_id| function.dfg.is_constant(*value_id)) { - self.dfg[block_id].instructions_mut().push(instruction_id); - continue; - } + let terminator = entry_block.take_terminator(); + let TerminatorInstruction::Return { return_values, call_stack: _ } = terminator else { + return false; + }; - // Replace the instruction results with the constant values we got - let current_results = self.dfg.instruction_results(instruction_id).to_vec(); - assert_eq!(return_values.len(), current_results.len()); + // Sanity check: make sure all returned values are constant + if !return_values.iter().all(|value_id| function.dfg.is_constant(*value_id)) { + return false; + } - for (current_result_id, return_value_id) in - current_results.iter().zip(return_values) - { - let new_return_value_id = - copy_constant_to_function(&function, self, return_value_id); - self.dfg.set_value_from_id(*current_result_id, new_return_value_id); - } - } + // Replace the instruction results with the constant values we got + let current_results = self.dfg.instruction_results(instruction_id).to_vec(); + assert_eq!(return_values.len(), current_results.len()); + + for (current_result_id, return_value_id) in current_results.iter().zip(return_values) { + let new_return_value_id = copy_constant_to_function(&function, self, return_value_id); + self.dfg.set_value_from_id(*current_result_id, new_return_value_id); } + + true } } +/// Copies a constant from one function to another. +/// Though it might seem we can just take a value out of `from_function` and call `make_value` on `to_function`, +/// if the constant is an array the values will still keep pointing to `from_function`. So, this function +/// recursively copies the array values too. fn copy_constant_to_function( from_function: &Function, to_function: &mut Function, - argument_id: ValueId, + constant_id: ValueId, ) -> ValueId { - if let Some((constant, typ)) = from_function.dfg.get_numeric_constant_with_type(argument_id) { + if let Some((constant, typ)) = from_function.dfg.get_numeric_constant_with_type(constant_id) { to_function.dfg.make_constant(constant, typ) - } else if let Some((constants, typ)) = from_function.dfg.get_array_constant(argument_id) { + } else if let Some((constants, typ)) = from_function.dfg.get_array_constant(constant_id) { let new_constants = constants .iter() .map(|constant_id| copy_constant_to_function(from_function, to_function, *constant_id)) From fd15065f2c36b6548e0f4220cdd1297aad0ac42a Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Thu, 7 Nov 2024 08:09:45 -0300 Subject: [PATCH 08/16] Run optimization passes by reusing existing code --- compiler/noirc_evaluator/src/ssa.rs | 104 +++++++++++------- .../src/ssa/opt/inline_const_brillig_calls.rs | 75 +++++++------ 2 files changed, 103 insertions(+), 76 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa.rs b/compiler/noirc_evaluator/src/ssa.rs index 637aa1db8ff..b6a18baf6ea 100644 --- a/compiler/noirc_evaluator/src/ssa.rs +++ b/compiler/noirc_evaluator/src/ssa.rs @@ -86,49 +86,14 @@ pub(crate) fn optimize_into_acir( let ssa_gen_span = span!(Level::TRACE, "ssa_generation"); let ssa_gen_span_guard = ssa_gen_span.enter(); - let mut ssa = SsaBuilder::new( + let builder = SsaBuilder::new( program, options.enable_ssa_logging, options.force_brillig_output, options.print_codegen_timings, &options.emit_ssa, - )? - .run_pass(Ssa::defunctionalize, "After Defunctionalization:") - .run_pass(Ssa::remove_paired_rc, "After Removing Paired rc_inc & rc_decs:") - .run_pass(Ssa::separate_runtime, "After Runtime Separation:") - .run_pass(Ssa::resolve_is_unconstrained, "After Resolving IsUnconstrained:") - .run_pass(|ssa| ssa.inline_functions(options.inliner_aggressiveness), "After Inlining:") - .run_pass(Ssa::inline_const_brillig_calls, "After Inlining Const Brillig Calls:") - // Run mem2reg with the CFG separated into blocks - .run_pass(Ssa::mem2reg, "After Mem2Reg (1st):") - .run_pass(Ssa::simplify_cfg, "After Simplifying (1st):") - .run_pass(Ssa::as_slice_optimization, "After `as_slice` optimization") - .try_run_pass( - Ssa::evaluate_static_assert_and_assert_constant, - "After `static_assert` and `assert_constant`:", - )? - .try_run_pass(Ssa::unroll_loops_iteratively, "After Unrolling:")? - .run_pass(Ssa::simplify_cfg, "After Simplifying (2nd):") - .run_pass(Ssa::flatten_cfg, "After Flattening:") - .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 (2nd):") - // 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. - // This pass must come immediately following `mem2reg` as the succeeding passes - // may create an SSA which inlining fails to handle. - .run_pass( - |ssa| ssa.inline_functions_with_no_predicates(options.inliner_aggressiveness), - "After Inlining:", - ) - .run_pass(Ssa::remove_if_else, "After Remove IfElse:") - .run_pass(Ssa::fold_constants, "After Constant Folding:") - .run_pass(Ssa::remove_enable_side_effects, "After EnableSideEffectsIf removal:") - .run_pass(Ssa::fold_constants_using_constraints, "After Constraint Folding:") - .run_pass(Ssa::dead_instruction_elimination, "After Dead Instruction Elimination:") - .run_pass(Ssa::simplify_cfg, "After Simplifying:") - .run_pass(Ssa::array_set_optimization, "After Array Set Optimizations:") - .finish(); + )?; + let mut ssa = optimize_ssa(builder, options.inliner_aggressiveness)?; let ssa_level_warnings = if options.skip_underconstrained_check { vec![] @@ -150,6 +115,69 @@ pub(crate) fn optimize_into_acir( Ok(ArtifactsAndWarnings(artifacts, ssa_level_warnings)) } +fn optimize_ssa(builder: SsaBuilder, inliner_aggressiveness: i64) -> Result { + let builder = builder + .run_pass(Ssa::defunctionalize, "After Defunctionalization:") + .run_pass(Ssa::remove_paired_rc, "After Removing Paired rc_inc & rc_decs:") + .run_pass(Ssa::separate_runtime, "After Runtime Separation:") + .run_pass(Ssa::resolve_is_unconstrained, "After Resolving IsUnconstrained:") + .run_pass(|ssa| ssa.inline_functions(inliner_aggressiveness), "After Inlining:") + .run_pass( + |ssa| ssa.inline_const_brillig_calls(inliner_aggressiveness), + "After Inlining Const Brillig Calls:", + ); + let ssa = optimize_ssa_after_inline_const_brillig_calls( + builder, + inliner_aggressiveness, + true, // inline functions with no predicates + )?; + Ok(ssa) +} + +fn optimize_ssa_after_inline_const_brillig_calls( + builder: SsaBuilder, + inliner_aggressiveness: i64, + inline_functions_with_no_predicates: bool, +) -> Result { + let builder = builder + // Run mem2reg with the CFG separated into blocks + .run_pass(Ssa::mem2reg, "After Mem2Reg (1st):") + .run_pass(Ssa::simplify_cfg, "After Simplifying (1st):") + .run_pass(Ssa::as_slice_optimization, "After `as_slice` optimization") + .try_run_pass( + Ssa::evaluate_static_assert_and_assert_constant, + "After `static_assert` and `assert_constant`:", + )? + .try_run_pass(Ssa::unroll_loops_iteratively, "After Unrolling:")? + .run_pass(Ssa::simplify_cfg, "After Simplifying (2nd):") + .run_pass(Ssa::flatten_cfg, "After Flattening:") + .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 (2nd):"); + let builder = if inline_functions_with_no_predicates { + // 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. + // This pass must come immediately following `mem2reg` as the succeeding passes + // may create an SSA which inlining fails to handle. + builder.run_pass( + |ssa| ssa.inline_functions_with_no_predicates(inliner_aggressiveness), + "After Inlining:", + ) + } else { + builder + }; + let ssa = builder + .run_pass(Ssa::remove_if_else, "After Remove IfElse:") + .run_pass(Ssa::fold_constants, "After Constant Folding:") + .run_pass(Ssa::remove_enable_side_effects, "After EnableSideEffectsIf removal:") + .run_pass(Ssa::fold_constants_using_constraints, "After Constraint Folding:") + .run_pass(Ssa::dead_instruction_elimination, "After Dead Instruction Elimination:") + .run_pass(Ssa::simplify_cfg, "After Simplifying:") + .run_pass(Ssa::array_set_optimization, "After Array Set Optimizations:") + .finish(); + Ok(ssa) +} + // Helper to time SSA passes fn time(name: &str, print_timings: bool, f: impl FnOnce() -> T) -> T { let start_time = chrono::Utc::now().time(); diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index d92f7bf58c2..1d91f541587 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -1,8 +1,8 @@ //! This pass tries to inline calls to brillig functions that have all constant arguments. use std::collections::BTreeMap; -use fxhash::FxHashMap; -use noirc_frontend::monomorphization::ast::InlineType; +use acvm::acir::circuit::ErrorSelector; +use noirc_frontend::{monomorphization::ast::InlineType, Type}; use crate::{ errors::RuntimeError, @@ -12,14 +12,15 @@ use crate::{ instruction::{Instruction, InstructionId, TerminatorInstruction}, value::{Value, ValueId}, }, - opt::flatten_cfg::flatten_function_cfg, - Ssa, + optimize_ssa_after_inline_const_brillig_calls, Ssa, SsaBuilder, }, }; impl Ssa { #[tracing::instrument(level = "trace", skip(self))] - pub(crate) fn inline_const_brillig_calls(mut self) -> Self { + pub(crate) fn inline_const_brillig_calls(mut self, inliner_aggressiveness: i64) -> Self { + let error_selector_to_type = &self.error_selector_to_type; + // Collect all brillig functions so that later we can find them when processing a call instruction let mut brillig_functions = BTreeMap::::new(); for (func_id, func) in &self.functions { @@ -30,7 +31,11 @@ impl Ssa { } for func in self.functions.values_mut() { - func.inline_const_brillig_calls(&brillig_functions); + func.inline_const_brillig_calls( + &brillig_functions, + inliner_aggressiveness, + error_selector_to_type, + ); } self } @@ -40,10 +45,17 @@ impl Function { pub(crate) fn inline_const_brillig_calls( &mut self, brillig_functions: &BTreeMap, + inliner_aggressiveness: i64, + error_selector_to_type: &BTreeMap, ) { for block_id in self.reachable_blocks() { for instruction_id in self.dfg[block_id].take_instructions() { - if !self.optimize_const_brillig_call(instruction_id, brillig_functions) { + if !self.optimize_const_brillig_call( + instruction_id, + brillig_functions, + inliner_aggressiveness, + error_selector_to_type, + ) { self.dfg[block_id].instructions_mut().push(instruction_id); } } @@ -57,6 +69,8 @@ impl Function { &mut self, instruction_id: InstructionId, brillig_functions: &BTreeMap, + inliner_aggressiveness: i64, + error_selector_to_type: &BTreeMap, ) -> bool { let instruction = &self.dfg[instruction_id]; let Instruction::Call { func: func_id, arguments } = instruction else { @@ -95,9 +109,10 @@ impl Function { } // Try to fully optimize the function. If we can't, we can't inline it's constant value. - if optimize(&mut function).is_err() { + let Ok(mut function) = optimize(function, inliner_aggressiveness, error_selector_to_type) + else { return false; - } + }; let entry_block = &mut function.dfg[entry_block_id]; @@ -155,35 +170,19 @@ fn copy_constant_to_function( /// after the `inline_const_brillig_calls` pass. /// The function is changed to be an ACIR function so the function can potentially /// be optimized into a single return terminator. -fn optimize(function: &mut Function) -> Result<(), RuntimeError> { +fn optimize( + mut function: Function, + inliner_aggressiveness: i64, + error_selector_to_type: &BTreeMap, +) -> Result { function.set_runtime(RuntimeType::Acir(InlineType::InlineAlways)); - function.mem2reg(); - function.simplify_function(); - function.as_slice_optimization(); - function.evaluate_static_assert_and_assert_constant()?; - - let mut errors = Vec::new(); - function.try_to_unroll_loops(&mut errors); - if !errors.is_empty() { - return Err(errors.swap_remove(0)); - } - - function.simplify_function(); - - let mut no_predicates = FxHashMap::default(); - no_predicates.insert(function.id(), function.is_no_predicates()); - flatten_function_cfg(function, &no_predicates); - - function.remove_bit_shifts(); - function.mem2reg(); - function.remove_if_else(); - function.constant_fold(false); - function.remove_enable_side_effects(); - function.constant_fold(true); - function.dead_instruction_elimination(true); - function.simplify_function(); - function.array_set_optimization(); - - Ok(()) + let ssa = Ssa::new(vec![function], error_selector_to_type.clone()); + let builder = SsaBuilder { ssa, print_ssa_passes: false, print_codegen_timings: false }; + let mut ssa = optimize_ssa_after_inline_const_brillig_calls( + builder, + inliner_aggressiveness, + false, // don't inline functions with no predicates + )?; + Ok(ssa.functions.pop_first().unwrap().1) } From 59aed416af71399cc1aa71ac4ee5760b291cbbf8 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Thu, 7 Nov 2024 10:18:55 -0300 Subject: [PATCH 09/16] Remove brillig functions that were inlined everywhere --- .../src/ssa/opt/inline_const_brillig_calls.rs | 36 +++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index 1d91f541587..67bd8483a81 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -1,5 +1,5 @@ //! This pass tries to inline calls to brillig functions that have all constant arguments. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashSet}; use acvm::acir::circuit::ErrorSelector; use noirc_frontend::{monomorphization::ast::InlineType, Type}; @@ -27,16 +27,40 @@ impl Ssa { if let RuntimeType::Brillig(..) = func.runtime() { let cloned_function = Function::clone_with_id(*func_id, func); brillig_functions.insert(*func_id, cloned_function); - } + }; } + // Keep track of which brillig functions we couldn't completely inline: we'll remove the ones we could. + let mut brillig_functions_we_could_not_inline = HashSet::new(); + for func in self.functions.values_mut() { func.inline_const_brillig_calls( &brillig_functions, + &mut brillig_functions_we_could_not_inline, inliner_aggressiveness, error_selector_to_type, ); } + + // Remove the brillig functions that are no longer called + for func_id in brillig_functions.keys() { + // We never want to remove the main function (it could be brillig if `--force-brillig` was given) + if self.main_id == *func_id { + continue; + } + + if brillig_functions_we_could_not_inline.contains(func_id) { + continue; + } + + // We also don't want to remove entry points + if self.entry_point_to_generated_index.contains_key(func_id) { + continue; + } + + self.functions.remove(func_id); + } + self } } @@ -45,6 +69,7 @@ impl Function { pub(crate) fn inline_const_brillig_calls( &mut self, brillig_functions: &BTreeMap, + brillig_functions_we_could_not_inline: &mut HashSet, inliner_aggressiveness: i64, error_selector_to_type: &BTreeMap, ) { @@ -53,6 +78,7 @@ impl Function { if !self.optimize_const_brillig_call( instruction_id, brillig_functions, + brillig_functions_we_could_not_inline, inliner_aggressiveness, error_selector_to_type, ) { @@ -69,6 +95,7 @@ impl Function { &mut self, instruction_id: InstructionId, brillig_functions: &BTreeMap, + brillig_functions_we_could_not_inline: &mut HashSet, inliner_aggressiveness: i64, error_selector_to_type: &BTreeMap, ) -> bool { @@ -87,6 +114,7 @@ impl Function { }; if !arguments.iter().all(|argument| self.dfg.is_constant(*argument)) { + brillig_functions_we_could_not_inline.insert(*func_id); return false; } @@ -111,6 +139,7 @@ impl Function { // Try to fully optimize the function. If we can't, we can't inline it's constant value. let Ok(mut function) = optimize(function, inliner_aggressiveness, error_selector_to_type) else { + brillig_functions_we_could_not_inline.insert(*func_id); return false; }; @@ -118,16 +147,19 @@ impl Function { // If the entry block has instructions, we can't inline it (we need a terminator) if !entry_block.instructions().is_empty() { + brillig_functions_we_could_not_inline.insert(*func_id); return false; } let terminator = entry_block.take_terminator(); let TerminatorInstruction::Return { return_values, call_stack: _ } = terminator else { + brillig_functions_we_could_not_inline.insert(*func_id); return false; }; // Sanity check: make sure all returned values are constant if !return_values.iter().all(|value_id| function.dfg.is_constant(*value_id)) { + brillig_functions_we_could_not_inline.insert(*func_id); return false; } From 5a37d5db496e6de729e7b6663b8995c63313f8ef Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Thu, 7 Nov 2024 10:58:04 -0300 Subject: [PATCH 10/16] Don't panic on loop unrolling if we can't find an induction variable --- compiler/noirc_evaluator/src/ssa/opt/unrolling.rs | 7 +++---- .../unconstrained_loop_with_break/Nargo.toml | 7 +++++++ .../unconstrained_loop_with_break/src/main.nr | 11 +++++++++++ 3 files changed, 21 insertions(+), 4 deletions(-) create mode 100644 test_programs/execution_success/unconstrained_loop_with_break/Nargo.toml create mode 100644 test_programs/execution_success/unconstrained_loop_with_break/src/main.nr diff --git a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs index 97c45fd4ddd..003877e85f7 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs @@ -222,7 +222,7 @@ fn unroll_loop( let mut jump_value = get_induction_variable(function, unroll_into)?; while let Some(context) = unroll_loop_header(function, loop_, unroll_into, jump_value)? { - let (last_block, last_value) = context.unroll_loop_iteration(); + let (last_block, last_value) = context.unroll_loop_iteration()?; unroll_into = last_block; jump_value = last_value; } @@ -367,7 +367,7 @@ 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) { + fn unroll_loop_iteration(mut self) -> Result<(BasicBlockId, ValueId), CallStack> { let mut next_blocks = self.unroll_loop_block(); while let Some(block) = next_blocks.pop() { @@ -380,8 +380,7 @@ impl<'f> LoopIteration<'f> { } } - self.induction_value - .expect("Expected to find the induction variable by end of loop iteration") + self.induction_value.ok_or_else(CallStack::new) } /// Unroll a single block in the current iteration of the loop diff --git a/test_programs/execution_success/unconstrained_loop_with_break/Nargo.toml b/test_programs/execution_success/unconstrained_loop_with_break/Nargo.toml new file mode 100644 index 00000000000..f789a4b7d5f --- /dev/null +++ b/test_programs/execution_success/unconstrained_loop_with_break/Nargo.toml @@ -0,0 +1,7 @@ +[package] +name = "unconstrained_loop_with_break" +version = "0.1.0" +type = "bin" +authors = [""] + +[dependencies] diff --git a/test_programs/execution_success/unconstrained_loop_with_break/src/main.nr b/test_programs/execution_success/unconstrained_loop_with_break/src/main.nr new file mode 100644 index 00000000000..ae28eb23c58 --- /dev/null +++ b/test_programs/execution_success/unconstrained_loop_with_break/src/main.nr @@ -0,0 +1,11 @@ +fn main() { + unsafe { foo() }; +} + +unconstrained fn foo() { + for i in 0..1 { + if i == 0 { + break; + } + } +} From dd29e65f0cae2708131dd704a7dd4cc43c4de864 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Thu, 7 Nov 2024 13:29:59 -0300 Subject: [PATCH 11/16] Improve comment --- .../src/ssa/opt/inline_const_brillig_calls.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index 67bd8483a81..e2bee50f63d 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -89,8 +89,9 @@ impl Function { } /// Tries to optimize an instruction if it's a call that points to a brillig function, - /// and all its arguments are constant. If the optimization is successful, returns true. - /// Otherwise returns false. The given instruction is not removed from the function. + /// and all its arguments are constant. If the optimization is successful, the + /// values returned by the brillig call are replaced by the constant values that the + /// function returns, and this method returns `true`. fn optimize_const_brillig_call( &mut self, instruction_id: InstructionId, From a68df4cd73a29d8c527ec26d3aafd811b55920f0 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Thu, 7 Nov 2024 13:31:33 -0300 Subject: [PATCH 12/16] Explain why we don't inline --- .../src/ssa/opt/inline_const_brillig_calls.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index e2bee50f63d..d8fbc8bd335 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -215,7 +215,11 @@ fn optimize( let mut ssa = optimize_ssa_after_inline_const_brillig_calls( builder, inliner_aggressiveness, - false, // don't inline functions with no predicates + // Don't inline functions with no predicates. + // The reason for this is that in this specific context the `Ssa` object only holds + // a single function. For inlining to work we need to know all other functions that + // exist (so we can inline them). Here we choose to skip this optimization for simplicity reasons. + false, )?; Ok(ssa.functions.pop_first().unwrap().1) } From 03b4d05e397b83563671befd4bd279bba5aff0da Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Thu, 7 Nov 2024 13:38:45 -0300 Subject: [PATCH 13/16] Turn `copy_constant_to_function` into a method --- .../src/ssa/opt/inline_const_brillig_calls.rs | 40 +++++++++---------- 1 file changed, 18 insertions(+), 22 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index d8fbc8bd335..24ebb5d4879 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -133,7 +133,7 @@ impl Function { // Replace the ValueId of parameters with the ValueId of arguments for (parameter_id, argument_id) in entry_block_parameters.iter().zip(arguments) { // Lookup the argument in the current function and insert it in the function copy - let new_argument_id = copy_constant_to_function(self, &mut function, *argument_id); + let new_argument_id = self.copy_constant_to_function(*argument_id, &mut function); function.dfg.set_value_from_id(*parameter_id, new_argument_id); } @@ -169,33 +169,29 @@ impl Function { assert_eq!(return_values.len(), current_results.len()); for (current_result_id, return_value_id) in current_results.iter().zip(return_values) { - let new_return_value_id = copy_constant_to_function(&function, self, return_value_id); + let new_return_value_id = function.copy_constant_to_function(return_value_id, self); self.dfg.set_value_from_id(*current_result_id, new_return_value_id); } true } -} -/// Copies a constant from one function to another. -/// Though it might seem we can just take a value out of `from_function` and call `make_value` on `to_function`, -/// if the constant is an array the values will still keep pointing to `from_function`. So, this function -/// recursively copies the array values too. -fn copy_constant_to_function( - from_function: &Function, - to_function: &mut Function, - constant_id: ValueId, -) -> ValueId { - if let Some((constant, typ)) = from_function.dfg.get_numeric_constant_with_type(constant_id) { - to_function.dfg.make_constant(constant, typ) - } else if let Some((constants, typ)) = from_function.dfg.get_array_constant(constant_id) { - let new_constants = constants - .iter() - .map(|constant_id| copy_constant_to_function(from_function, to_function, *constant_id)) - .collect(); - to_function.dfg.make_array(new_constants, typ) - } else { - unreachable!("A constant should be either a numeric constant or an array constant") + /// Copies a constant from this function to another one. + /// Though it might seem we can just take a value out of `self` and call `make_value` on `function`, + /// if the constant is an array the values will still keep pointing to `self`. So, this function + /// recursively copies the array values too. + fn copy_constant_to_function(&self, constant_id: ValueId, function: &mut Function) -> ValueId { + if let Some((constant, typ)) = self.dfg.get_numeric_constant_with_type(constant_id) { + function.dfg.make_constant(constant, typ) + } else if let Some((constants, typ)) = self.dfg.get_array_constant(constant_id) { + let new_constants = constants + .iter() + .map(|constant_id| self.copy_constant_to_function(*constant_id, function)) + .collect(); + function.dfg.make_array(new_constants, typ) + } else { + unreachable!("A constant should be either a numeric constant or an array constant") + } } } From 5298bf34c9b737f89b7ac8778014104b8ca40db6 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Mon, 11 Nov 2024 11:19:40 -0300 Subject: [PATCH 14/16] Introduce OptimizeResult --- .../src/ssa/opt/inline_const_brillig_calls.rs | 82 +++++++++++-------- 1 file changed, 50 insertions(+), 32 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index 24ebb5d4879..cf2ac25f1da 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -65,6 +65,18 @@ impl Ssa { } } +/// Result of trying to optimize an instruction (any instruction) in this pass. +enum OptimizeResult { + /// Nothing was done because the instruction wasn't a call to a brillig function, + /// or some arguments to it were not constants. + NotABrilligCall, + /// The instruction was a call to a brillig function, but we couldn't optimize it. + CannotOptimize(FunctionId), + /// The instruction was a call to a brillig function and we were able to optimize it, + /// returning the optimized function and the constant values it returned. + Optimized(Function, Vec), +} + impl Function { pub(crate) fn inline_const_brillig_calls( &mut self, @@ -75,23 +87,41 @@ impl Function { ) { for block_id in self.reachable_blocks() { for instruction_id in self.dfg[block_id].take_instructions() { - if !self.optimize_const_brillig_call( + let optimize_result = self.optimize_const_brillig_call( instruction_id, brillig_functions, brillig_functions_we_could_not_inline, inliner_aggressiveness, error_selector_to_type, - ) { - self.dfg[block_id].instructions_mut().push(instruction_id); + ); + match optimize_result { + OptimizeResult::NotABrilligCall => { + self.dfg[block_id].instructions_mut().push(instruction_id); + } + OptimizeResult::CannotOptimize(func_id) => { + self.dfg[block_id].instructions_mut().push(instruction_id); + brillig_functions_we_could_not_inline.insert(func_id); + } + OptimizeResult::Optimized(function, return_values) => { + // Replace the instruction results with the constant values we got + let current_results = self.dfg.instruction_results(instruction_id).to_vec(); + assert_eq!(return_values.len(), current_results.len()); + + for (current_result_id, return_value_id) in + current_results.iter().zip(return_values) + { + let new_return_value_id = + function.copy_constant_to_function(return_value_id, self); + self.dfg.set_value_from_id(*current_result_id, new_return_value_id); + } + } } } } } /// Tries to optimize an instruction if it's a call that points to a brillig function, - /// and all its arguments are constant. If the optimization is successful, the - /// values returned by the brillig call are replaced by the constant values that the - /// function returns, and this method returns `true`. + /// and all its arguments are constant. fn optimize_const_brillig_call( &mut self, instruction_id: InstructionId, @@ -99,29 +129,29 @@ impl Function { brillig_functions_we_could_not_inline: &mut HashSet, inliner_aggressiveness: i64, error_selector_to_type: &BTreeMap, - ) -> bool { + ) -> OptimizeResult { let instruction = &self.dfg[instruction_id]; let Instruction::Call { func: func_id, arguments } = instruction else { - return false; + return OptimizeResult::NotABrilligCall; }; let func_value = &self.dfg[*func_id]; let Value::Function(func_id) = func_value else { - return false; + return OptimizeResult::NotABrilligCall; }; + let func_id = *func_id; - let Some(function) = brillig_functions.get(func_id) else { - return false; + let Some(function) = brillig_functions.get(&func_id) else { + return OptimizeResult::NotABrilligCall; }; if !arguments.iter().all(|argument| self.dfg.is_constant(*argument)) { - brillig_functions_we_could_not_inline.insert(*func_id); - return false; + return OptimizeResult::CannotOptimize(func_id); } // The function we have is already a copy of the original function, but we need to clone // it again because there might be multiple calls to the same brillig function. - let mut function = Function::clone_with_id(*func_id, function); + let mut function = Function::clone_with_id(func_id, function); // Find the entry block and remove its parameters let entry_block_id = function.entry_block(); @@ -140,40 +170,28 @@ impl Function { // Try to fully optimize the function. If we can't, we can't inline it's constant value. let Ok(mut function) = optimize(function, inliner_aggressiveness, error_selector_to_type) else { - brillig_functions_we_could_not_inline.insert(*func_id); - return false; + return OptimizeResult::CannotOptimize(func_id); }; let entry_block = &mut function.dfg[entry_block_id]; // If the entry block has instructions, we can't inline it (we need a terminator) if !entry_block.instructions().is_empty() { - brillig_functions_we_could_not_inline.insert(*func_id); - return false; + brillig_functions_we_could_not_inline.insert(func_id); + return OptimizeResult::CannotOptimize(func_id); } let terminator = entry_block.take_terminator(); let TerminatorInstruction::Return { return_values, call_stack: _ } = terminator else { - brillig_functions_we_could_not_inline.insert(*func_id); - return false; + return OptimizeResult::CannotOptimize(func_id); }; // Sanity check: make sure all returned values are constant if !return_values.iter().all(|value_id| function.dfg.is_constant(*value_id)) { - brillig_functions_we_could_not_inline.insert(*func_id); - return false; - } - - // Replace the instruction results with the constant values we got - let current_results = self.dfg.instruction_results(instruction_id).to_vec(); - assert_eq!(return_values.len(), current_results.len()); - - for (current_result_id, return_value_id) in current_results.iter().zip(return_values) { - let new_return_value_id = function.copy_constant_to_function(return_value_id, self); - self.dfg.set_value_from_id(*current_result_id, new_return_value_id); + return OptimizeResult::CannotOptimize(func_id); } - true + OptimizeResult::Optimized(function, return_values) } /// Copies a constant from this function to another one. From e83fdbdf9c47141ad8930ac627610934e9392bbc Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Mon, 11 Nov 2024 11:20:47 -0300 Subject: [PATCH 15/16] More accurate comment regarding main brillig functions --- .../noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index cf2ac25f1da..e29842768ae 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -44,7 +44,8 @@ impl Ssa { // Remove the brillig functions that are no longer called for func_id in brillig_functions.keys() { - // We never want to remove the main function (it could be brillig if `--force-brillig` was given) + // We never want to remove the main function (it could be `unconstrained` or it + // could have been turned into brillig if `--force-brillig` was given) if self.main_id == *func_id { continue; } From 63aa63d847274def0bb9328383ec79f1a3698ee2 Mon Sep 17 00:00:00 2001 From: Ary Borenszweig Date: Mon, 11 Nov 2024 11:45:02 -0300 Subject: [PATCH 16/16] Introduce UnrollMode to preserve original checks in acir mode --- compiler/noirc_evaluator/src/ssa.rs | 14 ++-- .../src/ssa/opt/inline_const_brillig_calls.rs | 3 +- compiler/noirc_evaluator/src/ssa/opt/mod.rs | 2 +- .../noirc_evaluator/src/ssa/opt/unrolling.rs | 80 ++++++++++++++----- 4 files changed, 70 insertions(+), 29 deletions(-) diff --git a/compiler/noirc_evaluator/src/ssa.rs b/compiler/noirc_evaluator/src/ssa.rs index b6a18baf6ea..33cfb918ece 100644 --- a/compiler/noirc_evaluator/src/ssa.rs +++ b/compiler/noirc_evaluator/src/ssa.rs @@ -33,6 +33,7 @@ use noirc_frontend::{ hir_def::{function::FunctionSignature, types::Type as HirType}, monomorphization::ast::Program, }; +use opt::unrolling::UnrollMode; use tracing::{span, Level}; use self::{ @@ -130,6 +131,7 @@ fn optimize_ssa(builder: SsaBuilder, inliner_aggressiveness: i64) -> Result Result { let builder = builder // Run mem2reg with the CFG separated into blocks @@ -148,7 +151,7 @@ fn optimize_ssa_after_inline_const_brillig_calls( Ssa::evaluate_static_assert_and_assert_constant, "After `static_assert` and `assert_constant`:", )? - .try_run_pass(Ssa::unroll_loops_iteratively, "After Unrolling:")? + .try_run_pass(|ssa| Ssa::unroll_loops_iteratively(ssa, unroll_mode), "After Unrolling:")? .run_pass(Ssa::simplify_cfg, "After Simplifying (2nd):") .run_pass(Ssa::flatten_cfg, "After Flattening:") .run_pass(Ssa::remove_bit_shifts, "After Removing Bit Shifts:") @@ -457,11 +460,10 @@ impl SsaBuilder { } /// The same as `run_pass` but for passes that may fail - fn try_run_pass( - mut self, - pass: fn(Ssa) -> Result, - msg: &str, - ) -> Result { + fn try_run_pass(mut self, pass: F, msg: &str) -> Result + where + F: FnOnce(Ssa) -> Result, + { self.ssa = time(msg, self.print_codegen_timings, || pass(self.ssa))?; Ok(self.print(msg)) } diff --git a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs index e29842768ae..3b28c78e7fe 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/inline_const_brillig_calls.rs @@ -12,7 +12,7 @@ use crate::{ instruction::{Instruction, InstructionId, TerminatorInstruction}, value::{Value, ValueId}, }, - optimize_ssa_after_inline_const_brillig_calls, Ssa, SsaBuilder, + optimize_ssa_after_inline_const_brillig_calls, Ssa, SsaBuilder, UnrollMode, }, }; @@ -235,6 +235,7 @@ fn optimize( // a single function. For inlining to work we need to know all other functions that // exist (so we can inline them). Here we choose to skip this optimization for simplicity reasons. false, + UnrollMode::Brillig, )?; Ok(ssa.functions.pop_first().unwrap().1) } diff --git a/compiler/noirc_evaluator/src/ssa/opt/mod.rs b/compiler/noirc_evaluator/src/ssa/opt/mod.rs index c9aa5e17efe..9cd71e08973 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/mod.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/mod.rs @@ -21,4 +21,4 @@ mod remove_if_else; mod resolve_is_unconstrained; mod runtime_separation; mod simplify_cfg; -mod unrolling; +pub(crate) mod unrolling; diff --git a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs index 003877e85f7..55681618853 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs @@ -37,13 +37,29 @@ use crate::{ }; use fxhash::FxHashMap as HashMap; +/// In this mode we are unrolling. +#[derive(Debug, Clone, Copy)] +pub(crate) enum UnrollMode { + /// This is the normal unroll mode, where we are unrolling a brillig function. + Acir, + /// There's one optimization, `inline_const_brillig_calls`, where we try to optimize + /// brillig functions with all constant arguments. For that we turn the brillig function + /// into an acir one and try to optimize it. This brillig function could then have + /// `break` statements that can't be possible in acir, so in this mode we don't panic + /// if we end up in unexpected situations that might have been produced by `break`. + Brillig, +} + 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. - pub(crate) fn unroll_loops_iteratively(mut ssa: Ssa) -> Result { + pub(crate) fn unroll_loops_iteratively( + mut ssa: Ssa, + mode: UnrollMode, + ) -> Result { // Try to unroll loops first: let mut unroll_errors; - (ssa, unroll_errors) = ssa.try_to_unroll_loops(); + (ssa, unroll_errors) = ssa.try_to_unroll_loops(mode); // Keep unrolling until no more errors are found while !unroll_errors.is_empty() { @@ -58,7 +74,7 @@ impl Ssa { ssa = ssa.mem2reg(); // Unroll again - (ssa, unroll_errors) = ssa.try_to_unroll_loops(); + (ssa, unroll_errors) = ssa.try_to_unroll_loops(mode); // 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)); @@ -71,22 +87,22 @@ impl Ssa { /// 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) { + pub(crate) fn try_to_unroll_loops(mut self, mode: UnrollMode) -> (Ssa, Vec) { let mut errors = vec![]; for function in self.functions.values_mut() { - function.try_to_unroll_loops(&mut errors); + function.try_to_unroll_loops(mode, &mut errors); } (self, errors) } } impl Function { - pub(crate) fn try_to_unroll_loops(&mut self, errors: &mut Vec) { + pub(crate) fn try_to_unroll_loops(&mut self, mode: UnrollMode, 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)); + errors.extend(find_all_loops(self).unroll_each_loop(self, mode)); } } } @@ -151,7 +167,7 @@ fn find_all_loops(function: &Function) -> Loops { impl Loops { /// Unroll all loops within a given function. /// Any loops which fail to be unrolled (due to using non-constant indices) will be unmodified. - fn unroll_each_loop(mut self, function: &mut Function) -> Vec { + fn unroll_each_loop(mut self, function: &mut Function, mode: UnrollMode) -> Vec { let mut unroll_errors = vec![]; while let Some(next_loop) = self.yet_to_unroll.pop() { // If we've previously modified a block in this loop we need to refresh the context. @@ -161,13 +177,13 @@ impl Loops { new_context.failed_to_unroll = self.failed_to_unroll; return unroll_errors .into_iter() - .chain(new_context.unroll_each_loop(function)) + .chain(new_context.unroll_each_loop(function, mode)) .collect(); } // Don't try to unroll the loop again if it is known to fail if !self.failed_to_unroll.contains(&next_loop.header) { - match unroll_loop(function, &self.cfg, &next_loop) { + match unroll_loop(function, &self.cfg, &next_loop, mode) { Ok(_) => self.modified_blocks.extend(next_loop.blocks), Err(call_stack) => { self.failed_to_unroll.insert(next_loop.header); @@ -217,12 +233,13 @@ fn unroll_loop( function: &mut Function, cfg: &ControlFlowGraph, loop_: &Loop, + mode: UnrollMode, ) -> Result<(), CallStack> { - let mut unroll_into = get_pre_header(cfg, loop_)?; + let mut unroll_into = get_pre_header(cfg, loop_, mode)?; let mut jump_value = get_induction_variable(function, unroll_into)?; while let Some(context) = unroll_loop_header(function, loop_, unroll_into, jump_value)? { - let (last_block, last_value) = context.unroll_loop_iteration()?; + let (last_block, last_value) = context.unroll_loop_iteration(mode)?; unroll_into = last_block; jump_value = last_value; } @@ -233,16 +250,28 @@ fn unroll_loop( /// The loop pre-header is the block that comes before the loop begins. Generally a header block /// is expected to have 2 predecessors: the pre-header and the final block of the loop which jumps /// back to the beginning. -fn get_pre_header(cfg: &ControlFlowGraph, loop_: &Loop) -> Result { +fn get_pre_header( + cfg: &ControlFlowGraph, + loop_: &Loop, + mode: UnrollMode, +) -> Result { let mut pre_header = cfg .predecessors(loop_.header) .filter(|predecessor| *predecessor != loop_.back_edge_start) .collect::>(); - if pre_header.len() == 1 { - Ok(pre_header.remove(0)) - } else { - Err(CallStack::new()) + match mode { + UnrollMode::Acir => { + assert_eq!(pre_header.len(), 1); + Ok(pre_header.remove(0)) + } + UnrollMode::Brillig => { + if pre_header.len() == 1 { + Ok(pre_header.remove(0)) + } else { + Err(CallStack::new()) + } + } } } @@ -367,7 +396,10 @@ 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) -> Result<(BasicBlockId, ValueId), CallStack> { + fn unroll_loop_iteration( + mut self, + mode: UnrollMode, + ) -> Result<(BasicBlockId, ValueId), CallStack> { let mut next_blocks = self.unroll_loop_block(); while let Some(block) = next_blocks.pop() { @@ -380,7 +412,12 @@ impl<'f> LoopIteration<'f> { } } - self.induction_value.ok_or_else(CallStack::new) + match mode { + UnrollMode::Acir => Ok(self + .induction_value + .expect("Expected to find the induction variable by end of loop iteration")), + UnrollMode::Brillig => self.induction_value.ok_or_else(CallStack::new), + } } /// Unroll a single block in the current iteration of the loop @@ -512,6 +549,7 @@ mod tests { use crate::ssa::{ function_builder::FunctionBuilder, ir::{instruction::BinaryOp, map::Id, types::Type}, + opt::unrolling::UnrollMode, }; #[test] @@ -632,7 +670,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) = ssa.try_to_unroll_loops(UnrollMode::Acir); assert_eq!(errors.len(), 0, "All loops should be unrolled"); assert_eq!(ssa.main().reachable_blocks().len(), 5); } @@ -682,7 +720,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) = ssa.try_to_unroll_loops(UnrollMode::Acir); assert_eq!(errors.len(), 1, "Expected to fail to unroll loop"); } }