Auto merge of #142531 - ohadravid:better-storage-calls-copy-prop, r=<try>

Remove fewer Storage calls in `copy_prop`

Modify the `copy_prop` MIR optimization pass to remove fewer `Storage{Live,Dead}` calls, allowing for better optimizations by LLVM - see #141649.

### Details

This is my attempt to fix the mentioned issue (this is the first part, I also implemented a similar solution for GVN in [this branch](https://github.com/rust-lang/rust/compare/master...ohadravid:rust:better-storage-calls-gvn-v2?expand=1)).

The idea is to use the `MaybeStorageDead` analysis and remove only the storage calls of `head`s that are maybe-storage-dead when the associated `local` is accessed (or, conversely, keep the storage of `head`s that are for-sure alive in _every_ relevant access).

When combined with the GVN change, the final example in the issue (#141649 (comment)) is optimized as expected by LLVM. I also measured the effect on a few functions in `rav1d` (where I originally saw the issue) and observed reduced stack usage in several of them.

This is my first attempt at working with MIR optimizations, so it's possible this isn't the right approach — but all tests pass, and the resulting diffs appear correct.

r? tmiasko

since he commented on the issue and pointed to these passes.

Author: bors

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -1,8 +1,12 @@
+use std::borrow::Cow;
+
 use rustc_index::IndexSlice;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
+use rustc_mir_dataflow::impls::{MaybeStorageDead, always_storage_live_locals};
+use rustc_mir_dataflow::{Analysis, ResultsCursor};
 use tracing::{debug, instrument};
 
 use crate::ssa::SsaLocals;
@@ -16,7 +20,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 `_a`, either copied or moved.
+/// 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 {
@@ -34,15 +38,40 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
         let fully_moved = fully_moved_locals(&ssa, body);
         debug!(?fully_moved);
 
-        let mut storage_to_remove = DenseBitSet::new_empty(fully_moved.domain_size());
+        let mut head_storage_to_check = DenseBitSet::new_empty(fully_moved.domain_size());
+
         for (local, &head) in ssa.copy_classes().iter_enumerated() {
             if local != head {
-                storage_to_remove.insert(head);
+                // We need to determine if we can keep the head's storage statements (which enables better optimizations).
+                // For every local's usage location, if the head is in `maybe_storage_dead`, we have to remove the storage statements for it.
+                head_storage_to_check.insert(head);
             }
         }
 
         let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);
 
+        let storage_to_remove = if any_replacement {
+            let always_live_locals = &always_storage_live_locals(body);
+
+            let maybe_storage_dead = MaybeStorageDead::new(Cow::Borrowed(always_live_locals))
+                .iterate_to_fixpoint(tcx, body, None)
+                .into_results_cursor(body);
+
+            let mut storage_checker = StorageChecker {
+                copy_classes: ssa.copy_classes(),
+                maybe_storage_dead,
+                head_storage_to_check,
+                storage_to_remove: DenseBitSet::new_empty(fully_moved.domain_size()),
+            };
+
+            storage_checker.visit_body(body);
+
+            storage_checker.storage_to_remove
+        } else {
+            // Will be empty anyway.
+            head_storage_to_check
+        };
+
         Replacer {
             tcx,
             copy_classes: ssa.copy_classes(),
@@ -119,6 +148,7 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
         if self.borrowed_locals.contains(*local) {
             return;
         }
+
         match ctxt {
             // Do not modify the local in storage statements.
             PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead) => {}
@@ -172,3 +202,36 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
         }
     }
 }
+
+struct StorageChecker<'a, 'tcx> {
+    storage_to_remove: DenseBitSet<Local>,
+    head_storage_to_check: DenseBitSet<Local>,
+    maybe_storage_dead: ResultsCursor<'a, 'tcx, MaybeStorageDead<'a>>,
+    copy_classes: &'a IndexSlice<Local, Local>,
+}
+
+impl<'a, 'tcx> Visitor<'tcx> for StorageChecker<'a, 'tcx> {
+    fn visit_local(&mut self, local: Local, context: PlaceContext, location: Location) {
+        if !context.is_use() {
+            return;
+        }
+
+        let head = self.copy_classes[local];
+        if self.head_storage_to_check.contains(head) {
+            self.maybe_storage_dead.seek_after_primary_effect(location);
+
+            if self.maybe_storage_dead.get().contains(head) {
+                debug!(
+                    ?location,
+                    ?local,
+                    ?head,
+                    "found use of local with head at a location in which head is maybe dead, marking head for storage removal"
+                );
+                self.storage_to_remove.insert(head);
+
+                // Once we found a use of the head that is maybe dead, we do not need to check it again.
+                self.head_storage_to_check.remove(head);
+            }
+        }
+    }
+}

tests/mir-opt/copy-prop/cycle.main.CopyProp.panic-abort.diff

@@ -26,7 +26,7 @@
       }
 
       bb1: {
--         StorageLive(_2);
+          StorageLive(_2);
           _2 = copy _1;
 -         StorageLive(_3);
 -         _3 = copy _2;
@@ -46,7 +46,7 @@
           StorageDead(_5);
           _0 = const ();
 -         StorageDead(_3);
--         StorageDead(_2);
+          StorageDead(_2);
           StorageDead(_1);
           return;
       }

tests/mir-opt/copy-prop/cycle.main.CopyProp.panic-unwind.diff

@@ -26,7 +26,7 @@
       }
 
       bb1: {
--         StorageLive(_2);
+          StorageLive(_2);
           _2 = copy _1;
 -         StorageLive(_3);
 -         _3 = copy _2;
@@ -46,7 +46,7 @@
           StorageDead(_5);
           _0 = const ();
 -         StorageDead(_3);
--         StorageDead(_2);
+          StorageDead(_2);
           StorageDead(_1);
           return;
       }

tests/mir-opt/copy-prop/dead_stores_79191.f.CopyProp.after.panic-abort.mir

@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
     }
 
     bb0: {
+        StorageLive(_2);
         _2 = copy _1;
         _1 = const 5_usize;
         _1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
 
     bb1: {
         StorageDead(_4);
+        StorageDead(_2);
         return;
     }
 }

tests/mir-opt/copy-prop/dead_stores_79191.f.CopyProp.after.panic-unwind.mir

@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
     }
 
     bb0: {
+        StorageLive(_2);
         _2 = copy _1;
         _1 = const 5_usize;
         _1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
 
     bb1: {
         StorageDead(_4);
+        StorageDead(_2);
         return;
     }
 }

tests/mir-opt/copy-prop/dead_stores_better.f.CopyProp.after.panic-abort.mir

@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
     }
 
     bb0: {
+        StorageLive(_2);
         _2 = copy _1;
         _1 = const 5_usize;
         _1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
 
     bb1: {
         StorageDead(_4);
+        StorageDead(_2);
         return;
     }
 }

tests/mir-opt/copy-prop/dead_stores_better.f.CopyProp.after.panic-unwind.mir

@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
     }
 
     bb0: {
+        StorageLive(_2);
         _2 = copy _1;
         _1 = const 5_usize;
         _1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
 
     bb1: {
         StorageDead(_4);
+        StorageDead(_2);
         return;
     }
 }

tests/mir-opt/copy-prop/issue_107511.main.CopyProp.panic-abort.diff

@@ -86,7 +86,7 @@
       }
 
       bb6: {
--         StorageLive(_16);
+          StorageLive(_16);
           _16 = copy ((_11 as Some).0: usize);
           StorageLive(_17);
 -         StorageLive(_18);
@@ -116,7 +116,7 @@
           StorageDead(_17);
 -         StorageDead(_18);
 -         _10 = const ();
--         StorageDead(_16);
+          StorageDead(_16);
           StorageDead(_13);
           StorageDead(_11);
 -         StorageDead(_10);

tests/mir-opt/copy-prop/issue_107511.main.CopyProp.panic-unwind.diff

@@ -86,7 +86,7 @@
       }
 
       bb6: {
--         StorageLive(_16);
+          StorageLive(_16);
           _16 = copy ((_11 as Some).0: usize);
           StorageLive(_17);
 -         StorageLive(_18);
@@ -116,7 +116,7 @@
           StorageDead(_17);
 -         StorageDead(_18);
 -         _10 = const ();
--         StorageDead(_16);
+          StorageDead(_16);
           StorageDead(_13);
           StorageDead(_11);
 -         StorageDead(_10);

tests/mir-opt/copy-prop/issue_141649.main.CopyProp.panic-abort.diff

@@ -0,0 +1,100 @@
+- // MIR for `main` before CopyProp
++ // MIR for `main` after CopyProp
+  
+  fn main() -> () {
+      let mut _0: ();
+      let _1: ();
+      let _2: main::S;
+      let _3: ();
+      let mut _4: main::S;
+      let _5: ();
+      let _6: main::S;
+      let _7: ();
+      let mut _8: main::S;
+      let _9: ();
+      let _10: main::C;
+      let _11: ();
+      let mut _12: main::C;
+      let _13: main::C;
+      let _14: ();
+      let mut _15: main::C;
+      scope 1 {
+          debug s1 => _2;
+      }
+      scope 2 {
+          debug s2 => _6;
+      }
+      scope 3 {
+          debug c1 => _10;
+      }
+      scope 4 {
+          debug c2 => _13;
+      }
+  
+      bb0: {
+-         StorageLive(_1);
+          StorageLive(_2);
+          _2 = S(const 1_usize, const 2_usize);
+          StorageLive(_3);
+-         StorageLive(_4);
+-         _4 = move _2;
+-         _3 = std::mem::drop::<S>(move _4) -> [return: bb1, unwind unreachable];
++         _3 = std::mem::drop::<S>(move _2) -> [return: bb1, unwind unreachable];
+      }
+  
+      bb1: {
+-         StorageDead(_4);
+          StorageDead(_3);
+-         _1 = const ();
+          StorageDead(_2);
+-         StorageDead(_1);
+-         StorageLive(_5);
+          StorageLive(_6);
+          _6 = S(const 3_usize, const 4_usize);
+          StorageLive(_7);
+-         StorageLive(_8);
+-         _8 = move _6;
+-         _7 = std::mem::drop::<S>(move _8) -> [return: bb2, unwind unreachable];
++         _7 = std::mem::drop::<S>(move _6) -> [return: bb2, unwind unreachable];
+      }
+  
+      bb2: {
+-         StorageDead(_8);
+          StorageDead(_7);
+-         _5 = const ();
+          StorageDead(_6);
+-         StorageDead(_5);
+-         StorageLive(_9);
+          StorageLive(_10);
+          _10 = C(const 1_usize, const 2_usize);
+          StorageLive(_11);
+-         StorageLive(_12);
+-         _12 = copy _10;
+-         _11 = std::mem::drop::<C>(move _12) -> [return: bb3, unwind unreachable];
++         _11 = std::mem::drop::<C>(copy _10) -> [return: bb3, unwind unreachable];
+      }
+  
+      bb3: {
+-         StorageDead(_12);
+          StorageDead(_11);
+-         _9 = const ();
+          StorageDead(_10);
+-         StorageDead(_9);
+          StorageLive(_13);
+          _13 = C(const 3_usize, const 4_usize);
+          StorageLive(_14);
+-         StorageLive(_15);
+-         _15 = copy _13;
+-         _14 = std::mem::drop::<C>(move _15) -> [return: bb4, unwind unreachable];
++         _14 = std::mem::drop::<C>(copy _13) -> [return: bb4, unwind unreachable];
+      }
+  
+      bb4: {
+-         StorageDead(_15);
+          StorageDead(_14);
+          _0 = const ();
+          StorageDead(_13);
+          return;
+      }
+  }
+