bjoernager/rust
bjoernager/rust
1
Fork 0
Code Activity

Auto merge of #106908 - cjgillot:copyprop-ssa, r=oli-obk

Browse source 
Implement simple CopyPropagation based on SSA analysis

This PR extracts the "copy propagation" logic from https://github.com/rust-lang/rust/pull/106285.

MIR may produce chains of assignment between locals, like `_x = move? _y`.
This PR attempts to remove such chains by unifying locals.

The current implementation is a bit overzealous in turning moves into copies, and in removing storage statements.
This commit is contained in:
bors 2023-01-29 13:01:06 +00:00
commit 2a4b00beaa
68 changed files with 1948 additions and 278 deletions
Download patch file Download diff file Expand all files Collapse all files

178
compiler/rustc_mir_transform/src/copy_prop.rs Normal file
View file

@ -0,0 +1,178 @@
use rustc_index::bit_set::BitSet;
use rustc_index::vec::IndexVec;
use rustc_middle::mir::visit::*;
use rustc_middle::mir::*;
use rustc_middle::ty::TyCtxt;
use rustc_mir_dataflow::impls::borrowed_locals;
use crate::ssa::SsaLocals;
use crate::MirPass;
/// Unify locals that copy each other.
///
/// We consider patterns of the form
/// _a = rvalue
/// _b = move? _a
/// _c = move? _a
/// _d = move? _c
/// where each of the locals is only assigned once.
///
/// We want to replace all those locals by `_a`, either copied or moved.
pub struct CopyProp;
impl<'tcx> MirPass<'tcx> for CopyProp {
fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
sess.mir_opt_level() >= 4
}
#[instrument(level = "trace", skip(self, tcx, body))]
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
debug!(def_id = ?body.source.def_id());
propagate_ssa(tcx, body);
}
}
fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
let borrowed_locals = borrowed_locals(body);
let ssa = SsaLocals::new(tcx, param_env, body, &borrowed_locals);
let fully_moved = fully_moved_locals(&ssa, body);
debug!(?fully_moved);
let mut storage_to_remove = BitSet::new_empty(fully_moved.domain_size());
for (local, &head) in ssa.copy_classes().iter_enumerated() {
if local != head {
storage_to_remove.insert(head);
}
}
let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);
Replacer {
tcx,
copy_classes: &ssa.copy_classes(),
fully_moved,
borrowed_locals,
storage_to_remove,
}
.visit_body_preserves_cfg(body);
if any_replacement {
crate::simplify::remove_unused_definitions(body);
}
}
/// `SsaLocals` computed equivalence classes between locals considering copy/move assignments.
///
/// This function also returns whether all the `move?` in the pattern are `move` and not copies.
/// A local which is in the bitset can be replaced by `move _a`. Otherwise, it must be
/// replaced by `copy _a`, as we cannot move multiple times from `_a`.
///
/// If an operand copies `_c`, it must happen before the assignment `_d = _c`, otherwise it is UB.
/// This means that replacing it by a copy of `_a` if ok, since this copy happens before `_c` is
/// moved, and therefore that `_d` is moved.
#[instrument(level = "trace", skip(ssa, body))]
fn fully_moved_locals(ssa: &SsaLocals, body: &Body<'_>) -> BitSet<Local> {
let mut fully_moved = BitSet::new_filled(body.local_decls.len());
for (_, rvalue) in ssa.assignments(body) {
let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) | Rvalue::CopyForDeref(place))
= rvalue
else { continue };
let Some(rhs) = place.as_local() else { continue };
if !ssa.is_ssa(rhs) {
continue;
}
if let Rvalue::Use(Operand::Copy(_)) | Rvalue::CopyForDeref(_) = rvalue {
fully_moved.remove(rhs);
}
}
ssa.meet_copy_equivalence(&mut fully_moved);
fully_moved
}
/// Utility to help performing subtitution of `*pattern` by `target`.
struct Replacer<'a, 'tcx> {
tcx: TyCtxt<'tcx>,
fully_moved: BitSet<Local>,
storage_to_remove: BitSet<Local>,
borrowed_locals: BitSet<Local>,
copy_classes: &'a IndexVec<Local, Local>,
}
impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
fn tcx(&self) -> TyCtxt<'tcx> {
self.tcx
}
fn visit_local(&mut self, local: &mut Local, ctxt: PlaceContext, _: Location) {
let new_local = self.copy_classes[*local];
match ctxt {
// Do not modify the local in storage statements.
PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead) => {}
// The local should have been marked as non-SSA.
PlaceContext::MutatingUse(_) => assert_eq!(*local, new_local),
// We access the value.
_ => *local = new_local,
}
}
fn visit_place(&mut self, place: &mut Place<'tcx>, ctxt: PlaceContext, loc: Location) {
if let Some(new_projection) = self.process_projection(&place.projection, loc) {
place.projection = self.tcx().intern_place_elems(&new_projection);
}
let observes_address = match ctxt {
PlaceContext::NonMutatingUse(
NonMutatingUseContext::SharedBorrow
| NonMutatingUseContext::ShallowBorrow
| NonMutatingUseContext::UniqueBorrow
| NonMutatingUseContext::AddressOf,
) => true,
// For debuginfo, merging locals is ok.
PlaceContext::NonUse(NonUseContext::VarDebugInfo) => {
self.borrowed_locals.contains(place.local)
}
_ => false,
};
if observes_address && !place.is_indirect() {
// We observe the address of `place.local`. Do not replace it.
} else {
self.visit_local(
&mut place.local,
PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy),
loc,
)
}
}
fn visit_operand(&mut self, operand: &mut Operand<'tcx>, loc: Location) {
if let Operand::Move(place) = *operand
&& let Some(local) = place.as_local()
&& !self.fully_moved.contains(local)
{
*operand = Operand::Copy(place);
}
self.super_operand(operand, loc);
}
fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
if let StatementKind::StorageDead(l) = stmt.kind
&& self.storage_to_remove.contains(l)
{
stmt.make_nop();
} else if let StatementKind::Assign(box (ref place, ref mut rvalue)) = stmt.kind
&& place.as_local().is_some()
{
// Do not replace assignments.
self.visit_rvalue(rvalue, loc)
} else {
self.super_statement(stmt, loc);
}
}
}

5
compiler/rustc_mir_transform/src/dest_prop.rs
View file

@ -328,7 +328,8 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Merger<'a, 'tcx> {
match &statement.kind {
StatementKind::Assign(box (dest, rvalue)) => {
match rvalue {
Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) => {
Rvalue::CopyForDeref(place)
| Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) => {
// These might've been turned into self-assignments by the replacement
// (this includes the original statement we wanted to eliminate).
if dest == place {
@ -755,7 +756,7 @@ impl<'tcx> Visitor<'tcx> for FindAssignments<'_, '_, 'tcx> {
fn visit_statement(&mut self, statement: &Statement<'tcx>, _: Location) {
if let StatementKind::Assign(box (
lhs,
Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)),
Rvalue::CopyForDeref(rhs) | Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)),
)) = &statement.kind
{
let Some((src, dest)) = places_to_candidate_pair(*lhs, *rhs, self.body) else {

3
compiler/rustc_mir_transform/src/lib.rs
View file

@ -54,6 +54,7 @@ mod const_debuginfo;
mod const_goto;
mod const_prop;
mod const_prop_lint;
mod copy_prop;
mod coverage;
mod ctfe_limit;
mod dataflow_const_prop;
@ -87,6 +88,7 @@ mod required_consts;
mod reveal_all;
mod separate_const_switch;
mod shim;
mod ssa;
// This pass is public to allow external drivers to perform MIR cleanup
pub mod simplify;
mod simplify_branches;
@ -563,6 +565,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
&instcombine::InstCombine,
&separate_const_switch::SeparateConstSwitch,
&simplify::SimplifyLocals::new("before-const-prop"),
&copy_prop::CopyProp,
//
// FIXME(#70073): This pass is responsible for both optimization as well as some lints.
&const_prop::ConstProp,

16
compiler/rustc_mir_transform/src/simplify.rs
View file

@ -404,6 +404,18 @@ impl<'tcx> MirPass<'tcx> for SimplifyLocals {
}
}
pub fn remove_unused_definitions<'tcx>(body: &mut Body<'tcx>) {
// First, we're going to get a count of *actual* uses for every `Local`.
let mut used_locals = UsedLocals::new(body);
// Next, we're going to remove any `Local` with zero actual uses. When we remove those
// `Locals`, we're also going to subtract any uses of other `Locals` from the `used_locals`
// count. For example, if we removed `_2 = discriminant(_1)`, then we'll subtract one from
// `use_counts[_1]`. That in turn might make `_1` unused, so we loop until we hit a
// fixedpoint where there are no more unused locals.
remove_unused_definitions_helper(&mut used_locals, body);
}
pub fn simplify_locals<'tcx>(body: &mut Body<'tcx>, tcx: TyCtxt<'tcx>) {
// First, we're going to get a count of *actual* uses for every `Local`.
let mut used_locals = UsedLocals::new(body);
@ -413,7 +425,7 @@ pub fn simplify_locals<'tcx>(body: &mut Body<'tcx>, tcx: TyCtxt<'tcx>) {
// count. For example, if we removed `_2 = discriminant(_1)`, then we'll subtract one from
// `use_counts[_1]`. That in turn might make `_1` unused, so we loop until we hit a
// fixedpoint where there are no more unused locals.
remove_unused_definitions(&mut used_locals, body);
remove_unused_definitions_helper(&mut used_locals, body);
// Finally, we'll actually do the work of shrinking `body.local_decls` and remapping the `Local`s.
let map = make_local_map(&mut body.local_decls, &used_locals);
@ -548,7 +560,7 @@ impl<'tcx> Visitor<'tcx> for UsedLocals {
}
/// Removes unused definitions. Updates the used locals to reflect the changes made.
fn remove_unused_definitions(used_locals: &mut UsedLocals, body: &mut Body<'_>) {
fn remove_unused_definitions_helper(used_locals: &mut UsedLocals, body: &mut Body<'_>) {
// The use counts are updated as we remove the statements. A local might become unused
// during the retain operation, leading to a temporary inconsistency (storage statements or
// definitions referencing the local might remain). For correctness it is crucial that this

222
compiler/rustc_mir_transform/src/ssa.rs Normal file
View file

@ -0,0 +1,222 @@
use either::Either;
use rustc_data_structures::graph::dominators::Dominators;
use rustc_index::bit_set::BitSet;
use rustc_index::vec::IndexVec;
use rustc_middle::middle::resolve_lifetime::Set1;
use rustc_middle::mir::visit::*;
use rustc_middle::mir::*;
use rustc_middle::ty::{ParamEnv, TyCtxt};
#[derive(Debug)]
pub struct SsaLocals {
/// Assignments to each local. This defines whether the local is SSA.
assignments: IndexVec<Local, Set1<LocationExtended>>,
/// We visit the body in reverse postorder, to ensure each local is assigned before it is used.
/// We remember the order in which we saw the assignments to compute the SSA values in a single
/// pass.
assignment_order: Vec<Local>,
/// Copy equivalence classes between locals. See `copy_classes` for documentation.
copy_classes: IndexVec<Local, Local>,
}
impl SsaLocals {
pub fn new<'tcx>(
tcx: TyCtxt<'tcx>,
param_env: ParamEnv<'tcx>,
body: &Body<'tcx>,
borrowed_locals: &BitSet<Local>,
) -> SsaLocals {
let assignment_order = Vec::new();
let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls);
let dominators = body.basic_blocks.dominators();
let mut visitor = SsaVisitor { assignments, assignment_order, dominators };
for (local, decl) in body.local_decls.iter_enumerated() {
if matches!(body.local_kind(local), LocalKind::Arg) {
visitor.assignments[local] = Set1::One(LocationExtended::Arg);
}
if borrowed_locals.contains(local) && !decl.ty.is_freeze(tcx, param_env) {
visitor.assignments[local] = Set1::Many;
}
}
for (bb, data) in traversal::reverse_postorder(body) {
visitor.visit_basic_block_data(bb, data);
}
for var_debug_info in &body.var_debug_info {
visitor.visit_var_debug_info(var_debug_info);
}
debug!(?visitor.assignments);
visitor
.assignment_order
.retain(|&local| matches!(visitor.assignments[local], Set1::One(_)));
debug!(?visitor.assignment_order);
let copy_classes = compute_copy_classes(&visitor, body);
SsaLocals {
assignments: visitor.assignments,
assignment_order: visitor.assignment_order,
copy_classes,
}
}
pub fn is_ssa(&self, local: Local) -> bool {
matches!(self.assignments[local], Set1::One(_))
}
pub fn assignments<'a, 'tcx>(
&'a self,
body: &'a Body<'tcx>,
) -> impl Iterator<Item = (Local, &'a Rvalue<'tcx>)> + 'a {
self.assignment_order.iter().filter_map(|&local| {
if let Set1::One(LocationExtended::Plain(loc)) = self.assignments[local] {
// `loc` must point to a direct assignment to `local`.
let Either::Left(stmt) = body.stmt_at(loc) else { bug!() };
let Some((target, rvalue)) = stmt.kind.as_assign() else { bug!() };
assert_eq!(target.as_local(), Some(local));
Some((local, rvalue))
} else {
None
}
})
}
/// Compute the equivalence classes for locals, based on copy statements.
///
/// The returned vector maps each local to the one it copies. In the following case:
/// _a = &mut _0
/// _b = move? _a
/// _c = move? _a
/// _d = move? _c
/// We return the mapping
/// _a => _a // not a copy so, represented by itself
/// _b => _a
/// _c => _a
/// _d => _a // transitively through _c
///
/// Exception: we do not see through the return place, as it cannot be substituted.
pub fn copy_classes(&self) -> &IndexVec<Local, Local> {
&self.copy_classes
}
/// Make a property uniform on a copy equivalence class by removing elements.
pub fn meet_copy_equivalence(&self, property: &mut BitSet<Local>) {
// Consolidate to have a local iff all its copies are.
//
// `copy_classes` defines equivalence classes between locals. The `local`s that recursively
// move/copy the same local all have the same `head`.
for (local, &head) in self.copy_classes.iter_enumerated() {
// If any copy does not have `property`, then the head is not.
if !property.contains(local) {
property.remove(head);
}
}
for (local, &head) in self.copy_classes.iter_enumerated() {
// If any copy does not have `property`, then the head doesn't either,
// then no copy has `property`.
if !property.contains(head) {
property.remove(local);
}
}
// Verify that we correctly computed equivalence classes.
#[cfg(debug_assertions)]
for (local, &head) in self.copy_classes.iter_enumerated() {
assert_eq!(property.contains(local), property.contains(head));
}
}
}
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum LocationExtended {
Plain(Location),
Arg,
}
struct SsaVisitor {
dominators: Dominators<BasicBlock>,
assignments: IndexVec<Local, Set1<LocationExtended>>,
assignment_order: Vec<Local>,
}
impl<'tcx> Visitor<'tcx> for SsaVisitor {
fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) {
match ctxt {
PlaceContext::MutatingUse(MutatingUseContext::Store) => {
self.assignments[local].insert(LocationExtended::Plain(loc));
self.assignment_order.push(local);
}
// Anything can happen with raw pointers, so remove them.
PlaceContext::NonMutatingUse(NonMutatingUseContext::AddressOf)
| PlaceContext::MutatingUse(_) => self.assignments[local] = Set1::Many,
// Immutable borrows are taken into account in `SsaLocals::new` by
// removing non-freeze locals.
PlaceContext::NonMutatingUse(_) => {
let set = &mut self.assignments[local];
let assign_dominates = match *set {
Set1::Empty | Set1::Many => false,
Set1::One(LocationExtended::Arg) => true,
Set1::One(LocationExtended::Plain(assign)) => {
assign.successor_within_block().dominates(loc, &self.dominators)
}
};
// We are visiting a use that is not dominated by an assignment.
// Either there is a cycle involved, or we are reading for uninitialized local.
// Bail out.
if !assign_dominates {
*set = Set1::Many;
}
}
PlaceContext::NonUse(_) => {}
}
}
}
#[instrument(level = "trace", skip(ssa, body))]
fn compute_copy_classes(ssa: &SsaVisitor, body: &Body<'_>) -> IndexVec<Local, Local> {
let mut copies = IndexVec::from_fn_n(|l| l, body.local_decls.len());
for &local in &ssa.assignment_order {
debug!(?local);
if local == RETURN_PLACE {
// `_0` is special, we cannot rename it.
continue;
}
// This is not SSA: mark that we don't know the value.
debug!(assignments = ?ssa.assignments[local]);
let Set1::One(LocationExtended::Plain(loc)) = ssa.assignments[local] else { continue };
// `loc` must point to a direct assignment to `local`.
let Either::Left(stmt) = body.stmt_at(loc) else { bug!() };
let Some((_target, rvalue)) = stmt.kind.as_assign() else { bug!() };
assert_eq!(_target.as_local(), Some(local));
let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) | Rvalue::CopyForDeref(place))
= rvalue
else { continue };
let Some(rhs) = place.as_local() else { continue };
let Set1::One(_) = ssa.assignments[rhs] else { continue };
// We visit in `assignment_order`, ie. reverse post-order, so `rhs` has been
// visited before `local`, and we just have to copy the representing local.
copies[local] = copies[rhs];
}
debug!(?copies);
// Invariant: `copies` must point to the head of an equivalence class.
#[cfg(debug_assertions)]
for &head in copies.iter() {
assert_eq!(copies[head], head);
}
copies
}