Skip to content

[MIR] Enhance the SimplifyCfg pass to merge consecutive blocks #29838

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

Closed
wants to merge 1 commit into from
Closed
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
80 changes: 52 additions & 28 deletions src/librustc_mir/transform/simplify_cfg.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,28 +21,7 @@ impl SimplifyCfg {
SimplifyCfg
}

fn remove_dead_blocks(&self, mir: &mut Mir) {
let mut seen = vec![false; mir.basic_blocks.len()];

// These blocks are always required.
seen[START_BLOCK.index()] = true;
seen[END_BLOCK.index()] = true;
seen[DIVERGE_BLOCK.index()] = true;

let mut worklist = vec![START_BLOCK];
while let Some(bb) = worklist.pop() {
for succ in mir.basic_block_data(bb).terminator.successors() {
if !seen[succ.index()] {
seen[succ.index()] = true;
worklist.push(*succ);
}
}
}

util::retain_basic_blocks(mir, &seen);
}

fn remove_goto_chains(&self, mir: &mut Mir) -> bool {
fn merge_consecutive_blocks(&self, mir: &mut Mir) -> bool {

// Find the target at the end of the jump chain, return None if there is a loop
fn final_target(mir: &Mir, mut target: BasicBlock) -> Option<BasicBlock> {
Expand All @@ -65,25 +44,71 @@ impl SimplifyCfg {
Some(target)
}

let mut predecessor_map = util::build_predecessor_map(mir);

let mut changed = false;
for bb in mir.all_basic_blocks() {
let mut seen = vec![false; mir.basic_blocks.len()];
let mut worklist = vec![START_BLOCK];
while let Some(bb) = worklist.pop() {
// Temporarily swap out the terminator we're modifying to keep borrowck happy
let mut terminator = Terminator::Diverge;
mem::swap(&mut terminator, &mut mir.basic_block_data_mut(bb).terminator);

// Shortcut chains of empty blocks that just jump from one to the next
for target in terminator.successors_mut() {
let new_target = match final_target(mir, *target) {
Some(new_target) => new_target,
None if mir.basic_block_data(bb).statements.is_empty() => bb,
None => continue
};
changed |= *target != new_target;
*target = new_target;

if *target != new_target {
changed = true;
predecessor_map.remove_predecessor(*target, bb);
predecessor_map.add_predecessor(new_target, bb);
*target = new_target;
}
}

mir.basic_block_data_mut(bb).terminator = terminator;
// See if we can merge the target block into this one
match terminator {
Terminator::Goto { target } if target.index() > DIVERGE_BLOCK.index() &&
predecessor_map.predecessors(target).len() == 1 => {
changed = true;
let mut other_data = BasicBlockData {
statements: Vec::new(),
terminator: Terminator::Goto { target: target}
};
mem::swap(&mut other_data, mir.basic_block_data_mut(target));

predecessor_map.replace_predecessor(target, bb, target);
for succ in other_data.terminator.successors() {
predecessor_map.replace_predecessor(*succ, target, bb);
}

let data = mir.basic_block_data_mut(bb);
data.statements.append(&mut other_data.statements);
mem::swap(&mut data.terminator, &mut other_data.terminator);
}
_ => mir.basic_block_data_mut(bb).terminator = terminator
}

for succ in mir.basic_block_data(bb).terminator.successors() {
if !seen[succ.index()] {
seen[succ.index()] = true;
worklist.push(*succ);
}
}
}

// These blocks must be retained, so mark them seen even if we didn't see them
seen[START_BLOCK.index()] = true;
seen[END_BLOCK.index()] = true;
seen[DIVERGE_BLOCK.index()] = true;

// Now get rid of all the blocks we never saw
util::retain_basic_blocks(mir, &seen);

changed
}

Expand Down Expand Up @@ -125,8 +150,7 @@ impl MirPass for SimplifyCfg {
let mut changed = true;
while changed {
changed = self.simplify_branches(mir);
changed |= self.remove_goto_chains(mir);
self.remove_dead_blocks(mir);
changed |= self.merge_consecutive_blocks(mir);
Copy link
Contributor

Choose a reason for hiding this comment

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

I think that doing merge_consecutive_blocks without a remove_dead_blocks pass beforehand is missing out on some optimizations. For example:

Entry -> B <- A

Since A is dead, B only has one "real" predecessor, and so can be merged into Entry. However, since dead blocks aren't removed until the afterwards, merge_consecutive_blocks sees B as having 2 predecessors and doesn't merge it into Entry.

}

// FIXME: Should probably be moved into some kind of pass manager
Expand Down
68 changes: 68 additions & 0 deletions src/librustc_mir/transform/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
// except according to those terms.

use repr::*;
use visit::*;

/// Update basic block ids in all terminators using the given replacements,
/// useful e.g. after removal of several basic blocks to update all terminators
Expand Down Expand Up @@ -49,3 +50,70 @@ pub fn retain_basic_blocks(mir: &mut Mir, keep: &[bool]) {

update_basic_block_ids(mir, &replacements);
}

// A simple map to perform quick lookups of the predecessors of a BasicBlock.
// Since BasicBlocks usually only have a small number of predecessors, we use a
// simple vector. Also, if a block has the same target more than once, for
// example in a switch, it will appear in the target's predecessor list multiple
// times. This allows to update the map more easily when modifying the graph.
pub struct PredecessorMap {
map: Vec<Vec<BasicBlock>>,
Copy link
Contributor

Choose a reason for hiding this comment

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

As far as I can tell, this pass only cares about how many predecessors a block has, not what those predecessors are. It would be simpler and more efficient to just store a Vec<usize>.

}

impl PredecessorMap {
pub fn new(num_blocks: usize) -> PredecessorMap {
PredecessorMap {
map: vec![Vec::new(); num_blocks],
}
}

pub fn predecessors(&self, block: BasicBlock) -> &[BasicBlock] {
&self.map[block.index()]
}

pub fn add_predecessor(&mut self, block: BasicBlock, predecessor: BasicBlock) {
self.map[block.index()].push(predecessor);
}

pub fn remove_predecessor(&mut self, block: BasicBlock, predecessor: BasicBlock) {
let pos = self.map[block.index()].iter().position(|&p| p == predecessor).expect(
&format!("{:?} is not registered as a predecessor of {:?}", predecessor, block));

self.map[block.index()].swap_remove(pos);
}

pub fn replace_predecessor(&mut self, block: BasicBlock, old: BasicBlock, new: BasicBlock) {
self.remove_predecessor(block, old);
self.add_predecessor(block, new);
}
}

struct PredecessorVisitor {
predecessor_map: PredecessorMap,
}

impl PredecessorVisitor {
fn new(num_blocks: usize) -> PredecessorVisitor {
PredecessorVisitor {
predecessor_map: PredecessorMap::new(num_blocks),
}
}
}

impl<'tcx> Visitor<'tcx> for PredecessorVisitor {
fn visit_mir(&mut self, mir: &Mir<'tcx>) {
self.predecessor_map = PredecessorMap::new(mir.basic_blocks.len());
self.super_mir(mir);
}

fn visit_branch(&mut self, source: BasicBlock, target: BasicBlock) {
self.predecessor_map.add_predecessor(target, source);
}
}

pub fn build_predecessor_map(mir: &Mir) -> PredecessorMap {
let mut v = PredecessorVisitor::new(mir.basic_blocks.len());
v.visit_mir(mir);

v.predecessor_map
}