Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Hoist constant allocation outside of loops #6158

Merged
merged 2 commits into from
Sep 26, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -26,20 +26,23 @@ pub(crate) struct ConstantAllocation {
constant_usage: HashMap<ValueId, HashMap<BasicBlockId, Vec<InstructionLocation>>>,
allocation_points: HashMap<BasicBlockId, HashMap<InstructionLocation, Vec<ValueId>>>,
dominator_tree: DominatorTree,
blocks_within_loops: HashSet<BasicBlockId>,
}

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
}
Expand Down Expand Up @@ -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()
Expand All @@ -114,29 +117,105 @@ impl ConstantAllocation {
};

self.allocation_points
.entry(common_dominator)
.entry(allocation_point)
.or_default()
.entry(location)
.or_default()
.push(*constant_id);
}
}

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(&current_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<BasicBlockId> {
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(
TomAFrench marked this conversation as resolved.
Show resolved Hide resolved
header: BasicBlockId,
back_edge_start: BasicBlockId,
cfg: &ControlFlowGraph,
) -> HashSet<BasicBlockId> {
let mut blocks = HashSet::default();
blocks.insert(header);

let mut insert = |block, stack: &mut Vec<BasicBlockId>| {
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
}
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down
Loading