Skip to content

coverage: Don't rely on the custom traversal to find enclosing loops #132091

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.

Sign up for GitHub

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

Merged
merged 1 commit into from
Oct 27, 2024
Merged

coverage: Don't rely on the custom traversal to find enclosing loops #132091

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
31 changes: 13 additions & 18 deletions compiler/rustc_mir_transform/src/coverage/counters.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,13 +280,12 @@ impl<'a> CountersBuilder<'a> {
//
// The traversal tries to ensure that, when a loop is encountered, all
// nodes within the loop are visited before visiting any nodes outside
// the loop. It also keeps track of which loop(s) the traversal is
// currently inside.
// the loop.
let mut traversal = TraverseCoverageGraphWithLoops::new(self.graph);
while let Some(bcb) = traversal.next() {
let _span = debug_span!("traversal", ?bcb).entered();
if self.bcb_needs_counter.contains(bcb) {
self.make_node_counter_and_out_edge_counters(&traversal, bcb);
self.make_node_counter_and_out_edge_counters(bcb);
}
}

Expand All @@ -299,12 +298,8 @@ impl<'a> CountersBuilder<'a> {

/// Make sure the given node has a node counter, and then make sure each of
/// its out-edges has an edge counter (if appropriate).
#[instrument(level = "debug", skip(self, traversal))]
fn make_node_counter_and_out_edge_counters(
&mut self,
traversal: &TraverseCoverageGraphWithLoops<'_>,
from_bcb: BasicCoverageBlock,
) {
#[instrument(level = "debug", skip(self))]
fn make_node_counter_and_out_edge_counters(&mut self, from_bcb: BasicCoverageBlock) {
// First, ensure that this node has a counter of some kind.
// We might also use that counter to compute one of the out-edge counters.
let node_counter = self.get_or_make_node_counter(from_bcb);
Expand Down Expand Up @@ -337,8 +332,7 @@ impl<'a> CountersBuilder<'a> {

// If there are out-edges without counters, choose one to be given an expression
// (computed from this node and the other out-edges) instead of a physical counter.
let Some(target_bcb) =
self.choose_out_edge_for_expression(traversal, &candidate_successors)
let Some(target_bcb) = self.choose_out_edge_for_expression(from_bcb, &candidate_successors)
else {
return;
};
Expand Down Expand Up @@ -462,12 +456,12 @@ impl<'a> CountersBuilder<'a> {
/// choose one to be given a counter expression instead of a physical counter.
fn choose_out_edge_for_expression(
&self,
traversal: &TraverseCoverageGraphWithLoops<'_>,
from_bcb: BasicCoverageBlock,
candidate_successors: &[BasicCoverageBlock],
) -> Option<BasicCoverageBlock> {
// Try to find a candidate that leads back to the top of a loop,
// because reloop edges tend to be executed more times than loop-exit edges.
if let Some(reloop_target) = self.find_good_reloop_edge(traversal, &candidate_successors) {
if let Some(reloop_target) = self.find_good_reloop_edge(from_bcb, &candidate_successors) {
debug!("Selecting reloop target {reloop_target:?} to get an expression");
return Some(reloop_target);
}
Expand All @@ -486,7 +480,7 @@ impl<'a> CountersBuilder<'a> {
/// for them to be able to avoid a physical counter increment.
fn find_good_reloop_edge(
&self,
traversal: &TraverseCoverageGraphWithLoops<'_>,
from_bcb: BasicCoverageBlock,
candidate_successors: &[BasicCoverageBlock],
) -> Option<BasicCoverageBlock> {
// If there are no candidates, avoid iterating over the loop stack.
Expand All @@ -495,14 +489,15 @@ impl<'a> CountersBuilder<'a> {
}

// Consider each loop on the current traversal context stack, top-down.
for reloop_bcbs in traversal.reloop_bcbs_per_loop() {
for loop_header_node in self.graph.loop_headers_containing(from_bcb) {
// Try to find a candidate edge that doesn't exit this loop.
for &target_bcb in candidate_successors {
// An edge is a reloop edge if its target dominates any BCB that has
// an edge back to the loop header. (Otherwise it's an exit edge.)
let is_reloop_edge = reloop_bcbs
.iter()
.any(|&reloop_bcb| self.graph.dominates(target_bcb, reloop_bcb));
let is_reloop_edge = self
.graph
.reloop_predecessors(loop_header_node)
.any(|reloop_bcb| self.graph.dominates(target_bcb, reloop_bcb));
if is_reloop_edge {
// We found a good out-edge to be given an expression.
return Some(target_bcb);
Expand Down
117 changes: 72 additions & 45 deletions compiler/rustc_mir_transform/src/coverage/graph.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
use std::cmp::Ordering;
use std::collections::VecDeque;
use std::iter;
use std::ops::{Index, IndexMut};

use rustc_data_structures::captures::Captures;
use rustc_data_structures::fx::FxHashSet;
use rustc_data_structures::graph::dominators::{self, Dominators};
use rustc_data_structures::graph::dominators::Dominators;
use rustc_data_structures::graph::{self, DirectedGraph, StartNode};
use rustc_index::IndexVec;
use rustc_index::bit_set::BitSet;
Expand All @@ -20,11 +21,17 @@ pub(crate) struct CoverageGraph {
bb_to_bcb: IndexVec<BasicBlock, Option<BasicCoverageBlock>>,
pub(crate) successors: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>,
pub(crate) predecessors: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>,

dominators: Option<Dominators<BasicCoverageBlock>>,
/// Allows nodes to be compared in some total order such that _if_
/// `a` dominates `b`, then `a < b`. If neither node dominates the other,
/// their relative order is consistent but arbitrary.
dominator_order_rank: IndexVec<BasicCoverageBlock, u32>,
/// A loop header is a node that dominates one or more of its predecessors.
is_loop_header: BitSet<BasicCoverageBlock>,
/// For each node, the loop header node of its nearest enclosing loop.
/// This forms a linked list that can be traversed to find all enclosing loops.
enclosing_loop_header: IndexVec<BasicCoverageBlock, Option<BasicCoverageBlock>>,
}

impl CoverageGraph {
Expand Down Expand Up @@ -66,17 +73,38 @@ impl CoverageGraph {
predecessors,
dominators: None,
dominator_order_rank: IndexVec::from_elem_n(0, num_nodes),
is_loop_header: BitSet::new_empty(num_nodes),
enclosing_loop_header: IndexVec::from_elem_n(None, num_nodes),
};
assert_eq!(num_nodes, this.num_nodes());

this.dominators = Some(dominators::dominators(&this));
// Set the dominators first, because later init steps rely on them.
this.dominators = Some(graph::dominators::dominators(&this));

// The dominator rank of each node is just its index in a reverse-postorder traversal.
let reverse_post_order = graph::iterate::reverse_post_order(&this, this.start_node());
// Iterate over all nodes, such that dominating nodes are visited before
// the nodes they dominate. Either preorder or reverse postorder is fine.
let dominator_order = graph::iterate::reverse_post_order(&this, this.start_node());
// The coverage graph is created by traversal, so all nodes are reachable.
assert_eq!(reverse_post_order.len(), this.num_nodes());
for (rank, bcb) in (0u32..).zip(reverse_post_order) {
assert_eq!(dominator_order.len(), this.num_nodes());
for (rank, bcb) in (0u32..).zip(dominator_order) {
// The dominator rank of each node is its index in a dominator-order traversal.
this.dominator_order_rank[bcb] = rank;

// A node is a loop header if it dominates any of its predecessors.
if this.reloop_predecessors(bcb).next().is_some() {
this.is_loop_header.insert(bcb);
}

// If the immediate dominator is a loop header, that's our enclosing loop.
// Otherwise, inherit the immediate dominator's enclosing loop.
// (Dominator order ensures that we already processed the dominator.)
if let Some(dom) = this.dominators().immediate_dominator(bcb) {
this.enclosing_loop_header[bcb] = this
.is_loop_header
.contains(dom)
.then_some(dom)
.or_else(|| this.enclosing_loop_header[dom]);
}
}

// The coverage graph's entry-point node (bcb0) always starts with bb0,
Expand Down Expand Up @@ -172,9 +200,14 @@ impl CoverageGraph {
if bb.index() < self.bb_to_bcb.len() { self.bb_to_bcb[bb] } else { None }
}

#[inline(always)]
fn dominators(&self) -> &Dominators<BasicCoverageBlock> {
self.dominators.as_ref().unwrap()
}

#[inline(always)]
pub(crate) fn dominates(&self, dom: BasicCoverageBlock, node: BasicCoverageBlock) -> bool {
self.dominators.as_ref().unwrap().dominates(dom, node)
self.dominators().dominates(dom, node)
}

#[inline(always)]
Expand Down Expand Up @@ -214,6 +247,36 @@ impl CoverageGraph {
None
}
}

/// For each loop that contains the given node, yields the "loop header"
/// node representing that loop, from innermost to outermost. If the given
/// node is itself a loop header, it is yielded first.
pub(crate) fn loop_headers_containing(
&self,
bcb: BasicCoverageBlock,
) -> impl Iterator<Item = BasicCoverageBlock> + Captures<'_> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just a side note here, the Rust 1.82 release notes mention that we now have dedicated use<> syntax to replace Captures, and this eventually becoming the default in Rust 2024.
Though I believe + '_ would also work here, as the &self lifelime is the only one captured here?

I also don’t really know what the current conventions for compiler code are related to all this.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, is that stable now? Nice.

I think I'm going to stick with Captures for now. After the compiler switches over to Rust 2024, someone will probably go through and try to rip out unnecessary uses of Captures, and it will be easier to find and migrate this code if it matches the prevailing style in the rest of the compiler.

let self_if_loop_header = self.is_loop_header.contains(bcb).then_some(bcb).into_iter();

let mut curr = Some(bcb);
let strictly_enclosing = iter::from_fn(move || {
let enclosing = self.enclosing_loop_header[curr?];
curr = enclosing;
enclosing
});

self_if_loop_header.chain(strictly_enclosing)
}

/// For the given node, yields the subset of its predecessor nodes that
/// it dominates. If that subset is non-empty, the node is a "loop header",
/// and each of those predecessors represents an in-edge that jumps back to
/// the top of its loop.
pub(crate) fn reloop_predecessors(
&self,
to_bcb: BasicCoverageBlock,
) -> impl Iterator<Item = BasicCoverageBlock> + Captures<'_> {
self.predecessors[to_bcb].iter().copied().filter(move |&pred| self.dominates(to_bcb, pred))
}
}

impl Index<BasicCoverageBlock> for CoverageGraph {
Expand Down Expand Up @@ -439,15 +502,12 @@ struct TraversalContext {
pub(crate) struct TraverseCoverageGraphWithLoops<'a> {
basic_coverage_blocks: &'a CoverageGraph,

backedges: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>,
context_stack: Vec<TraversalContext>,
visited: BitSet<BasicCoverageBlock>,
}

impl<'a> TraverseCoverageGraphWithLoops<'a> {
pub(crate) fn new(basic_coverage_blocks: &'a CoverageGraph) -> Self {
let backedges = find_loop_backedges(basic_coverage_blocks);

let worklist = VecDeque::from([basic_coverage_blocks.start_node()]);
let context_stack = vec![TraversalContext { loop_header: None, worklist }];

Expand All @@ -456,17 +516,7 @@ impl<'a> TraverseCoverageGraphWithLoops<'a> {
// of the stack as loops are entered, and popped off of the stack when a loop's worklist is
// exhausted.
let visited = BitSet::new_empty(basic_coverage_blocks.num_nodes());
Self { basic_coverage_blocks, backedges, context_stack, visited }
}

/// For each loop on the loop context stack (top-down), yields a list of BCBs
/// within that loop that have an outgoing edge back to the loop header.
pub(crate) fn reloop_bcbs_per_loop(&self) -> impl Iterator<Item = &[BasicCoverageBlock]> {
self.context_stack
.iter()
.rev()
.filter_map(|context| context.loop_header)
.map(|header_bcb| self.backedges[header_bcb].as_slice())
Self { basic_coverage_blocks, context_stack, visited }
}

pub(crate) fn next(&mut self) -> Option<BasicCoverageBlock> {
Expand All @@ -488,7 +538,7 @@ impl<'a> TraverseCoverageGraphWithLoops<'a> {
}
debug!("Visiting {bcb:?}");

if self.backedges[bcb].len() > 0 {
if self.basic_coverage_blocks.is_loop_header.contains(bcb) {
debug!("{bcb:?} is a loop header! Start a new TraversalContext...");
self.context_stack
.push(TraversalContext { loop_header: Some(bcb), worklist: VecDeque::new() });
Expand Down Expand Up @@ -566,29 +616,6 @@ impl<'a> TraverseCoverageGraphWithLoops<'a> {
}
}

fn find_loop_backedges(
basic_coverage_blocks: &CoverageGraph,
) -> IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>> {
let num_bcbs = basic_coverage_blocks.num_nodes();
let mut backedges = IndexVec::from_elem_n(Vec::<BasicCoverageBlock>::new(), num_bcbs);

// Identify loops by their backedges.
for (bcb, _) in basic_coverage_blocks.iter_enumerated() {
for &successor in &basic_coverage_blocks.successors[bcb] {
if basic_coverage_blocks.dominates(successor, bcb) {
let loop_header = successor;
let backedge_from_bcb = bcb;
debug!(
"Found BCB backedge: {:?} -> loop_header: {:?}",
backedge_from_bcb, loop_header
);
backedges[loop_header].push(backedge_from_bcb);
}
}
}
backedges
}

fn short_circuit_preorder<'a, 'tcx, F, Iter>(
body: &'a mir::Body<'tcx>,
filtered_successors: F,
Expand Down
Loading