Skip to content

Commit

Permalink
Rollup merge of #129720 - nnethercote:simplify-dest_prop-mm, r=cjgillot
Browse files Browse the repository at this point in the history
Simplify DestProp memory management

The DestProp MIR pass has some convoluted memory management. This PR simplifies it.

r? ```@davidtwco```
  • Loading branch information
matthiaskrgr authored Sep 5, 2024
2 parents 11085aa + 1be2204 commit 15e7a67
Showing 1 changed file with 46 additions and 67 deletions.
113 changes: 46 additions & 67 deletions compiler/rustc_mir_transform/src/dest_prop.rs
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,8 @@ impl<'tcx> crate::MirPass<'tcx> for DestinationPropagation {

fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
let def_id = body.source.def_id();
let mut allocations = Allocations::default();
let mut candidates = Candidates::default();
let mut write_info = WriteInfo::default();
trace!(func = ?tcx.def_path_str(def_id));

let borrowed = rustc_mir_dataflow::impls::borrowed_locals(body);
Expand Down Expand Up @@ -191,20 +192,15 @@ impl<'tcx> crate::MirPass<'tcx> for DestinationPropagation {
loop {
// PERF: Can we do something smarter than recalculating the candidates and liveness
// results?
let mut candidates = find_candidates(
body,
&borrowed,
&mut allocations.candidates,
&mut allocations.candidates_reverse,
);
candidates.reset_and_find(body, &borrowed);
trace!(?candidates);
dest_prop_mir_dump(tcx, body, &points, &live, round_count);

FilterInformation::filter_liveness(
&mut candidates,
&points,
&live,
&mut allocations.write_info,
&mut write_info,
body,
);

Expand Down Expand Up @@ -253,20 +249,8 @@ impl<'tcx> crate::MirPass<'tcx> for DestinationPropagation {
}
}

/// Container for the various allocations that we need.
///
/// We store these here and hand out `&mut` access to them, instead of dropping and recreating them
/// frequently. Everything with a `&'alloc` lifetime points into here.
#[derive(Default)]
struct Allocations {
candidates: FxIndexMap<Local, Vec<Local>>,
candidates_reverse: FxIndexMap<Local, Vec<Local>>,
write_info: WriteInfo,
// PERF: Do this for `MaybeLiveLocals` allocations too.
}

#[derive(Debug)]
struct Candidates<'alloc> {
#[derive(Debug, Default)]
struct Candidates {
/// The set of candidates we are considering in this optimization.
///
/// We will always merge the key into at most one of its values.
Expand All @@ -281,11 +265,12 @@ struct Candidates<'alloc> {
///
/// We will still report that we would like to merge `_1` and `_2` in an attempt to allow us to
/// remove that assignment.
c: &'alloc mut FxIndexMap<Local, Vec<Local>>,
c: FxIndexMap<Local, Vec<Local>>,

/// A reverse index of the `c` set; if the `c` set contains `a => Place { local: b, proj }`,
/// then this contains `b => a`.
// PERF: Possibly these should be `SmallVec`s?
reverse: &'alloc mut FxIndexMap<Local, Vec<Local>>,
reverse: FxIndexMap<Local, Vec<Local>>,
}

//////////////////////////////////////////////////////////
Expand Down Expand Up @@ -358,19 +343,40 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Merger<'a, 'tcx> {
//
// This section enforces bullet point 2

struct FilterInformation<'a, 'body, 'alloc, 'tcx> {
body: &'body Body<'tcx>,
struct FilterInformation<'a, 'tcx> {
body: &'a Body<'tcx>,
points: &'a DenseLocationMap,
live: &'a SparseIntervalMatrix<Local, PointIndex>,
candidates: &'a mut Candidates<'alloc>,
write_info: &'alloc mut WriteInfo,
candidates: &'a mut Candidates,
write_info: &'a mut WriteInfo,
at: Location,
}

// We first implement some utility functions which we will expose removing candidates according to
// different needs. Throughout the liveness filtering, the `candidates` are only ever accessed
// through these methods, and not directly.
impl<'alloc> Candidates<'alloc> {
impl Candidates {
/// Collects the candidates for merging.
///
/// This is responsible for enforcing the first and third bullet point.
fn reset_and_find<'tcx>(&mut self, body: &Body<'tcx>, borrowed: &BitSet<Local>) {
self.c.clear();
self.reverse.clear();
let mut visitor = FindAssignments { body, candidates: &mut self.c, borrowed };
visitor.visit_body(body);
// Deduplicate candidates.
for (_, cands) in self.c.iter_mut() {
cands.sort();
cands.dedup();
}
// Generate the reverse map.
for (src, cands) in self.c.iter() {
for dest in cands.iter().copied() {
self.reverse.entry(dest).or_default().push(*src);
}
}
}

/// Just `Vec::retain`, but the condition is inverted and we add debugging output
fn vec_filter_candidates(
src: Local,
Expand Down Expand Up @@ -445,7 +451,7 @@ enum CandidateFilter {
Remove,
}

impl<'a, 'body, 'alloc, 'tcx> FilterInformation<'a, 'body, 'alloc, 'tcx> {
impl<'a, 'tcx> FilterInformation<'a, 'tcx> {
/// Filters the set of candidates to remove those that conflict.
///
/// The steps we take are exactly those that are outlined at the top of the file. For each
Expand All @@ -463,12 +469,12 @@ impl<'a, 'body, 'alloc, 'tcx> FilterInformation<'a, 'body, 'alloc, 'tcx> {
/// before the statement/terminator will correctly report locals that are read in the
/// statement/terminator to be live. We are additionally conservative by treating all written to
/// locals as also being read from.
fn filter_liveness<'b>(
candidates: &mut Candidates<'alloc>,
fn filter_liveness(
candidates: &mut Candidates,
points: &DenseLocationMap,
live: &SparseIntervalMatrix<Local, PointIndex>,
write_info_alloc: &'alloc mut WriteInfo,
body: &'b Body<'tcx>,
write_info: &mut WriteInfo,
body: &Body<'tcx>,
) {
let mut this = FilterInformation {
body,
Expand All @@ -477,7 +483,7 @@ impl<'a, 'body, 'alloc, 'tcx> FilterInformation<'a, 'body, 'alloc, 'tcx> {
candidates,
// We don't actually store anything at this scope, we just keep things here to be able
// to reuse the allocation.
write_info: write_info_alloc,
write_info,
// Doesn't matter what we put here, will be overwritten before being used
at: Location::START,
};
Expand Down Expand Up @@ -734,40 +740,13 @@ fn places_to_candidate_pair<'tcx>(
Some((a, b))
}

/// Collects the candidates for merging
///
/// This is responsible for enforcing the first and third bullet point.
fn find_candidates<'alloc, 'tcx>(
body: &Body<'tcx>,
borrowed: &BitSet<Local>,
candidates: &'alloc mut FxIndexMap<Local, Vec<Local>>,
candidates_reverse: &'alloc mut FxIndexMap<Local, Vec<Local>>,
) -> Candidates<'alloc> {
candidates.clear();
candidates_reverse.clear();
let mut visitor = FindAssignments { body, candidates, borrowed };
visitor.visit_body(body);
// Deduplicate candidates
for (_, cands) in candidates.iter_mut() {
cands.sort();
cands.dedup();
}
// Generate the reverse map
for (src, cands) in candidates.iter() {
for dest in cands.iter().copied() {
candidates_reverse.entry(dest).or_default().push(*src);
}
}
Candidates { c: candidates, reverse: candidates_reverse }
}

struct FindAssignments<'a, 'alloc, 'tcx> {
struct FindAssignments<'a, 'tcx> {
body: &'a Body<'tcx>,
candidates: &'alloc mut FxIndexMap<Local, Vec<Local>>,
candidates: &'a mut FxIndexMap<Local, Vec<Local>>,
borrowed: &'a BitSet<Local>,
}

impl<'tcx> Visitor<'tcx> for FindAssignments<'_, '_, 'tcx> {
impl<'tcx> Visitor<'tcx> for FindAssignments<'_, 'tcx> {
fn visit_statement(&mut self, statement: &Statement<'tcx>, _: Location) {
if let StatementKind::Assign(box (
lhs,
Expand Down Expand Up @@ -819,9 +798,9 @@ fn is_local_required(local: Local, body: &Body<'_>) -> bool {
/////////////////////////////////////////////////////////
// MIR Dump

fn dest_prop_mir_dump<'body, 'tcx>(
fn dest_prop_mir_dump<'tcx>(
tcx: TyCtxt<'tcx>,
body: &'body Body<'tcx>,
body: &Body<'tcx>,
points: &DenseLocationMap,
live: &SparseIntervalMatrix<Local, PointIndex>,
round: usize,
Expand Down

0 comments on commit 15e7a67

Please sign in to comment.