[MCP] Optimize copies when src is used during backward propagation

Before this patch, redundant COPY couldn't be removed
for the following case:
  $R0 = OP ...
  ... // Read of %R0
  $R1 = COPY killed $R0

This patch adds support for tracking the users of
the source register during backward propagation, so
that we can remove the redundant COPY in the above case
and optimize it to:
  $R1 = OP ...
  ... // Replace all uses of %R0 with $R1

Upstream PR:
llvm/llvm-project#111130

Signed-off-by: Vladimir Radosavljevic <[email protected]>

llvm/lib/CodeGen/MachineCopyPropagation.cpp

@@ -109,6 +109,9 @@ class CopyTracker {
   struct CopyInfo {
     MachineInstr *MI = nullptr;
     MachineInstr *LastSeenUseInCopy = nullptr;
+    // EraVM local begin
+    SmallPtrSet<MachineInstr *, 4> SrcUsers;
+    // EraVM local end
     SmallVector<MCRegister, 4> DefRegs;
     bool Avail = false;
   };
@@ -182,6 +185,45 @@ class CopyTracker {
     }
   }
 
+  // EraVM local begin
+  /// Track copy's src users, and return false if that can't be done.
+  /// We can only track if we have a COPY instruction which source is
+  /// the same as the Reg.
+  bool trackSrcUsers(MCRegister Reg, MachineInstr &MI,
+                     const TargetRegisterInfo &TRI, const TargetInstrInfo &TII,
+                     bool UseCopyInstr) {
+    MCRegUnit RU = *TRI.regunits(Reg).begin();
+    MachineInstr *AvailCopy = findCopyDefViaUnit(RU, TRI);
+    if (!AvailCopy)
+      return false;
+
+    std::optional<DestSourcePair> CopyOperands =
+        isCopyInstr(*AvailCopy, TII, UseCopyInstr);
+    Register Src = CopyOperands->Source->getReg();
+
+    // Bail out, if the source of the copy is not the same as the Reg.
+    if (Src != Reg)
+      return false;
+
+    auto I = Copies.find(RU);
+    if (I == Copies.end())
+      return false;
+
+    I->second.SrcUsers.insert(&MI);
+    return true;
+  }
+
+  /// Return the users for a given register.
+  SmallPtrSet<MachineInstr *, 4> getSrcUsers(MCRegister Reg,
+                                             const TargetRegisterInfo &TRI) {
+    MCRegUnit RU = *TRI.regunits(Reg).begin();
+    auto I = Copies.find(RU);
+    if (I == Copies.end())
+      return {};
+    return I->second.SrcUsers;
+  }
+  // EraVM local end
+
   /// Add this copy's registers into the tracker's copy maps.
   void trackCopy(MachineInstr *MI, const TargetRegisterInfo &TRI,
                  const TargetInstrInfo &TII, bool UseCopyInstr) {
@@ -194,7 +236,9 @@ class CopyTracker {
 
     // Remember Def is defined by the copy.
     for (MCRegUnit Unit : TRI.regunits(Def))
-      Copies[Unit] = {MI, nullptr, {}, true};
+      // EraVM local begin
+      Copies[Unit] = {MI, nullptr, {}, {}, true};
+    // EraVM local end
 
     // Remember source that's copied to Def. Once it's clobbered, then
     // it's no longer available for copy propagation.
@@ -385,6 +429,10 @@ class MachineCopyPropagation : public MachineFunctionPass {
   bool hasImplicitOverlap(const MachineInstr &MI, const MachineOperand &Use);
   bool hasOverlappingMultipleDef(const MachineInstr &MI,
                                  const MachineOperand &MODef, Register Def);
+  // EraVM local begin
+  bool canUpdateSrcUsers(const MachineInstr &Copy,
+                         const MachineOperand &CopySrc);
+  // EraVM local end
 
   /// Candidates for deletion.
   SmallSetVector<MachineInstr *, 8> MaybeDeadCopies;
@@ -625,6 +673,28 @@ bool MachineCopyPropagation::hasOverlappingMultipleDef(
   return false;
 }
 
+// EraVM local begin
+/// Return true if it is safe to update the users of the source register of the
+/// copy.
+bool MachineCopyPropagation::canUpdateSrcUsers(const MachineInstr &Copy,
+                                               const MachineOperand &CopySrc) {
+  for (auto *SrcUser : Tracker.getSrcUsers(CopySrc.getReg(), *TRI)) {
+    if (hasImplicitOverlap(*SrcUser, CopySrc))
+      return false;
+
+    for (MachineOperand &MO : SrcUser->uses()) {
+      if (!MO.isReg() || !MO.isUse() || MO.getReg() != CopySrc.getReg())
+        continue;
+      if (MO.isTied() || !MO.isRenamable() ||
+          !isBackwardPropagatableRegClassCopy(Copy, *SrcUser,
+                                              MO.getOperandNo()))
+        return false;
+    }
+  }
+  return true;
+}
+// EraVM local end
+
 /// Look for available copies whose destination register is used by \p MI and
 /// replace the use in \p MI with the copy's source register.
 void MachineCopyPropagation::forwardUses(MachineInstr &MI) {
@@ -995,13 +1065,29 @@ void MachineCopyPropagation::propagateDefs(MachineInstr &MI) {
     if (hasOverlappingMultipleDef(MI, MODef, Def))
       continue;
 
+    // EraVM local begin
+    if (!canUpdateSrcUsers(*Copy, *CopyOperands->Source))
+      continue;
+    // EraVM local end
+
     LLVM_DEBUG(dbgs() << "MCP: Replacing " << printReg(MODef.getReg(), TRI)
                       << "\n     with " << printReg(Def, TRI) << "\n     in "
                       << MI << "     from " << *Copy);
 
     MODef.setReg(Def);
     MODef.setIsRenamable(CopyOperands->Destination->isRenamable());
 
+    // EraVM local begin
+    for (auto *SrcUser : Tracker.getSrcUsers(Src, *TRI)) {
+      for (MachineOperand &MO : SrcUser->uses()) {
+        if (!MO.isReg() || !MO.isUse() || MO.getReg() != Src)
+          continue;
+        MO.setReg(Def);
+        MO.setIsRenamable(CopyOperands->Destination->isRenamable());
+      }
+    }
+    // EraVM local end
+
     LLVM_DEBUG(dbgs() << "MCP: After replacement: " << MI << "\n");
     MaybeDeadCopies.insert(Copy);
     Changed = true;
@@ -1067,7 +1153,11 @@ void MachineCopyPropagation::BackwardCopyPropagateBlock(
               CopyDbgUsers[Copy].insert(&MI);
             }
           }
-        } else {
+          // EraVM local begin
+        } else if (!Tracker.trackSrcUsers(MO.getReg().asMCReg(), MI, *TRI, *TII,
+                                          UseCopyInstr)) {
+          // If we can't track the source users, invalidate the register.
+          // EraVM local end
           Tracker.invalidateRegister(MO.getReg().asMCReg(), *TRI, *TII,
                                      UseCopyInstr);
         }

llvm/test/CodeGen/EraVM/machine-cp-backward-users.mir

@@ -18,9 +18,8 @@ body:             |
     ; CHECK-LABEL: name: test
     ; CHECK: liveins: $r1, $r2, $r4
     ; CHECK-NEXT: {{  $}}
-    ; CHECK-NEXT: renamable $r3 = ADDirr_s i256 1, killed renamable $r1, i256 0
-    ; CHECK-NEXT: dead $r0 = SUBrrr_v renamable $r3, renamable $r2, i256 0, implicit-def $flags
-    ; CHECK-NEXT: renamable $r1 = COPY killed renamable $r3
+    ; CHECK-NEXT: renamable $r1 = ADDirr_s i256 1, killed renamable $r1, i256 0
+    ; CHECK-NEXT: dead $r0 = SUBrrr_v renamable $r1, renamable $r2, i256 0, implicit-def $flags
     ; CHECK-NEXT: RET 0, implicit $r1
     renamable $r3 = ADDirr_s i256 1, killed renamable $r1, i256 0
     dead $r0 = SUBrrr_v renamable $r3, renamable $r2, i256 0, implicit-def $flags