Auto merge of #111556 - cjgillot:copy-prop-nrvo, r=oli-obk

Merge return place with other locals in CopyProp. This reintroduces a limited form of NRVO. r? wg-mir-opt
rust-lang · May 16, 2023 · 5c3a336 · 5c3a336
2 parents e77366b + adfffc7
commit 5c3a336
Show file tree

Hide file tree

Showing 19 changed files with 617 additions and 240 deletions.
diff --git a/compiler/rustc_mir_transform/src/copy_prop.rs b/compiler/rustc_mir_transform/src/copy_prop.rs
@@ -162,20 +162,22 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
     }
 
     fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
-        match stmt.kind {
-            // When removing storage statements, we need to remove both (#107511).
-            StatementKind::StorageLive(l) | StatementKind::StorageDead(l)
-                if self.storage_to_remove.contains(l) =>
-            {
-                stmt.make_nop()
-            }
-            StatementKind::Assign(box (ref place, ref mut rvalue))
-                if place.as_local().is_some() =>
-            {
-                // Do not replace assignments.
-                self.visit_rvalue(rvalue, loc)
-            }
-            _ => self.super_statement(stmt, loc),
+        // When removing storage statements, we need to remove both (#107511).
+        if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = stmt.kind
+            && self.storage_to_remove.contains(l)
+        {
+            stmt.make_nop();
+            return
+        }
+
+        self.super_statement(stmt, loc);
+
+        // Do not leave tautological assignments around.
+        if let StatementKind::Assign(box (lhs, ref rhs)) = stmt.kind
+            && let Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)) | Rvalue::CopyForDeref(rhs) = *rhs
+            && lhs == rhs
+        {
+            stmt.make_nop();
         }
     }
 }
diff --git a/compiler/rustc_mir_transform/src/ssa.rs b/compiler/rustc_mir_transform/src/ssa.rs
@@ -101,14 +101,15 @@ impl SsaLocals {
             .retain(|&local| matches!(visitor.assignments[local], Set1::One(_)));
         debug!(?visitor.assignment_order);
 
-        let copy_classes = compute_copy_classes(&mut visitor, body);
-
-        SsaLocals {
+        let mut ssa = SsaLocals {
             assignments: visitor.assignments,
             assignment_order: visitor.assignment_order,
             direct_uses: visitor.direct_uses,
-            copy_classes,
-        }
+            // This is filled by `compute_copy_classes`.
+            copy_classes: IndexVec::default(),
+        };
+        compute_copy_classes(&mut ssa, body);
+        ssa
     }
 
     pub fn num_locals(&self) -> usize {
@@ -261,49 +262,54 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor {
 }
 
 #[instrument(level = "trace", skip(ssa, body))]
-fn compute_copy_classes(ssa: &mut SsaVisitor, body: &Body<'_>) -> IndexVec<Local, Local> {
+fn compute_copy_classes(ssa: &mut SsaLocals, body: &Body<'_>) {
+    let mut direct_uses = std::mem::take(&mut ssa.direct_uses);
     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));
-
+    for (local, 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 };
-        let Set1::One(_) = ssa.assignments[rhs] else { continue };
+        if !ssa.is_ssa(rhs) {
+            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];
-        ssa.direct_uses[rhs] -= 1;
+        let head = copies[rhs];
+
+        if local == RETURN_PLACE {
+            // `_0` is special, we cannot rename it. Instead, rename the class of `rhs` to
+            // `RETURN_PLACE`. This is only possible if the class head is a temporary, not an
+            // argument.
+            if body.local_kind(head) != LocalKind::Temp {
+                continue;
+            }
+            for h in copies.iter_mut() {
+                if *h == head {
+                    *h = RETURN_PLACE;
+                }
+            }
+        } else {
+            copies[local] = head;
+        }
+        direct_uses[rhs] -= 1;
     }
 
     debug!(?copies);
-    debug!(?ssa.direct_uses);
+    debug!(?direct_uses);
 
     // Invariant: `copies` must point to the head of an equivalence class.
     #[cfg(debug_assertions)]
     for &head in copies.iter() {
         assert_eq!(copies[head], head);
     }
+    debug_assert_eq!(copies[RETURN_PLACE], RETURN_PLACE);
 
-    copies
+    ssa.direct_uses = direct_uses;
+    ssa.copy_classes = copies;
 }
 
 #[derive(Debug)]

diff --git a/tests/codegen/fewer-names.rs b/tests/codegen/fewer-names.rs
@@ -13,8 +13,8 @@ pub fn sum(x: u32, y: u32) -> u32 {
 
     // NO-LABEL: define{{.*}}i32 @sum(i32 noundef %x, i32 noundef %y)
     // NO-NEXT:  start:
-    // NO-NEXT:    %z = add i32 %y, %x
-    // NO-NEXT:    ret i32 %z
+    // NO-NEXT:    %0 = add i32 %y, %x
+    // NO-NEXT:    ret i32 %0
     let z = x + y;
     z
 }
diff --git a/tests/codegen/mem-replace-big-type.rs b/tests/codegen/mem-replace-big-type.rs
@@ -13,8 +13,7 @@ pub struct Big([u64; 7]);
 pub fn replace_big(dst: &mut Big, src: Big) -> Big {
     // Back in 1.68, this emitted six `memcpy`s.
     // `read_via_copy` in 1.69 got that down to three.
-    // `write_via_move` it was originally down to the essential two, however
-    // with nrvo disabled it is back at 3
+    // `write_via_move` and nvro get this down to the essential two.
     std::mem::replace(dst, src)
 }
 
@@ -26,11 +25,9 @@ pub fn replace_big(dst: &mut Big, src: Big) -> Big {
 // For a large type, we expect exactly three `memcpy`s
 // CHECK-LABEL: define internal void @{{.+}}mem{{.+}}replace{{.+}}sret(%Big)
 // CHECK-NOT: call void @llvm.memcpy
-// CHECK: call void @llvm.memcpy.{{.+}}({{i8\*|ptr}} align 8 %result, {{i8\*|ptr}} align 8 %dest, i{{.*}} 56, i1 false)
+// CHECK: call void @llvm.memcpy.{{.+}}({{i8\*|ptr}} align 8 %0, {{i8\*|ptr}} align 8 %dest, i{{.*}} 56, i1 false)
 // CHECK-NOT: call void @llvm.memcpy
 // CHECK: call void @llvm.memcpy.{{.+}}({{i8\*|ptr}} align 8 %dest, {{i8\*|ptr}} align 8 %src, i{{.*}} 56, i1 false)
 // CHECK-NOT: call void @llvm.memcpy
-// CHECK: call void @llvm.memcpy.{{.+}}({{i8\*|ptr}} align 8 %0, {{i8\*|ptr}} align 8 %result, i{{.*}} 56, i1 false)
-// CHECK-NOT: call void @llvm.memcpy
 
 // CHECK-NOT: call void @llvm.memcpy
diff --git a/tests/codegen/var-names.rs b/tests/codegen/var-names.rs
@@ -9,7 +9,7 @@ pub fn test(a: u32, b: u32) -> u32 {
     // CHECK: %c = add i32 %a, %b
     let d = c;
     let e = d * a;
-    // CHECK-NEXT: %e = mul i32 %c, %a
+    // CHECK-NEXT: %0 = mul i32 %c, %a
     e
-    // CHECK-NEXT: ret i32 %e
+    // CHECK-NEXT: ret i32 %0
 }
diff --git a/tests/mir-opt/copy-prop/copy_propagation_arg.arg_src.CopyProp.diff b/tests/mir-opt/copy-prop/copy_propagation_arg.arg_src.CopyProp.diff
@@ -6,15 +6,17 @@
       let mut _0: i32;                     // return place in scope 0 at $DIR/copy_propagation_arg.rs:+0:27: +0:30
       let _2: i32;                         // in scope 0 at $DIR/copy_propagation_arg.rs:+1:9: +1:10
       scope 1 {
-          debug y => _2;                   // in scope 1 at $DIR/copy_propagation_arg.rs:+1:9: +1:10
+-         debug y => _2;                   // in scope 1 at $DIR/copy_propagation_arg.rs:+1:9: +1:10
++         debug y => _0;                   // in scope 1 at $DIR/copy_propagation_arg.rs:+1:9: +1:10
       }
 
       bb0: {
-          StorageLive(_2);                 // scope 0 at $DIR/copy_propagation_arg.rs:+1:9: +1:10
-          _2 = _1;                         // scope 0 at $DIR/copy_propagation_arg.rs:+1:13: +1:14
+-         StorageLive(_2);                 // scope 0 at $DIR/copy_propagation_arg.rs:+1:9: +1:10
+-         _2 = _1;                         // scope 0 at $DIR/copy_propagation_arg.rs:+1:13: +1:14
++         _0 = _1;                         // scope 0 at $DIR/copy_propagation_arg.rs:+1:13: +1:14
           _1 = const 123_i32;              // scope 1 at $DIR/copy_propagation_arg.rs:+2:5: +2:12
-          _0 = _2;                         // scope 1 at $DIR/copy_propagation_arg.rs:+3:5: +3:6
-          StorageDead(_2);                 // scope 0 at $DIR/copy_propagation_arg.rs:+4:1: +4:2
+-         _0 = _2;                         // scope 1 at $DIR/copy_propagation_arg.rs:+3:5: +3:6
+-         StorageDead(_2);                 // scope 0 at $DIR/copy_propagation_arg.rs:+4:1: +4:2
           return;                          // scope 0 at $DIR/copy_propagation_arg.rs:+4:2: +4:2
       }
   }

diff --git a/tests/mir-opt/inline/inline_into_box_place.main.Inline.diff b/tests/mir-opt/inline/inline_into_box_place.main.Inline.diff
@@ -16,8 +16,7 @@
 +         let mut _4: usize;               // in scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
 +         let mut _5: usize;               // in scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
 +         let mut _6: *mut u8;             // in scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-+         let mut _7: std::boxed::Box<std::vec::Vec<u32>>; // in scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-+         let mut _8: *const std::vec::Vec<u32>; // in scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
++         let mut _7: *const std::vec::Vec<u32>; // in scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
 +         scope 4 {
 +         }
 +     }
@@ -66,12 +65,9 @@
       bb3: {
 -         StorageDead(_1);                 // scope 0 at $DIR/inline_into_box_place.rs:+2:1: +2:2
 -         return;                          // scope 0 at $DIR/inline_into_box_place.rs:+2:2: +2:2
-+         StorageLive(_7);                 // scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-+         _7 = ShallowInitBox(move _6, std::vec::Vec<u32>); // scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-+         _8 = (((_7.0: std::ptr::Unique<std::vec::Vec<u32>>).0: std::ptr::NonNull<std::vec::Vec<u32>>).0: *const std::vec::Vec<u32>); // scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-+         (*_8) = move _2;                 // scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-+         _1 = move _7;                    // scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-+         StorageDead(_7);                 // scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
++         _1 = ShallowInitBox(move _6, std::vec::Vec<u32>); // scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
++         _7 = (((_1.0: std::ptr::Unique<std::vec::Vec<u32>>).0: std::ptr::NonNull<std::vec::Vec<u32>>).0: *const std::vec::Vec<u32>); // scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
++         (*_7) = move _2;                 // scope 3 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
 +         StorageDead(_2);                 // scope 0 at $DIR/inline_into_box_place.rs:+1:48: +1:49
 +         _0 = const ();                   // scope 0 at $DIR/inline_into_box_place.rs:+0:11: +2:2
 +         drop(_1) -> [return: bb1, unwind: bb2]; // scope 0 at $DIR/inline_into_box_place.rs:+2:1: +2:2

diff --git a/tests/mir-opt/inline/issue_58867_inline_as_ref_as_mut.b.Inline.after.mir b/tests/mir-opt/inline/issue_58867_inline_as_ref_as_mut.b.Inline.after.mir
@@ -8,22 +8,18 @@ fn b(_1: &mut Box<T>) -> &mut T {
     let mut _4: &mut std::boxed::Box<T>; // in scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
     scope 1 (inlined <Box<T> as AsMut<T>>::as_mut) { // at $DIR/issue_58867_inline_as_ref_as_mut.rs:8:7: 8:15
         debug self => _4;                // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        let mut _5: &mut T;              // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        let mut _6: std::boxed::Box<T>;  // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        let mut _7: *const T;            // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
+        let mut _5: std::boxed::Box<T>;  // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
+        let mut _6: *const T;            // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
     }
 
     bb0: {
         StorageLive(_2);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
         StorageLive(_3);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
         StorageLive(_4);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
         _4 = &mut (*_1);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
-        StorageLive(_5);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:7: +1:15
-        _6 = deref_copy (*_4);           // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        _7 = (((_6.0: std::ptr::Unique<T>).0: std::ptr::NonNull<T>).0: *const T); // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        _5 = &mut (*_7);                 // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        _3 = _5;                         // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        StorageDead(_5);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:7: +1:15
+        _5 = deref_copy (*_4);           // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
+        _6 = (((_5.0: std::ptr::Unique<T>).0: std::ptr::NonNull<T>).0: *const T); // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
+        _3 = &mut (*_6);                 // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
         _2 = &mut (*_3);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
         StorageDead(_4);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:14: +1:15
         _0 = &mut (*_2);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15

diff --git a/tests/mir-opt/inline/issue_58867_inline_as_ref_as_mut.d.Inline.after.mir b/tests/mir-opt/inline/issue_58867_inline_as_ref_as_mut.d.Inline.after.mir
@@ -7,21 +7,17 @@ fn d(_1: &Box<T>) -> &T {
     let mut _3: &std::boxed::Box<T>;     // in scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
     scope 1 (inlined <Box<T> as AsRef<T>>::as_ref) { // at $DIR/issue_58867_inline_as_ref_as_mut.rs:18:7: 18:15
         debug self => _3;                // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        let _4: &T;                      // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        let mut _5: std::boxed::Box<T>;  // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        let mut _6: *const T;            // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
+        let mut _4: std::boxed::Box<T>;  // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
+        let mut _5: *const T;            // in scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
     }
 
     bb0: {
         StorageLive(_2);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
         StorageLive(_3);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
         _3 = &(*_1);                     // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
-        StorageLive(_4);                 // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        _5 = deref_copy (*_3);           // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        _6 = (((_5.0: std::ptr::Unique<T>).0: std::ptr::NonNull<T>).0: *const T); // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        _4 = &(*_6);                     // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        _2 = _4;                         // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
-        StorageDead(_4);                 // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
+        _4 = deref_copy (*_3);           // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
+        _5 = (((_4.0: std::ptr::Unique<T>).0: std::ptr::NonNull<T>).0: *const T); // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
+        _2 = &(*_5);                     // scope 1 at $SRC_DIR/alloc/src/boxed.rs:LL:COL
         _0 = &(*_2);                     // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:5: +1:15
         StorageDead(_3);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+1:14: +1:15
         StorageDead(_2);                 // scope 0 at $DIR/issue_58867_inline_as_ref_as_mut.rs:+2:1: +2:2