Improve copy_prop and GVN mir-opt passes to remove fewer storage calls

Author: ohadravid

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -3,9 +3,10 @@ use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
+use rustc_mir_dataflow::{Analysis, ResultsCursor};
 use tracing::{debug, instrument};
 
-use crate::ssa::SsaLocals;
+use crate::ssa::{MaybeUninitializedLocals, SsaLocals};
 
 /// Unify locals that copy each other.
 ///
@@ -16,7 +17,7 @@ use crate::ssa::SsaLocals;
 ///   _d = move? _c
 /// where each of the locals is only assigned once.
 ///
-/// We want to replace all those locals by `copy _a`.
+/// We want to replace all those locals by `_a` (the "head"), either copied or moved.
 pub(super) struct CopyProp;
 
 impl<'tcx> crate::MirPass<'tcx> for CopyProp {
@@ -30,15 +31,19 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
 
         let typing_env = body.typing_env(tcx);
         let ssa = SsaLocals::new(tcx, body, typing_env);
+
         debug!(borrowed_locals = ?ssa.borrowed_locals());
         debug!(copy_classes = ?ssa.copy_classes());
 
         let mut any_replacement = false;
         // Locals that participate in copy propagation either as a source or a destination.
         let mut unified = DenseBitSet::new_empty(body.local_decls.len());
+        let mut storage_to_remove = DenseBitSet::new_empty(body.local_decls.len());
+
         for (local, &head) in ssa.copy_classes().iter_enumerated() {
             if local != head {
                 any_replacement = true;
+                storage_to_remove.insert(head);
                 unified.insert(head);
                 unified.insert(local);
             }
@@ -48,7 +53,46 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
             return;
         }
 
-        Replacer { tcx, copy_classes: ssa.copy_classes(), unified }.visit_body_preserves_cfg(body);
+        // When emitting storage statements, we want to retain the head locals' storage statements,
+        // as this enables better optimizations. For each local use location, we mark the head for storage removal
+        // only if the head might be uninitialized at that point, or if the local is borrowed
+        // (since we cannot easily determine when it's used).
+        let storage_to_remove = if tcx.sess.emit_lifetime_markers() {
+            storage_to_remove.clear();
+
+            // If the local is borrowed, we cannot easily determine if it is used, so we have to remove the storage statements.
+            let borrowed_locals = ssa.borrowed_locals();
+
+            for (local, &head) in ssa.copy_classes().iter_enumerated() {
+                if local != head && borrowed_locals.contains(local) {
+                    storage_to_remove.insert(head);
+                }
+            }
+
+            let maybe_uninit = MaybeUninitializedLocals
+                .iterate_to_fixpoint(tcx, body, Some("mir_opt::copy_prop"))
+                .into_results_cursor(body);
+
+            let mut storage_checker = StorageChecker {
+                maybe_uninit,
+                copy_classes: ssa.copy_classes(),
+                storage_to_remove,
+            };
+
+            for (bb, data) in traversal::reachable(body) {
+                storage_checker.visit_basic_block_data(bb, data);
+            }
+
+            storage_checker.storage_to_remove
+        } else {
+            // Remove the storage statements of all the head locals.
+            storage_to_remove
+        };
+
+        debug!(?storage_to_remove);
+
+        Replacer { tcx, copy_classes: ssa.copy_classes(), unified, storage_to_remove }
+            .visit_body_preserves_cfg(body);
 
         crate::simplify::remove_unused_definitions(body);
     }
@@ -63,6 +107,7 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
 struct Replacer<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     unified: DenseBitSet<Local>,
+    storage_to_remove: DenseBitSet<Local>,
     copy_classes: &'a IndexSlice<Local, Local>,
 }
 
@@ -73,7 +118,13 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
 
     #[tracing::instrument(level = "trace", skip(self))]
     fn visit_local(&mut self, local: &mut Local, ctxt: PlaceContext, _: Location) {
-        *local = self.copy_classes[*local];
+        let new_local = self.copy_classes[*local];
+        match ctxt {
+            // Do not modify the local in storage statements.
+            PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead) => {}
+            // We access the value.
+            _ => *local = new_local,
+        }
     }
 
     #[tracing::instrument(level = "trace", skip(self))]
@@ -93,7 +144,7 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
     fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
         // When removing storage statements, we need to remove both (#107511).
         if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = stmt.kind
-            && self.unified.contains(l)
+            && self.storage_to_remove.contains(l)
         {
             stmt.make_nop(true);
         }
@@ -109,3 +160,39 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
         }
     }
 }
+
+// Marks heads of copy classes that are maybe uninitialized at the location of a local
+// as needing storage statement removal.
+struct StorageChecker<'a, 'tcx> {
+    maybe_uninit: ResultsCursor<'a, 'tcx, MaybeUninitializedLocals>,
+    copy_classes: &'a IndexSlice<Local, Local>,
+    storage_to_remove: DenseBitSet<Local>,
+}
+
+impl<'a, 'tcx> Visitor<'tcx> for StorageChecker<'a, 'tcx> {
+    fn visit_local(&mut self, local: Local, context: PlaceContext, loc: Location) {
+        if !context.is_use() {
+            return;
+        }
+
+        let head = self.copy_classes[local];
+
+        // If the local is the head, or if we already marked it for deletion, we do not need to check it.
+        if head == local || self.storage_to_remove.contains(head) {
+            return;
+        }
+
+        self.maybe_uninit.seek_before_primary_effect(loc);
+
+        if self.maybe_uninit.get().contains(head) {
+            debug!(
+                ?loc,
+                ?context,
+                ?local,
+                ?head,
+                "local's head is maybe uninit at this location, marking head for storage statement removal"
+            );
+            self.storage_to_remove.insert(head);
+        }
+    }
+}

compiler/rustc_mir_transform/src/gvn.rs

@@ -117,11 +117,12 @@ use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::layout::HasTypingEnv;
 use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_mir_dataflow::{Analysis, ResultsCursor};
 use rustc_span::DUMMY_SP;
 use smallvec::SmallVec;
 use tracing::{debug, instrument, trace};
 
-use crate::ssa::SsaLocals;
+use crate::ssa::{MaybeUninitializedLocals, SsaLocals};
 
 pub(super) struct GVN;
 
@@ -154,10 +155,34 @@ impl<'tcx> crate::MirPass<'tcx> for GVN {
             state.visit_basic_block_data(bb, data);
         }
 
-        // For each local that is reused (`y` above), we remove its storage statements do avoid any
-        // difficulty. Those locals are SSA, so should be easy to optimize by LLVM without storage
-        // statements.
-        StorageRemover { tcx, reused_locals: state.reused_locals }.visit_body_preserves_cfg(body);
+        // When emitting storage statements, we want to retain the reused locals' storage statements,
+        // as this enables better optimizations. For each local use location, we mark it for storage removal
+        // only if it might be uninitialized at that point.
+        let storage_to_remove = if tcx.sess.emit_lifetime_markers() {
+            let maybe_uninit = MaybeUninitializedLocals
+                .iterate_to_fixpoint(tcx, body, Some("mir_opt::gvn"))
+                .into_results_cursor(body);
+
+            let mut storage_checker = StorageChecker {
+                reused_locals: &state.reused_locals,
+                storage_to_remove: DenseBitSet::new_empty(body.local_decls.len()),
+                maybe_uninit,
+            };
+
+            for (bb, data) in traversal::reachable(body) {
+                storage_checker.visit_basic_block_data(bb, data);
+            }
+
+            storage_checker.storage_to_remove
+        } else {
+            // Remove the storage statements of all the reused locals.
+            state.reused_locals.clone()
+        };
+
+        debug!(?storage_to_remove);
+
+        StorageRemover { tcx, reused_locals: state.reused_locals, storage_to_remove }
+            .visit_body_preserves_cfg(body);
     }
 
     fn is_required(&self) -> bool {
@@ -2033,6 +2058,7 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
 struct StorageRemover<'tcx> {
     tcx: TyCtxt<'tcx>,
     reused_locals: DenseBitSet<Local>,
+    storage_to_remove: DenseBitSet<Local>,
 }
 
 impl<'tcx> MutVisitor<'tcx> for StorageRemover<'tcx> {
@@ -2053,11 +2079,53 @@ impl<'tcx> MutVisitor<'tcx> for StorageRemover<'tcx> {
         match stmt.kind {
             // When removing storage statements, we need to remove both (#107511).
             StatementKind::StorageLive(l) | StatementKind::StorageDead(l)
-                if self.reused_locals.contains(l) =>
+                if self.storage_to_remove.contains(l) =>
             {
                 stmt.make_nop(true)
             }
             _ => self.super_statement(stmt, loc),
         }
     }
 }
+
+struct StorageChecker<'a, 'tcx> {
+    reused_locals: &'a DenseBitSet<Local>,
+    storage_to_remove: DenseBitSet<Local>,
+    maybe_uninit: ResultsCursor<'a, 'tcx, MaybeUninitializedLocals>,
+}
+
+impl<'a, 'tcx> Visitor<'tcx> for StorageChecker<'a, 'tcx> {
+    fn visit_local(&mut self, local: Local, context: PlaceContext, location: Location) {
+        match context {
+            // These mutating uses do not require the local to be initialized,
+            // so we cannot use our maybe-uninit check on them.
+            // However, GVN doesn't introduce or move mutations,
+            // so this local must already have valid storage at this location.
+            PlaceContext::MutatingUse(MutatingUseContext::AsmOutput)
+            | PlaceContext::MutatingUse(MutatingUseContext::Call)
+            | PlaceContext::MutatingUse(MutatingUseContext::Store)
+            | PlaceContext::MutatingUse(MutatingUseContext::Yield)
+            | PlaceContext::NonUse(_) => {
+                return;
+            }
+            // Must check validity for other mutating usages and all non-mutating uses.
+            PlaceContext::MutatingUse(_) | PlaceContext::NonMutatingUse(_) => {}
+        }
+
+        // We only need to check reused locals which we haven't already removed storage for.
+        if !self.reused_locals.contains(local) || self.storage_to_remove.contains(local) {
+            return;
+        }
+
+        self.maybe_uninit.seek_before_primary_effect(location);
+
+        if self.maybe_uninit.get().contains(local) {
+            debug!(
+                ?location,
+                ?local,
+                "local is reused and is maybe uninit at this location, marking it for storage statement removal"
+            );
+            self.storage_to_remove.insert(local);
+        }
+    }
+}

compiler/rustc_mir_transform/src/ssa.rs

@@ -14,6 +14,7 @@ use rustc_middle::middle::resolve_bound_vars::Set1;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, TyCtxt};
+use rustc_mir_dataflow::Analysis;
 use tracing::{debug, instrument, trace};
 
 pub(super) struct SsaLocals {
@@ -391,3 +392,64 @@ impl StorageLiveLocals {
         matches!(self.storage_live[local], Set1::One(_))
     }
 }
+
+/// A dataflow analysis that tracks locals that are maybe uninitialized.
+///
+/// This is a simpler analysis than `MaybeUninitializedPlaces`, because it does not track
+/// individual fields.
+pub(crate) struct MaybeUninitializedLocals;
+
+impl<'tcx> Analysis<'tcx> for MaybeUninitializedLocals {
+    type Domain = DenseBitSet<Local>;
+
+    const NAME: &'static str = "maybe_uninit_locals";
+
+    fn bottom_value(&self, body: &Body<'tcx>) -> Self::Domain {
+        // bottom = all locals are initialized.
+        DenseBitSet::new_empty(body.local_decls.len())
+    }
+
+    fn initialize_start_block(&self, body: &Body<'tcx>, state: &mut Self::Domain) {
+        // All locals start as uninitialized...
+        state.insert_all();
+        // ...except for arguments, which are definitely initialized.
+        for arg in body.args_iter() {
+            state.remove(arg);
+        }
+    }
+
+    fn apply_primary_statement_effect(
+        &self,
+        state: &mut Self::Domain,
+        statement: &Statement<'tcx>,
+        _location: Location,
+    ) {
+        match statement.kind {
+            // An assignment makes a local initialized.
+            StatementKind::Assign(box (place, _)) => {
+                if let Some(local) = place.as_local() {
+                    state.remove(local);
+                }
+            }
+            // Storage{Live,Dead} makes a local uninitialized.
+            StatementKind::StorageLive(local) | StatementKind::StorageDead(local) => {
+                state.insert(local);
+            }
+            _ => {}
+        }
+    }
+
+    fn apply_call_return_effect(
+        &self,
+        state: &mut Self::Domain,
+        _block: BasicBlock,
+        return_places: CallReturnPlaces<'_, 'tcx>,
+    ) {
+        // The return place of a call is initialized.
+        return_places.for_each(|place| {
+            if let Some(local) = place.as_local() {
+                state.remove(local);
+            }
+        });
+    }
+}