From 180bfc99944cd42b3f44048213458d1399687cef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81lvaro=20Rodr=C3=ADguez?= Date: Thu, 26 Sep 2024 18:44:15 +0200 Subject: [PATCH] feat: Hoist constant allocation outside of loops (#6158) # Description ## Problem\* Avoids the runtime cost of reinitializing constants inside loops. ## Summary\* ## Additional Context ## Documentation\* Check one: - [x] No documentation needed. - [ ] Documentation included in this PR. - [ ] **[For Experimental Features]** Documentation to be submitted in a separate PR. # PR Checklist\* - [x] I have tested the changes locally. - [x] I have formatted the changes with [Prettier](https://prettier.io/) and/or `cargo fmt` on default settings. --- .../brillig_gen/constant_allocation.rs | 115 +++++++++++++++--- .../brillig/brillig_gen/variable_liveness.rs | 20 +-- 2 files changed, 110 insertions(+), 25 deletions(-) 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 cf484fa5038..f9ded224b33 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/constant_allocation.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/constant_allocation.rs @@ -1,7 +1,7 @@ //! This module analyzes the usage of constants in a given function and decides an allocation point for them. //! The allocation point will be the common dominator of all the places where the constant is used. //! By allocating in the common dominator, we can cache the constants for all subsequent uses. -use fxhash::FxHashMap as HashMap; +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; use crate::ssa::ir::{ basic_block::BasicBlockId, @@ -26,20 +26,23 @@ pub(crate) struct ConstantAllocation { constant_usage: HashMap>>, allocation_points: HashMap>>, dominator_tree: DominatorTree, + blocks_within_loops: HashSet, } impl ConstantAllocation { pub(crate) fn from_function(func: &Function) -> 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 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 mut instance = ConstantAllocation { constant_usage: HashMap::default(), allocation_points: HashMap::default(), dominator_tree, + blocks_within_loops, }; instance.collect_constant_usage(func); - instance.decide_allocation_points(); + instance.decide_allocation_points(func); instance } @@ -95,16 +98,16 @@ impl ConstantAllocation { } } - fn decide_allocation_points(&mut self) { + 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.iter().map(|(block_id, _)| *block_id).collect(); - let common_dominator = self.common_dominator(&block_ids); + let allocation_point = self.decide_allocation_point(*constant_id, &block_ids, func); - // If the common dominator is one of the places where it's used, we take the first usage in the common dominator. - // Otherwise, we allocate it at the terminator of the common dominator. + // If the allocation point is one of the places where it's used, we take the first usage in the allocation point. + // Otherwise, we allocate it at the terminator of the allocation point. let location = if let Some(locations_in_common_dominator) = - usage_in_blocks.get(&common_dominator) + usage_in_blocks.get(&allocation_point) { *locations_in_common_dominator .first() @@ -114,7 +117,7 @@ impl ConstantAllocation { }; self.allocation_points - .entry(common_dominator) + .entry(allocation_point) .or_default() .entry(location) .or_default() @@ -122,21 +125,97 @@ impl ConstantAllocation { } } - fn common_dominator(&self, block_ids: &[BasicBlockId]) -> BasicBlockId { - if block_ids.len() == 1 { - return block_ids[0]; - } - - let mut common_dominator = block_ids[0]; + fn decide_allocation_point( + &self, + constant_id: ValueId, + blocks_where_is_used: &[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); + } - for block_id in block_ids.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 + if func.dfg.is_constant(constant_id) { + self.exit_loops(common_dominator) + } else { + common_dominator } + } - common_dominator + /// Returns the nearest dominator that is outside of any loop. + fn exit_loops(&self, block: BasicBlockId) -> BasicBlockId { + let mut current_block = block; + while self.blocks_within_loops.contains(¤t_block) { + current_block = self + .dominator_tree + .immediate_dominator(current_block) + .expect("No dominator found when trying to allocate a constant outside of a loop"); + } + current_block } } pub(crate) fn is_constant_value(id: ValueId, dfg: &DataFlowGraph) -> bool { matches!(&dfg[dfg.resolve(id)], Value::NumericConstant { .. } | Value::Array { .. }) } + +/// 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, +) -> HashSet { + let mut blocks_in_loops = 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 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, +) -> HashSet { + let mut blocks = HashSet::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 73e88cee676..92595292bf0 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/variable_liveness.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/variable_liveness.rs @@ -570,28 +570,34 @@ mod test { let liveness = VariableLiveness::from_function(func, &constants); assert!(liveness.get_live_in(&func.entry_block()).is_empty()); - assert_eq!(liveness.get_live_in(&b1), &FxHashSet::from_iter([v0, v1, v3, v4].into_iter())); + assert_eq!( + liveness.get_live_in(&b1), + &FxHashSet::from_iter([v0, v1, v3, v4, twenty_seven, one].into_iter()) + ); assert_eq!(liveness.get_live_in(&b3), &FxHashSet::from_iter([v3].into_iter())); - assert_eq!(liveness.get_live_in(&b2), &FxHashSet::from_iter([v0, v1, v3, v4].into_iter())); + assert_eq!( + liveness.get_live_in(&b2), + &FxHashSet::from_iter([v0, v1, v3, v4, twenty_seven, one].into_iter()) + ); assert_eq!( liveness.get_live_in(&b4), - &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7].into_iter()) + &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7, twenty_seven, one].into_iter()) ); assert_eq!( liveness.get_live_in(&b6), - &FxHashSet::from_iter([v0, v1, v3, v4, one].into_iter()) + &FxHashSet::from_iter([v0, v1, v3, v4, twenty_seven, one].into_iter()) ); assert_eq!( liveness.get_live_in(&b5), - &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7, one].into_iter()) + &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7, twenty_seven, one].into_iter()) ); assert_eq!( liveness.get_live_in(&b7), - &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7, one].into_iter()) + &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7, twenty_seven, one].into_iter()) ); assert_eq!( liveness.get_live_in(&b8), - &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7, one].into_iter()) + &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7, twenty_seven, one].into_iter()) ); let block_3 = &func.dfg[b3];