diff --git a/compiler/noirc_errors/src/call_stack.rs b/compiler/noirc_errors/src/call_stack.rs index c982cae6686..91fb655e570 100644 --- a/compiler/noirc_errors/src/call_stack.rs +++ b/compiler/noirc_errors/src/call_stack.rs @@ -90,7 +90,7 @@ impl CallStackHelper { result } - /// Returns a new CallStackId which extends the call_stack with the provided call_stack. + /// Returns a new [CallStackId] which extends the `call_stack` with the provided `locations`. pub fn extend_call_stack( &mut self, mut call_stack: CallStackId, @@ -102,7 +102,7 @@ impl CallStackHelper { call_stack } - /// Adds a location to the call stack + /// Adds a location to the call stack, maintaining the location cache along the way. pub fn add_child(&mut self, call_stack: CallStackId, location: Location) -> CallStackId { let key = rustc_hash::FxBuildHasher.hash_one(location); if let Some(result) = self.locations[call_stack.index()].children_hash.get(&key) { @@ -142,7 +142,7 @@ impl CallStackHelper { } } - /// Get (or create) a CallStackId corresponding to the given locations + /// Get (or create) a CallStackId corresponding to the given locations. pub fn get_or_insert_locations(&mut self, locations: &CallStack) -> CallStackId { self.extend_call_stack(CallStackId::root(), locations) } diff --git a/compiler/noirc_evaluator/src/acir/mod.rs b/compiler/noirc_evaluator/src/acir/mod.rs index 0ec1317b264..01c387f5681 100644 --- a/compiler/noirc_evaluator/src/acir/mod.rs +++ b/compiler/noirc_evaluator/src/acir/mod.rs @@ -138,7 +138,7 @@ impl<'a> Context<'a> { ssa: &Ssa, function: &Function, ) -> Result>, RuntimeError> { - self.acir_context.set_call_stack_helper(self.brillig.call_stacks.to_owned()); + self.acir_context.set_call_stack_helper(self.brillig.call_stacks().clone()); match function.runtime() { RuntimeType::Acir(inline_type) => { match inline_type { diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs index 16f31de894e..a33e6406bc4 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs @@ -164,6 +164,9 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> { } } + // Allocate and initialize hoisted constants. These don't have a variable ID associated with them, + // so we return them explicitly, while the allocated global variables are in the `ssa_value_allocations` + // field of the `FunctionContext`. let mut new_hoisted_constants = HashMap::default(); for (constant, typ) in hoisted_global_constants.iter().copied() { let new_variable = allocate_value_with_type(self.brillig_context, Type::Numeric(typ)); @@ -315,6 +318,8 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> { } /// Converts an SSA instruction into a sequence of Brillig opcodes. + /// + /// If this is the last time a variable is used in this block, its memory slot gets deallocated, unless it's a global. fn convert_ssa_instruction( &mut self, instruction_id: InstructionId, @@ -510,8 +515,9 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> { for dead_variable in dead_variables { // Globals are reserved throughout the entirety of the program - let not_hoisted_global = self.get_hoisted_global(dfg, *dead_variable).is_none(); - if !dfg.is_global(*dead_variable) && not_hoisted_global { + let is_global = dfg.is_global(*dead_variable); + let is_hoisted_global = self.get_hoisted_global(dfg, *dead_variable).is_some(); + if !is_global && !is_hoisted_global { self.variables.remove_variable( dead_variable, self.function_context, @@ -520,6 +526,8 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> { } } } + + // Clear the call stack; it only applied to this instruction. self.brillig_context.set_call_stack(CallStackId::root()); } diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_globals.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_globals.rs index 67aeea12129..3bc9827ea3e 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_globals.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_globals.rs @@ -19,27 +19,39 @@ use crate::{ }, }; -/// Context structure for generating Brillig globals -/// it stores globals related data required for code generation of regular Brillig functions. -#[derive(Default)] -pub(crate) struct BrilligGlobals { +/// Map entry points to functions reachable from them and vice versa. +struct CallMap { + /// Maps a Brillig entry point to all functions called in that entry point. + /// This includes any nested calls as well, as we want to be able to associate + /// any Brillig function with the appropriate global allocations. + /// + /// This is the reverse of `inner_call_to_entry_point`. + entry_point_to_inner_calls: BTreeMap>, + /// Maps an inner call to its Brillig entry point, of which there should be only 1. + /// This is used to simplify fetching global allocations when compiling + /// individual Brillig functions. + /// + /// This is the reverse of `entry_point_to_inner_calls`. + inner_call_to_entry_point: HashMap>, +} + +/// Initial context structure for generating Brillig globals. +pub(crate) struct BrilligGlobalsInit { + call_map: CallMap, /// Both `used_globals` and `brillig_entry_points` need to be built /// from a function call graph. /// /// Maps a Brillig function to the globals used in that function. /// This includes all globals used in functions called internally. used_globals: HashMap>, - /// Maps a Brillig entry point to all functions called in that entry point. - /// This includes any nested calls as well, as we want to be able to associate - /// any Brillig function with the appropriate global allocations. - brillig_entry_points: BTreeMap>, /// Maps a Brillig entry point to constants shared across the entry point and its nested calls. - hoisted_global_constants: HashMap, + constant_usage: HashMap, +} - /// Maps an inner call to its Brillig entry point - /// This is simply used to simplify fetching global allocations when compiling - /// individual Brillig functions. - inner_call_to_entry_point: HashMap>, +/// Final context structure for generating Brillig globals. +/// It stores globals related data required for the code generation of regular Brillig functions. +pub(crate) struct BrilligGlobals { + call_map: CallMap, /// Final map that associated an entry point with its Brillig global allocations entry_point_globals_map: HashMap, /// Final map that associates an entry point with any local function constants @@ -49,50 +61,72 @@ pub(crate) struct BrilligGlobals { entry_point_hoisted_globals_map: HashMap, } -/// Mapping of SSA value ids to their Brillig allocations +/// Mapping of SSA value ids to their Brillig allocations. pub(crate) type SsaToBrilligGlobals = HashMap; -/// Mapping of constant values shared across functions hoisted to the global memory space + +/// Mapping of constant values shared across functions hoisted to the global memory space. pub(crate) type HoistedConstantsToBrilligGlobals = HashMap<(FieldElement, NumericType), BrilligVariable>; -/// Mapping of a constant value and the number of functions in which it occurs + +/// Mapping of a constant value to the number of functions in which it occurs. pub(crate) type ConstantCounterMap = HashMap<(FieldElement, NumericType), usize>; -impl BrilligGlobals { - pub(crate) fn new( - ssa: &Ssa, - mut used_globals: HashMap>, - main_id: FunctionId, - ) -> Self { +#[derive(Default)] +struct ConstantAllocationCache(HashMap); + +impl ConstantAllocationCache { + fn get_constants(&mut self, func: &Function) -> &ConstantAllocation { + self.0.entry(func.id()).or_insert_with(|| ConstantAllocation::from_function(func)) + } +} + +impl BrilligGlobalsInit { + /// Collect information about global usage for each Brillig entry point: + /// * which globals are used by an entry point and its callees, + /// * how many times each constant is used by an entry point and its callees. + /// + /// The population of allocation related information is deferred to [Self::declare_globals]. + pub(crate) fn new(ssa: &Ssa, main_id: FunctionId) -> Self { + let mut used_globals = ssa.used_globals_in_functions(); let call_graph = CallGraph::from_ssa(ssa); let brillig_entry_points = get_brillig_entry_points_with_reachability(&ssa.functions, main_id, &call_graph); - let mut hoisted_global_constants: HashMap = - HashMap::default(); + let mut constant_usage: HashMap = HashMap::default(); + + let mut constant_allocations = ConstantAllocationCache::default(); + // Mark any globals used in a Brillig entry point. - // Using the information collected we can determine which globals - // an entry point must initialize. + // Using the information collected we can determine which globals an entry point must initialize. for (entry_point, entry_point_inner_calls) in brillig_entry_points.iter() { + // Increment the use-count of local constants in this entry point. + let entry_func = &ssa.functions[entry_point]; Self::mark_globals_for_hoisting( - &mut hoisted_global_constants, + &mut constant_usage, *entry_point, - &ssa.functions[entry_point], + entry_func, + constant_allocations.get_constants(entry_func), ); + // Increment the use-count of local constants in all functions called by the entry point. for inner_call in entry_point_inner_calls.iter() { + let inner_func = &ssa.functions[inner_call]; Self::mark_globals_for_hoisting( - &mut hoisted_global_constants, + &mut constant_usage, *entry_point, - &ssa.functions[inner_call], + inner_func, + constant_allocations.get_constants(inner_func), ); + // Consider each global used by the inner function to be also used by the entry point. let inner_globals = used_globals .get(inner_call) - .expect("Should have a slot for each function") + .expect("ICE: inner function should be in used globals") .clone(); + used_globals .get_mut(entry_point) - .expect("ICE: should have func") + .expect("ICE: entry point should be in used globals") .extend(inner_globals); } } @@ -100,78 +134,88 @@ impl BrilligGlobals { let inner_call_to_entry_point = build_inner_call_to_entry_points(&brillig_entry_points); Self { + call_map: CallMap { + entry_point_to_inner_calls: brillig_entry_points, + inner_call_to_entry_point, + }, used_globals, - brillig_entry_points, - inner_call_to_entry_point, - hoisted_global_constants, - ..Default::default() + constant_usage, } } - pub(crate) fn entry_points(&self) -> &BTreeMap> { - &self.brillig_entry_points - } - /// Helper for marking that a constant was instantiated in a given function. /// For a given entry point, we want to determine which constants are shared across multiple functions. + /// + /// Increments the used-in-functions counter for each non-global constant, + /// which then can be considered for hoisting. fn mark_globals_for_hoisting( hoisted_global_constants: &mut HashMap, entry_point: FunctionId, function: &Function, + constants: &ConstantAllocation, ) { - // We can potentially have multiple local constants with the same value and type - let constants = ConstantAllocation::from_function(function); + let entry_const_usage = hoisted_global_constants.entry(entry_point).or_default(); + + // We can potentially have multiple local constants with the same value and type, + // in which case even one function will count as multiple occurrences. for constant in constants.get_constants() { - let value = function.dfg.get_numeric_constant(constant); - let value = value.unwrap(); - let typ = function.dfg.type_of_value(constant); + let value = function.dfg.get_numeric_constant_with_type(constant); + let (value, typ) = value.expect("it was found by constant allocation"); + // If the value is an actual global then there is nothing to hoist; + // otherwise increment the number of functions it is used in. if !function.dfg.is_global(constant) { - hoisted_global_constants - .entry(entry_point) - .or_default() - .entry((value, typ.unwrap_numeric())) + entry_const_usage + .entry((value, typ)) .and_modify(|counter| *counter += 1) .or_insert(1); } } } + /// A one-time initialization of the entry-points-to-globals maps. + /// + /// Selects the constants to be hoisted based on their usage count, then declares Brillig variables in the global memory space. + /// + /// The bytecode for each global is inserted into the [Brillig] structure. pub(crate) fn declare_globals( - &mut self, + mut self, globals_dfg: &DataFlowGraph, brillig: &mut Brillig, options: &BrilligOptions, - ) { + ) -> BrilligGlobals { let mut entry_point_globals_map = HashMap::default(); let mut entry_point_hoisted_globals_map = HashMap::default(); // We only need to generate globals for entry points - for (entry_point, _) in self.brillig_entry_points.iter() { - let entry_point = *entry_point; + for entry_point in self.call_map.entry_point_to_inner_calls.keys().copied() { + let used_globals = self + .used_globals + .remove(&entry_point) + .expect("entry point should be in used globals"); - let used_globals = self.used_globals.remove(&entry_point).unwrap_or_default(); - // Select set of constants which can be hoisted from function's to the global memory space - // for a given entry point. + // Select the set of constants which can be hoisted from the function to the global memory space + // for a given entry point: hoist anything that has more than 1 use. let hoisted_global_constants = self - .hoisted_global_constants + .constant_usage .remove(&entry_point) - .unwrap_or_default() + .expect("entry point should have hoisted constants") .iter() - .filter_map( - |(&value, &num_occurrences)| { - if num_occurrences > 1 { Some(value) } else { None } - }, - ) + .filter_map(|(&value, &num_occurrences)| (num_occurrences > 1).then_some(value)) .collect(); - let (artifact, brillig_globals, globals_size, hoisted_global_constants) = brillig - .convert_ssa_globals( - options, - globals_dfg, - &used_globals, - &hoisted_global_constants, - entry_point, - ); + // Compile a separate global bytecode and allocation for each entry point. + let BrilligGlobalsArtifact { + artifact, + brillig_globals, + globals_size, + hoisted_global_constants, + } = brillig.convert_ssa_globals( + options, + globals_dfg, + &used_globals, + &hoisted_global_constants, + entry_point, + ); entry_point_globals_map.insert(entry_point, brillig_globals); entry_point_hoisted_globals_map.insert(entry_point, hoisted_global_constants); @@ -180,29 +224,45 @@ impl BrilligGlobals { brillig.globals_memory_size.insert(entry_point, globals_size); } - self.entry_point_globals_map = entry_point_globals_map; - self.entry_point_hoisted_globals_map = entry_point_hoisted_globals_map; + BrilligGlobals { + call_map: self.call_map, + entry_point_globals_map, + entry_point_hoisted_globals_map, + } + } +} + +impl BrilligGlobals { + pub(crate) fn init(ssa: &Ssa, main_id: FunctionId) -> BrilligGlobalsInit { + BrilligGlobalsInit::new(ssa, main_id) } + /// Check whether a function is a Brillig entry point. + pub(crate) fn is_entry_point(&self, func_id: &FunctionId) -> bool { + self.call_map.entry_point_to_inner_calls.contains_key(func_id) + } /// Fetch the global allocations that can possibly be accessed /// by any given Brillig function (non-entry point or entry point). + /// /// The allocations available to a function are determined by its entry point. - /// For a given function id input, this function will search for that function's + /// For a given function ID input, this function will search for that function's /// entry point and fetch the global allocations associated with that entry point. + /// /// These allocations can then be used when compiling the Brillig function /// and resolving global variables. + /// + /// Panics if the function hasn't been prepared during initialization. pub(crate) fn get_brillig_globals( &self, brillig_function_id: FunctionId, - ) -> Option<(&SsaToBrilligGlobals, &HoistedConstantsToBrilligGlobals)> { + ) -> (&SsaToBrilligGlobals, &HoistedConstantsToBrilligGlobals) { // Check whether `brillig_function_id` is itself an entry point. // If so, return the global allocations directly. - let entry_point_globals = self.get_entry_point_globals(&brillig_function_id); - if entry_point_globals.is_some() { + if let Some(entry_point_globals) = self.get_entry_point_globals(&brillig_function_id) { return entry_point_globals; } - let entry_points = self.inner_call_to_entry_point.get(&brillig_function_id); + let entry_points = self.call_map.inner_call_to_entry_point.get(&brillig_function_id); let Some(entry_points) = entry_points else { unreachable!( "ICE: Expected global allocation to be set for function {brillig_function_id}" @@ -210,45 +270,46 @@ impl BrilligGlobals { }; // Sanity check: We should have guaranteed earlier that an inner call has only a single entry point - assert_eq!(entry_points.len(), 1, "{brillig_function_id} has multiple entry points"); + assert_eq!(entry_points.len(), 1, "ICE: {brillig_function_id} has multiple entry points"); let entry_point = entry_points.first().expect("ICE: Inner call should have an entry point"); - self.get_entry_point_globals(entry_point) + self.get_entry_point_globals(entry_point).expect("ICE: Entry point should have globals") } - /// Fetch the global allocations for a given entry point. + /// Try to fetch the global allocations for a given entry point. + /// /// This contains both the user specified globals, as well as any constants shared /// across functions that have been hoisted into the global space. + /// + /// Returns `None` if the function is not an entry point. fn get_entry_point_globals( &self, - entry_point: &FunctionId, + brillig_function_id: &FunctionId, ) -> Option<(&SsaToBrilligGlobals, &HoistedConstantsToBrilligGlobals)> { - if let (Some(globals), Some(hoisted_constants)) = ( - self.entry_point_globals_map.get(entry_point), - self.entry_point_hoisted_globals_map.get(entry_point), - ) { - Some((globals, hoisted_constants)) - } else { - None - } + let globals = self.entry_point_globals_map.get(brillig_function_id)?; + let hoisted_constants = self.entry_point_hoisted_globals_map.get(brillig_function_id)?; + Some((globals, hoisted_constants)) } } /// A globals artifact containing all information necessary for utilizing /// globals from SSA during Brillig code generation. -pub(crate) type BrilligGlobalsArtifact = ( - // The actual bytecode declaring globals and any metadata needing for linking - BrilligArtifact, - // The SSA value -> Brillig global allocations - // This will be used for fetching global values when compiling functions to Brillig. - HashMap, - // The size of the global memory - usize, - // Duplicate SSA constants local to a function -> Brillig global allocations - HashMap<(FieldElement, NumericType), BrilligVariable>, -); +pub(crate) struct BrilligGlobalsArtifact { + /// The actual bytecode declaring globals and any metadata needing for linking. + pub artifact: BrilligArtifact, + /// The SSA value -> Brillig global allocations. + /// This will be used for fetching global values when compiling functions to Brillig. + pub brillig_globals: HashMap, + /// The size of the global memory. + pub globals_size: usize, + /// Duplicate SSA constants local to a function -> Brillig global allocations. + pub hoisted_global_constants: HashMap<(FieldElement, NumericType), BrilligVariable>, +} impl Brillig { + /// Compile global opcodes and return the bytecode along with the memory allocations + /// of global variables and hoisted constants, so that we can use them as pre-allocated + /// registers when we compile functions. pub(crate) fn convert_ssa_globals( &mut self, options: &BrilligOptions, @@ -258,7 +319,7 @@ impl Brillig { entry_point: FunctionId, ) -> BrilligGlobalsArtifact { let mut brillig_context = BrilligContext::new_for_global_init(options, entry_point); - // The global space does not have globals itself + // The global space does not have globals itself. let empty_globals = HashMap::default(); // We can use any ID here as this context is only going to be used for globals which does not differentiate // by functions and blocks. The only Label that should be used in the globals context is `Label::globals_init()` @@ -288,17 +349,22 @@ impl Brillig { brillig_context.return_instruction(); + // The artifact contains the Brillig bytecode to initialize the global variables and constants. let artifact = brillig_context.artifact(); - // The global registers going to become pre-allocated registers when we compile other functions, + // The global registers are going to become pre-allocated registers when we compile functions, // so we can stop tracking their allocation life cycles and keep just the memory addresses. - let hoisted_global_constants = hoisted_global_constants .into_iter() .map(|(field, variable)| (field, variable.detach())) .collect(); - (artifact, function_context.ssa_value_allocations, globals_size, hoisted_global_constants) + BrilligGlobalsArtifact { + artifact, + brillig_globals: function_context.ssa_value_allocations, + globals_size, + hoisted_global_constants, + } } } @@ -578,10 +644,11 @@ mod tests { assert_eq!(brillig.globals.len(), 1, "Should have a single entry point"); for (func_id, artifact) in brillig.globals { assert_eq!(func_id.to_u32(), 1); + // We expect constants 0 and 1 to be hoisted. Not 20 because it only appears in one function. assert_eq!( artifact.byte_code.len(), 3, - "Expected enough opcodes to initialize the hoisted constants" + "Expected enough opcodes to initialize the hoisted constants:\n{artifact}" ); let Opcode::Const { destination, bit_size, value } = &artifact.byte_code[0] else { panic!("First opcode is expected to be `Const`"); @@ -594,7 +661,7 @@ mod tests { assert_eq!(*value, FieldElement::from(0u128)); let Opcode::Const { destination, bit_size, value } = &artifact.byte_code[1] else { - panic!("First opcode is expected to be `Const`"); + panic!("Second opcode is expected to be `Const`"); }; assert_eq!( destination.unwrap_direct(), diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/constant_allocation.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/constant_allocation.rs index 707781019da..43bc6c45838 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/constant_allocation.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/constant_allocation.rs @@ -6,41 +6,51 @@ use std::collections::{BTreeMap, BTreeSet}; use crate::ssa::ir::{ basic_block::BasicBlockId, - cfg::ControlFlowGraph, dfg::DataFlowGraph, dom::DominatorTree, function::Function, instruction::InstructionId, - post_order::PostOrder, value::{Value, ValueId}, }; -use super::variable_liveness::{collect_variables_of_value, variables_used_in_instruction}; +use super::variable_liveness::{is_variable, variables_used_in_instruction}; +use crate::ssa::opt::Loops; +/// Indicate where a variable was used in a block. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] pub(crate) enum InstructionLocation { Instruction(InstructionId), Terminator, } +/// Decisions about which block a constant should be allocated in. +/// +/// The allocation points can further be hoisted into the global +/// space if they are shared across functions. That algorithm is +/// based on these per-function results. #[derive(Default)] pub(crate) struct ConstantAllocation { + /// Map each constant to the blocks and instructions in which it is used. constant_usage: BTreeMap>>, + /// Map each block and instruction to the list of constants that should be allocated at that point. allocation_points: BTreeMap>>, + /// The dominator tree is used to find the common dominator of all the blocks that share a constant. dominator_tree: DominatorTree, + /// Further to finding the common dominator, we try to find a dominator that isn't part of a loop. blocks_within_loops: BTreeSet, } impl ConstantAllocation { + /// Run the constant allocation algorithm for a [Function] and return the decisions. pub(crate) fn from_function(func: &Function) -> Self { - let cfg = ControlFlowGraph::with_function(func); - let post_order = PostOrder::with_function(func); - let mut dominator_tree = DominatorTree::with_cfg_and_post_order(&cfg, &post_order); - let blocks_within_loops = find_all_blocks_within_loops(func, &cfg, &mut dominator_tree); + let loops = Loops::find_all(func); + let blocks_within_loops = + loops.yet_to_unroll.into_iter().flat_map(|_loop| _loop.blocks).collect(); + let mut instance = ConstantAllocation { constant_usage: BTreeMap::default(), allocation_points: BTreeMap::default(), - dominator_tree, + dominator_tree: loops.dom, blocks_within_loops, }; instance.collect_constant_usage(func); @@ -49,12 +59,14 @@ impl ConstantAllocation { instance } + /// Collect all constants allocated in a given block. pub(crate) fn allocated_in_block(&self, block_id: BasicBlockId) -> Vec { self.allocation_points.get(&block_id).map_or(Vec::default(), |allocations| { - allocations.iter().flat_map(|(_, constants)| constants.iter()).copied().collect() + allocations.iter().flat_map(|(_, constants)| constants).copied().collect() }) } + /// Collect all constants allocated in a given block at a specific location. pub(crate) fn allocated_at_location( &self, block_id: BasicBlockId, @@ -65,10 +77,11 @@ impl ConstantAllocation { }) } + /// Visit all constant variables in the function and record their locations. fn collect_constant_usage(&mut self, func: &Function) { let mut record_if_constant = |block_id: BasicBlockId, value_id: ValueId, location: InstructionLocation| { - if is_constant_value(value_id, &func.dfg) { + if is_numeric_constant(value_id, &func.dfg) { self.constant_usage .entry(value_id) .or_default() @@ -91,14 +104,16 @@ impl ConstantAllocation { } if let Some(terminator_instruction) = block.terminator() { terminator_instruction.for_each_value(|value_id| { - if let Some(variable) = collect_variables_of_value(value_id, &func.dfg) { - record_if_constant(block_id, variable, InstructionLocation::Terminator); + if is_variable(value_id, &func.dfg) { + record_if_constant(block_id, value_id, InstructionLocation::Terminator); } }); } } } + /// Based on the [Self::constant_usage] collected, find the common dominator of all the block where a constant is used + /// and mark it as the allocation point for the constant. fn decide_allocation_points(&mut self, func: &Function) { for (constant_id, usage_in_blocks) in self.constant_usage.iter() { let block_ids: Vec<_> = usage_in_blocks.keys().copied().collect(); @@ -126,26 +141,22 @@ impl ConstantAllocation { } } + /// Decide where to allocate a constant, based on the common dominator of the provided block list. fn decide_allocation_point( &self, constant_id: ValueId, - blocks_where_is_used: &[BasicBlockId], + used_in_blocks: &[BasicBlockId], func: &Function, ) -> BasicBlockId { // Find the common dominator of all the blocks where the constant is used. - let common_dominator = if blocks_where_is_used.len() == 1 { - blocks_where_is_used[0] - } else { - let mut common_dominator = blocks_where_is_used[0]; - - for block_id in blocks_where_is_used.iter().skip(1) { - common_dominator = - self.dominator_tree.common_dominator(common_dominator, *block_id); - } - - common_dominator - }; - // If the value only contains constants, it's safe to hoist outside of any loop + let common_dominator = used_in_blocks + .iter() + .copied() + .reduce(|a, b| self.dominator_tree.common_dominator(a, b)) + .unwrap_or(used_in_blocks[0]); + + // If the value only contains constants, it's safe to hoist outside of any loop. + // Technically we know this is going to be true, because we only collected values which are `Value::NumericConstant`. if func.dfg.is_constant(constant_id) { self.exit_loops(common_dominator) } else { @@ -165,62 +176,12 @@ impl ConstantAllocation { current_block } + /// Return the SSA [ValueId] of all constants (the same numeric constant might appear with multiple IDs). pub(crate) fn get_constants(&self) -> BTreeSet { self.constant_usage.keys().copied().collect() } } -pub(crate) fn is_constant_value(id: ValueId, dfg: &DataFlowGraph) -> bool { +fn is_numeric_constant(id: ValueId, dfg: &DataFlowGraph) -> bool { matches!(&dfg[id], Value::NumericConstant { .. }) } - -/// For a given function, finds all the blocks that are within loops -fn find_all_blocks_within_loops( - func: &Function, - cfg: &ControlFlowGraph, - dominator_tree: &mut DominatorTree, -) -> BTreeSet { - let mut blocks_in_loops = BTreeSet::default(); - for block_id in func.reachable_blocks() { - let block = &func.dfg[block_id]; - let successors = block.successors(); - for successor_id in successors { - if dominator_tree.dominates(successor_id, block_id) { - blocks_in_loops.extend(find_blocks_in_loop(successor_id, block_id, cfg)); - } - } - } - - blocks_in_loops -} - -/// Return each block that is in a loop starting in the given header block. -/// Expects back_edge_start -> header to be the back edge of the loop. -fn find_blocks_in_loop( - header: BasicBlockId, - back_edge_start: BasicBlockId, - cfg: &ControlFlowGraph, -) -> BTreeSet { - let mut blocks = BTreeSet::default(); - blocks.insert(header); - - let mut insert = |block, stack: &mut Vec| { - if !blocks.contains(&block) { - blocks.insert(block); - stack.push(block); - } - }; - - // Starting from the back edge of the loop, each predecessor of this block until - // the header is within the loop. - let mut stack = vec![]; - insert(back_edge_start, &mut stack); - - while let Some(block) = stack.pop() { - for predecessor in cfg.predecessors(block) { - insert(predecessor, &mut stack); - } - } - - blocks -} diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/variable_liveness.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/variable_liveness.rs index aadf294271e..b5b6578e2fb 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/variable_liveness.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/variable_liveness.rs @@ -1,80 +1,76 @@ //! This module analyzes the liveness of variables (non-constant values) throughout a function. //! It uses the approach detailed in the section 4.2 of this paper -use crate::ssa::ir::{ - basic_block::{BasicBlock, BasicBlockId}, - cfg::ControlFlowGraph, - dfg::DataFlowGraph, - dom::DominatorTree, - function::Function, - instruction::{Instruction, InstructionId}, - post_order::PostOrder, - value::{Value, ValueId}, +use std::collections::BTreeSet; + +use crate::ssa::{ + ir::{ + basic_block::{BasicBlock, BasicBlockId}, + cfg::ControlFlowGraph, + dfg::DataFlowGraph, + dom::DominatorTree, + function::Function, + instruction::{Instruction, InstructionId}, + post_order::PostOrder, + value::{Value, ValueId}, + }, + opt::Loops, }; use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet}; use super::constant_allocation::ConstantAllocation; +/// A set of [ValueId]s referring to SSA variables (not functions). +type Variables = HashSet; +/// The set variables which are dead after a given instruction (in a given block). +type LastUses = HashMap; +/// Maps a loop (identified by its header and back-edge) to the blocks in the loop body. +type LoopMap = HashMap>; + /// A back edge is an edge from a node to one of its ancestors. It denotes a loop in the CFG. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] struct BackEdge { + /// The header of the loop. header: BasicBlockId, + /// The back-edge of the loop. start: BasicBlockId, } -fn find_back_edges( - func: &Function, - cfg: &ControlFlowGraph, - post_order: &PostOrder, -) -> HashSet { - let mut tree = DominatorTree::with_cfg_and_post_order(cfg, post_order); - let mut back_edges = HashSet::default(); - - for block_id in func.reachable_blocks() { - let block = &func.dfg[block_id]; - let successors = block.successors(); - for successor_id in successors { - if tree.dominates(successor_id, block_id) { - back_edges.insert(BackEdge { start: block_id, header: successor_id }); - } - } - } - - back_edges -} - -/// Collects the underlying variables inside a value id. It might be more than one, for example in constant arrays that are constructed with multiple vars. -pub(crate) fn collect_variables_of_value( - value_id: ValueId, - dfg: &DataFlowGraph, -) -> Option { +/// Check if the [Value] behind the [ValueId] requires register allocation (like a function parameter), +/// rather than a global value like a user-defined function, intrinsic, or foreign function. +pub(super) fn is_variable(value_id: ValueId, dfg: &DataFlowGraph) -> bool { let value = &dfg[value_id]; match value { Value::Instruction { .. } | Value::Param { .. } | Value::NumericConstant { .. } - | Value::Global(_) => Some(value_id), + | Value::Global(_) => true, // Functions are not variables in a defunctionalized SSA. Only constant function values should appear. - Value::ForeignFunction(_) | Value::Function(_) | Value::Intrinsic(..) => None, + Value::ForeignFunction(_) | Value::Function(_) | Value::Intrinsic(..) => false, } } -pub(crate) fn variables_used_in_instruction( +/// Collect all [ValueId]s used in an [Instruction] which refer to variables (not functions). +pub(super) fn variables_used_in_instruction( instruction: &Instruction, dfg: &DataFlowGraph, ) -> Variables { let mut used = HashSet::default(); instruction.for_each_value(|value_id| { - let underlying_ids = collect_variables_of_value(value_id, dfg); - used.extend(underlying_ids); + if is_variable(value_id, dfg) { + used.insert(value_id); + } }); used } +/// Collect all [ValueId]s used in an [BasicBlock] which refer to [Variables]. +/// +/// Includes all the variables in the parameters, instructions and the terminator. fn variables_used_in_block(block: &BasicBlock, dfg: &DataFlowGraph) -> Variables { let mut used: Variables = block .instructions() @@ -90,241 +86,275 @@ fn variables_used_in_block(block: &BasicBlock, dfg: &DataFlowGraph) -> Variables if let Some(terminator) = block.terminator() { terminator.for_each_value(|value_id| { - used.extend(collect_variables_of_value(value_id, dfg)); + if is_variable(value_id, dfg) { + used.insert(value_id); + } }); } used } -type Variables = HashSet; - -fn compute_used_before_def( - block: &BasicBlock, - dfg: &DataFlowGraph, - defined_in_block: &Variables, -) -> Variables { - variables_used_in_block(block, dfg) - .into_iter() - .filter(|id| !defined_in_block.contains(id)) - .collect() -} - -type LastUses = HashMap; - /// A struct representing the liveness of variables throughout a function. #[derive(Default)] pub(crate) struct VariableLiveness { cfg: ControlFlowGraph, - post_order: PostOrder, - dominator_tree: DominatorTree, - /// The variables that are alive before the block starts executing + /// The variables that are alive before the block starts executing. live_in: HashMap, - /// The variables that stop being alive after each specific instruction + /// The variables that stop being alive after each specific instruction. last_uses: HashMap, - /// The list of block params the given block is defining. The order matters for the entry block, so it's a vec. + /// The list of block params the given block is defining. + /// The order matters for the entry block, so it's a vec. + /// + /// By _defining_ we mean the allocation of variables that appear in the parameter + /// list of some other block which this one immediately dominates, with values to be + /// assigned in the terminators of the predecessors of that block. param_definitions: HashMap>, } impl VariableLiveness { /// Computes the liveness of variables throughout a function. pub(crate) fn from_function(func: &Function, constants: &ConstantAllocation) -> Self { - let cfg = ControlFlowGraph::with_function(func); - let post_order = PostOrder::with_function(func); - let dominator_tree = DominatorTree::with_cfg_and_post_order(&cfg, &post_order); - - let mut instance = Self { - cfg, - post_order, - dominator_tree, + let loops = Loops::find_all(func); + + let back_edges: LoopMap = loops + .yet_to_unroll + .into_iter() + .map(|_loop| { + let back_edge = BackEdge { header: _loop.header, start: _loop.back_edge_start }; + let loop_body = _loop.blocks; + (back_edge, loop_body) + }) + .collect(); + + Self { + cfg: loops.cfg, live_in: HashMap::default(), last_uses: HashMap::default(), param_definitions: HashMap::default(), - }; - - instance.compute_block_param_definitions(func); - - instance.compute_live_in_of_blocks(func, constants); - - instance.compute_last_uses(func); - - instance + } + .compute_block_param_definitions(func, &loops.dom) + .compute_live_in_of_blocks(func, constants, back_edges) + .compute_last_uses(func) } - /// The set of values that are alive before the block starts executing + /// The set of values that are alive before the block starts executing. pub(crate) fn get_live_in(&self, block_id: &BasicBlockId) -> &Variables { - self.live_in.get(block_id).expect("Live ins should have been calculated") + self.live_in.get(block_id).expect("live_in should have been calculated") } - /// The set of values that are alive after the block has finished executed + /// The set of values that are alive after the block has finished executed. + /// + /// By definition this is the union of variables alive in the successors of the block. pub(crate) fn get_live_out(&self, block_id: &BasicBlockId) -> Variables { let mut live_out = HashSet::default(); for successor_id in self.cfg.successors(*block_id) { - live_out.extend(self.get_live_in(&successor_id)); + live_out.extend(self.get_live_in(&successor_id).iter().copied()); } live_out } - /// A map of instruction id to the set of values that die after the instruction has executed + /// For a given block, get the map of instruction ID to the set of values + /// that die after the instruction has executed. pub(crate) fn get_last_uses(&self, block_id: &BasicBlockId) -> &LastUses { - self.last_uses.get(block_id).expect("Last uses should have been calculated") + self.last_uses.get(block_id).expect("last_uses should have been calculated") } - /// Retrieves the list of block params the given block is defining. - /// Block params are defined before the block that owns them (since they are used by the predecessor blocks). They must be defined in the immediate dominator. - /// This is the last point where the block param can be allocated without it being allocated in different places in different branches. + /// Retrieves the list of block parameters the given block is defining. + /// + /// Block parameters are defined before the block that owns them + /// (since they are used by the predecessor blocks to pass values). + /// They must be defined in the immediate dominator, which is the + /// last point where the block parameter can be allocated + /// without it being allocated in different places in different branches. pub(crate) fn defined_block_params(&self, block_id: &BasicBlockId) -> Vec { self.param_definitions.get(block_id).cloned().unwrap_or_default() } - fn compute_block_param_definitions(&mut self, func: &Function) { - // Going in reverse post order to process the entry block first - let reverse_post_order = self.post_order.clone().into_vec_reverse(); + /// Compute [VariableLiveness::param_definitions]. + /// + /// Append the parameters of each block to the parameter definition list of + /// its immediate dominator. + fn compute_block_param_definitions(mut self, func: &Function, dom: &DominatorTree) -> Self { + assert!(self.param_definitions.is_empty(), "only define parameters once"); + + // Going in reverse post order to process the entry block first. + let reverse_post_order = PostOrder::with_function(func).into_vec_reverse(); for block in reverse_post_order { let params = func.dfg[block].parameters(); // If it has no dominator, it's the entry block - let dominator_block = - self.dominator_tree.immediate_dominator(block).unwrap_or(func.entry_block()); + let dominator_block = dom.immediate_dominator(block).unwrap_or(func.entry_block()); let definitions_for_the_dominator = self.param_definitions.entry(dominator_block).or_default(); definitions_for_the_dominator.extend(params.iter()); } - } - fn compute_live_in_of_blocks(&mut self, func: &Function, constants: &ConstantAllocation) { - let back_edges = find_back_edges(func, &self.cfg, &self.post_order); - - // First pass, propagate up the live_ins skipping back edges - self.compute_live_in_recursive(func, func.entry_block(), &back_edges, constants); + self + } - // Second pass, propagate header live_ins to the loop bodies - for back_edge in back_edges { - self.update_live_ins_within_loop(back_edge); + /// Compute [VariableLiveness::live_in]. + /// + /// Collect the variables which are alive before each block. + fn compute_live_in_of_blocks( + mut self, + func: &Function, + constants: &ConstantAllocation, + back_edges: LoopMap, + ) -> Self { + // First pass, propagate up the live_ins skipping back edges. + self.compute_live_in_recursive(func, func.entry_block(), constants, &back_edges); + + // Second pass, propagate header live_ins to the loop bodies. + for (back_edge, loop_body) in back_edges { + self.update_live_ins_within_loop(back_edge, loop_body); } + + self } + /// Starting with the entry block, traverse down all successors to compute their `live_in`, + /// then propagate the information back up towards the ancestors as `live_out`. + /// + /// The variables live at the *beginning* of a block are the variables used by the block, + /// plus the variables used by the successors of the block, minus the variables defined + /// in the block (by definition not alive at the beginning). fn compute_live_in_recursive( &mut self, func: &Function, block_id: BasicBlockId, - back_edges: &HashSet, constants: &ConstantAllocation, + back_edges: &LoopMap, ) { - let mut defined = self.compute_defined_variables(block_id, &func.dfg); - - defined.extend(constants.allocated_in_block(block_id)); - - let block: &BasicBlock = &func.dfg[block_id]; - - let used_before_def = compute_used_before_def(block, &func.dfg, &defined); - + let block = &func.dfg[block_id]; let mut live_out = HashSet::default(); + // Collect the `live_in` of successors; their union is the `live_in` of the parent. for successor_id in block.successors() { - if !back_edges.contains(&BackEdge { start: block_id, header: successor_id }) { - if !self.live_in.contains_key(&successor_id) { - self.compute_live_in_recursive(func, successor_id, back_edges, constants); - } - live_out.extend( - self.live_in - .get(&successor_id) - .expect("Live ins for successor should have been calculated"), - ); + // Skip back edges: do not revisit the header of the loop, to avoid infinite recursion. + // Because of this, we will have to revisit this block, to complete its definition of live_out, + // once we know what the live_in of the header is. + if back_edges.contains_key(&BackEdge { start: block_id, header: successor_id }) { + continue; + } + // If we haven't visited this successor yet, dive in. + if !self.live_in.contains_key(&successor_id) { + self.compute_live_in_recursive(func, successor_id, constants, back_edges); } + // Add the live_in of the successor to the union that forms the live_out of the parent. + // Note that because we skipped the back-edge, we couldn't call `get_live_out` here. + live_out.extend( + self.live_in + .get(&successor_id) + .expect("live_in for successor should have been calculated") + .iter() + .copied(), + ); } + // Based on the paper mentioned in the module docs, the definition would be: // live_in[BlockId] = before_def[BlockId] union (live_out[BlockId] - killed[BlockId]) - let passthrough_vars = live_out.difference(&defined).cloned().collect(); - self.live_in.insert(block_id, used_before_def.union(&passthrough_vars).cloned().collect()); + + // Variables used in this block, defined in this block or before. + let used = variables_used_in_block(block, &func.dfg); + + // Variables defined in this block are not alive at the beginning. + let defined = self.variables_defined_in_block(block_id, &func.dfg, constants); + + // Live at the beginning are the variables used, but not defined in this block, plus the ones + // it passes through to its successors, which are used by them, but not defined here. + // (Variables used by successors and defined in this block are part of `live_out`, but not `live_in`). + let live_in = used.union(&live_out).filter(|v| !defined.contains(v)).copied().collect(); + + self.live_in.insert(block_id, live_in); } - fn compute_defined_variables(&self, block_id: BasicBlockId, dfg: &DataFlowGraph) -> Variables { - let block: &BasicBlock = &dfg[block_id]; + /// Collects all the variables defined in a block, which includes: + /// * parameters of descendants this block immediately dominates + /// * the results of the instructions in the block + /// * constants which were allocated to this block + fn variables_defined_in_block( + &self, + block_id: BasicBlockId, + dfg: &DataFlowGraph, + constants: &ConstantAllocation, + ) -> Variables { + let block = &dfg[block_id]; let mut defined_vars = HashSet::default(); - for parameter in self.defined_block_params(&block_id) { - defined_vars.insert(parameter); - } + defined_vars.extend(self.defined_block_params(&block_id)); for instruction_id in block.instructions() { let result_values = dfg.instruction_results(*instruction_id); - for result_value in result_values { - defined_vars.insert(*result_value); - } + defined_vars.extend(result_values.iter().copied()); } + defined_vars.extend(constants.allocated_in_block(block_id)); + defined_vars } - fn update_live_ins_within_loop(&mut self, back_edge: BackEdge) { - let header_live_ins = self - .live_in - .get(&back_edge.header) - .expect("Live ins should have been calculated") - .clone(); + /// Once we know which variables are alive before the loop header, + /// we can append those variables to all of the loop's blocks. + /// Since we know that we have to come back + /// to the beginning of the loop, none of those blocks are allowed + /// anything but to keep these variables alive, so that the header + /// can use them again. + fn update_live_ins_within_loop( + &mut self, + back_edge: BackEdge, + loop_body: BTreeSet, + ) { + let header_live_in = self.get_live_in(&back_edge.header).clone(); - let body = self.compute_loop_body(back_edge); - for body_block_id in body { + for body_block_id in loop_body { self.live_in .get_mut(&body_block_id) .expect("Live ins should have been calculated") - .extend(&header_live_ins); - } - } - - fn compute_loop_body(&self, edge: BackEdge) -> HashSet { - let mut loop_blocks = HashSet::default(); - if edge.header == edge.start { - loop_blocks.insert(edge.header); - return loop_blocks; - } - loop_blocks.insert(edge.header); - loop_blocks.insert(edge.start); - - let mut stack = vec![edge.start]; - - while let Some(block) = stack.pop() { - for predecessor in self.cfg.predecessors(block) { - if !loop_blocks.contains(&predecessor) { - loop_blocks.insert(predecessor); - stack.push(predecessor); - } - } + .extend(header_live_in.iter().copied()); } - - loop_blocks } - fn compute_last_uses(&mut self, func: &Function) { + /// Compute [VariableLiveness::last_uses]. + /// + /// For each block, starting from the terminator than going backwards through the instructions, + /// take note of the first (technically last) instruction the value is used in. + fn compute_last_uses(mut self, func: &Function) -> Self { for block_id in func.reachable_blocks() { let block = &func.dfg[block_id]; let live_out = self.get_live_out(&block_id); + // Variables we have already visited, ie. they are used in "later" instructions or the terminator. let mut used_after: Variables = Default::default(); + // Variables becoming dead after each instruction. let mut block_last_uses: LastUses = Default::default(); - // First, handle the terminator + // First, handle the terminator; none of the instructions should cause these to go dead. if let Some(terminator_instruction) = block.terminator() { terminator_instruction.for_each_value(|value_id| { - let underlying_vars = collect_variables_of_value(value_id, &func.dfg); - used_after.extend(underlying_vars); + if is_variable(value_id, &func.dfg) { + used_after.insert(value_id); + } }); } - // Then, handle the instructions in reverse order to find the last use + // Then, handle the instructions in reverse order to find the last use. for instruction_id in block.instructions().iter().rev() { let instruction = &func.dfg[*instruction_id]; + // Collect the variables which will be dead after this instruction. let instruction_last_uses = variables_used_in_instruction(instruction, &func.dfg) .into_iter() .filter(|id| !used_after.contains(id) && !live_out.contains(id)) .collect(); - + // Remember that we have already handled these. used_after.extend(&instruction_last_uses); + // Remember that we can deallocate these after this instruction. block_last_uses.insert(*instruction_id, instruction_last_uses); } self.last_uses.insert(block_id, block_last_uses); } + + self } } @@ -336,10 +366,13 @@ mod test { use crate::brillig::brillig_gen::constant_allocation::ConstantAllocation; use crate::brillig::brillig_gen::variable_liveness::VariableLiveness; use crate::ssa::function_builder::FunctionBuilder; + use crate::ssa::ir::basic_block::BasicBlockId; use crate::ssa::ir::function::RuntimeType; use crate::ssa::ir::instruction::BinaryOp; use crate::ssa::ir::map::Id; use crate::ssa::ir::types::{NumericType, Type}; + use crate::ssa::ir::value::ValueId; + use crate::ssa::ssa_gen::Ssa; #[test] fn simple_back_propagation() { @@ -653,4 +686,47 @@ mod test { assert_eq!(liveness.get_live_in(&b1), &FxHashSet::from_iter([v1, v2].into_iter())); assert_eq!(liveness.get_live_in(&b2), &FxHashSet::from_iter([v1, v2].into_iter())); } + + /// A block parameter should be considered used, even if it's not actually used in the SSA. + #[test] + fn unused_block_parameter() { + let src = " + brillig(inline) fn main f0 { + b0(v0: u32, v1: u1): + jmpif v1 then: b1, else: b2 + b1(): + v7 = add v0, u32 10 + jmp b3(v0, v7) + b2(): + jmp b3(u32 1, u32 2) + b3(v2: u32, v3: u32): + return v3 + } + "; + let ssa = Ssa::from_str(src).unwrap(); + let func = ssa.main(); + let constants = ConstantAllocation::from_function(func); + let liveness = VariableLiveness::from_function(func, &constants); + + let [b0, b1, b2, b3] = block_ids(); + let [_v0, _v1, v2, v3] = value_ids(); + + // v2 and v3 are parameters of b3, but only v3 is used. + for p in [v2, v3] { + // Both should be allocated in b0, which is the immediate dominator of b3. + assert!(liveness.param_definitions[&b0].contains(&p), "{p} should be allocated in b0"); + for b in [b1, b2, b3] { + // Since they are defined in b0, they should be considered live all the way. + assert!(liveness.live_in[&b].contains(&p), "{p} should be live in {b}"); + } + } + } + + fn value_ids() -> [ValueId; N] { + std::array::from_fn(|i| ValueId::new(i as u32)) + } + + fn block_ids() -> [BasicBlockId; N] { + std::array::from_fn(|i| BasicBlockId::new(i as u32)) + } } diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir.rs index 5af25962d21..aacac530176 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_ir.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_ir.rs @@ -126,6 +126,7 @@ impl BrilligContext { MemoryAddress::Direct(GlobalSpace::start_with_layout(&self.layout())) } + /// If this flag is set, compile the array copy counter as a global. pub(crate) fn count_array_copies(&self) -> bool { self.count_arrays_copied } @@ -251,6 +252,7 @@ impl BrilligContext< /// Special brillig context to codegen compiler intrinsic shared procedures impl BrilligContext { + /// Create a [BrilligContext] with a [ScratchSpace] for passing procedure arguments. pub(crate) fn new_for_procedure( procedure_id: ProcedureId, options: &BrilligOptions, @@ -274,6 +276,7 @@ impl BrilligContext { /// Special brillig context to codegen global values initialization impl BrilligContext { + /// Create a [BrilligContext] with a [GlobalSpace] for memory allocations. pub(crate) fn new_for_global_init( options: &BrilligOptions, entry_point: FunctionId, @@ -292,6 +295,7 @@ impl BrilligContext { } } + /// Total size of the global memory space. pub(crate) fn global_space_size(&self) -> usize { // `GlobalSpace::start` is inclusive so we must add one to get the accurate total global memory size (self.registers().max_memory_address() + 1) - self.registers().start() diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/artifact.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/artifact.rs index 3309049d44f..401bac1acae 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_ir/artifact.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_ir/artifact.rs @@ -133,7 +133,7 @@ pub(crate) struct Label { } impl Label { - pub(crate) fn add_section(&self, section: usize) -> Self { + pub(crate) fn with_section(&self, section: usize) -> Self { Label { label_type: self.label_type.clone(), section: Some(section) } } diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/instructions.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/instructions.rs index 66eae9181b1..b48b0e21038 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_ir/instructions.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_ir/instructions.rs @@ -276,7 +276,7 @@ impl BrilligContext< /// Internal function used to compute the section labels fn compute_section_label(&self, section: usize) -> Label { - self.context_label.add_section(section) + self.context_label.with_section(section) } /// Returns the current section label diff --git a/compiler/noirc_evaluator/src/brillig/mod.rs b/compiler/noirc_evaluator/src/brillig/mod.rs index c28fe43e7fa..d5348e3dfbc 100644 --- a/compiler/noirc_evaluator/src/brillig/mod.rs +++ b/compiler/noirc_evaluator/src/brillig/mod.rs @@ -45,19 +45,21 @@ pub struct BrilligOptions { pub layout: LayoutConfig, } -/// Context structure for the brillig pass. -/// It stores brillig-related data required for brillig generation. +/// Context structure for the Brillig pass. +/// It stores Brillig-related data required for Brillig generation. #[derive(Default, Clone)] pub struct Brillig { /// Maps SSA function labels to their brillig artifact ssa_function_to_brillig: HashMap>, - pub call_stacks: CallStackHelper, + call_stacks: CallStackHelper, + /// Bytecode for globals for each Brillig entry point. globals: HashMap>, + /// The size of the global space for each Brillig entry point. globals_memory_size: HashMap, } impl Brillig { - /// Compiles a function into brillig and store the compilation artifacts + /// Compiles a function into brillig and store the compilation artifacts. pub(crate) fn compile( &mut self, func: &Function, @@ -76,7 +78,7 @@ impl Brillig { self.ssa_function_to_brillig.insert(func.id(), obj); } - /// Finds a brillig artifact by its label + /// Finds a brillig artifact by its label. pub(crate) fn find_by_label( &self, function_label: Label, @@ -127,6 +129,10 @@ impl Brillig { brillig_context.artifact() } + + pub fn call_stacks(&self) -> &CallStackHelper { + &self.call_stacks + } } impl std::ops::Index for Brillig { @@ -140,10 +146,9 @@ impl Ssa { /// Compile Brillig functions and ACIR functions reachable from them. #[tracing::instrument(level = "trace", skip_all)] pub fn to_brillig(&self, options: &BrilligOptions) -> Brillig { - let used_globals_map = self.used_globals_in_functions(); - - // Collect all the function ids that are reachable from brillig - // That means all the functions marked as brillig and ACIR functions called by them + // Collect all the function ids that are reachable from Brillig. + // That means all the functions marked as Brillig and ACIR functions called by them, + // but we should already have monomorphized ACIR functions as a Brillig as well. let brillig_reachable_function_ids = self .functions .iter() @@ -156,24 +161,25 @@ impl Ssa { return brillig; } - let mut brillig_globals = BrilligGlobals::new(self, used_globals_map, self.main_id); - // SSA Globals are computed once at compile time and shared across all functions, // thus we can just fetch globals from the main function. // This same globals graph will then be used to declare Brillig globals for the respective entry points. let globals = (*self.functions[&self.main_id].dfg.globals).clone(); let globals_dfg = DataFlowGraph::from(globals); - brillig_globals.declare_globals(&globals_dfg, &mut brillig, options); + + // Produce the globals Brillig bytecode and variable allocation for each entry point. + let brillig_globals = BrilligGlobals::init(self, self.main_id).declare_globals( + &globals_dfg, + &mut brillig, + options, + ); for brillig_function_id in brillig_reachable_function_ids { - let empty_allocations = HashMap::default(); - let empty_const_allocations = HashMap::default(); - let (globals_allocations, hoisted_constant_allocations) = brillig_globals - .get_brillig_globals(brillig_function_id) - .unwrap_or((&empty_allocations, &empty_const_allocations)); + let (globals_allocations, hoisted_constant_allocations) = + brillig_globals.get_brillig_globals(brillig_function_id); let func = &self.functions[&brillig_function_id]; - let is_entry_point = brillig_globals.entry_points().contains_key(&brillig_function_id); + let is_entry_point = brillig_globals.is_entry_point(&brillig_function_id); brillig.compile( func, diff --git a/compiler/noirc_evaluator/src/ssa/opt/evaluate_static_assert_and_assert_constant.rs b/compiler/noirc_evaluator/src/ssa/opt/evaluate_static_assert_and_assert_constant.rs index 6f9269c06aa..846d42a769d 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/evaluate_static_assert_and_assert_constant.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/evaluate_static_assert_and_assert_constant.rs @@ -26,12 +26,11 @@ use crate::{ instruction::{Instruction, InstructionId, Intrinsic}, value::ValueId, }, + opt::Loops, ssa_gen::Ssa, }, }; -use super::unrolling::Loops; - impl Ssa { /// See [`evaluate_static_assert_and_assert_constant`][self] module for more information. #[tracing::instrument(level = "trace", skip(self))] diff --git a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs index a0d899f31c5..e4949e54f75 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs @@ -207,7 +207,7 @@ fn flatten_cfg_pre_check(function: &Function) { if !function.runtime().is_acir() { return; } - let loops = super::unrolling::Loops::find_all(function); + let loops = super::Loops::find_all(function); assert_eq!(loops.yet_to_unroll.len(), 0); for block in function.reachable_blocks() { diff --git a/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs b/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs index d6206ff3575..ffade5e5a52 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs @@ -892,11 +892,11 @@ mod test { use crate::ssa::ir::function::RuntimeType; use crate::ssa::ir::instruction::{Instruction, Intrinsic, TerminatorInstruction}; use crate::ssa::ir::types::Type; + use crate::ssa::opt::Loops; use crate::ssa::opt::loop_invariant::{ CanBeHoistedResult, LoopContext, LoopInvariantContext, can_be_hoisted, }; use crate::ssa::opt::pure::Purity; - use crate::ssa::opt::unrolling::Loops; use crate::ssa::opt::{assert_normalized_ssa_equals, assert_ssa_does_not_change}; use acvm::AcirField; use noirc_frontend::monomorphization::ast::InlineType; diff --git a/compiler/noirc_evaluator/src/ssa/opt/mod.rs b/compiler/noirc_evaluator/src/ssa/opt/mod.rs index ce50422946c..b8ad1905447 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/mod.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/mod.rs @@ -41,6 +41,7 @@ mod unrolling; pub use constant_folding::DEFAULT_MAX_ITER as CONSTANT_FOLDING_MAX_ITER; pub use inlining::MAX_INSTRUCTIONS as INLINING_MAX_INSTRUCTIONS; +pub(crate) use unrolling::Loops; /// Asserts that the given SSA, after normalizing its IDs and printing it, /// is equal to the expected string. Normalization is done so the IDs don't diff --git a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs index f4faffcb7ae..3e4b4527e1f 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/unrolling.rs @@ -220,25 +220,27 @@ impl Function { /// Describe the blocks that constitute up a loop. #[derive(Debug)] -pub(super) struct Loop { +pub(crate) struct Loop { /// The header block of a loop is the block which dominates all the /// other blocks in the loop. - pub(super) header: BasicBlockId, + pub(crate) 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. - back_edge_start: BasicBlockId, + pub(crate) back_edge_start: BasicBlockId, /// All the blocks contained within the loop, including `header` and `back_edge_start`. - pub(super) blocks: BTreeSet, + pub(crate) blocks: BTreeSet, } /// All the unrolled loops in the SSA. -pub(super) struct Loops { +pub(crate) struct Loops { /// Loops that haven't been unrolled yet, which is all the loops currently in the CFG. - pub(super) yet_to_unroll: Vec, + pub(crate) yet_to_unroll: Vec, /// The CFG so we can query the predecessors of blocks when needed. - pub(super) cfg: ControlFlowGraph, + pub(crate) cfg: ControlFlowGraph, + /// The [DominatorTree] used during the discovery of loops. + pub(crate) dom: DominatorTree, } impl Loops { @@ -273,7 +275,7 @@ impl Loops { /// /// Returns all groups of blocks that look like a loop, even if we might not be able to unroll them, /// which we can use to check whether we were able to unroll all blocks. - pub(super) fn find_all(function: &Function) -> Self { + pub(crate) fn find_all(function: &Function) -> Self { let cfg = ControlFlowGraph::with_function(function); let post_order = PostOrder::with_function(function); let mut dom_tree = DominatorTree::with_cfg_and_post_order(&cfg, &post_order); @@ -296,7 +298,7 @@ impl Loops { // their loop range. We will start popping loops from the back. loops.sort_by_key(|loop_| loop_.blocks.len()); - Self { yet_to_unroll: loops, cfg } + Self { yet_to_unroll: loops, cfg, dom: dom_tree } } } diff --git a/compiler/noirc_evaluator/src/ssa/ssa_gen/program.rs b/compiler/noirc_evaluator/src/ssa/ssa_gen/program.rs index 312378c5157..8dfaa0f62c3 100644 --- a/compiler/noirc_evaluator/src/ssa/ssa_gen/program.rs +++ b/compiler/noirc_evaluator/src/ssa/ssa_gen/program.rs @@ -99,6 +99,7 @@ impl Ssa { function == self.main_id || self.functions[&function].runtime().is_entry_point() } + /// Collect all the global value IDs used per function. pub(crate) fn used_globals_in_functions(&self) -> HashMap> { fn add_value_to_globals_if_global( function: &Function,