[PostRASink]Add register dependency check for implicit operands

[junbuml]

Summary

This change extend the register dependency check for implicit operands in Copy instructions.
Fixes PR36902.

[junbuml]

Ping ?

[thegameg]

Sorry for the delay.

Thanks for fixing this. What do you think about using register units throughout the pass? I found them so much clearer to use when dealing with register aliases. @MatzeB has a great presentation on it: 2016 LLVM Developers' Meeting: M. Braun "Dealing with Register Hierarchies".

lib/CodeGen/MachineSink.cpp

1080

Would MO.readsReg() be a better fit here?

1128

I would comment on what the unsigned values represent.

[junbuml]

What do you think about using register units throughout the pass? I found them so much clearer to use when dealing with register aliases.

I think I saw the case where LiveRegUnits is used to find a scratch reg. It seems possible to use it here as well. Let me take a little closer look at it. I think it should be a separate NFC patch.

lib/CodeGen/MachineSink.cpp

1080

One case where readReg works better is when the source of Copy is undef, but I'm not perfectly clear if it's safe to skip the internal read.

[junbuml]

Just minor update in comments.

[junbuml]

lib/CodeGen/MachineSink.cpp

1086

I keep isUse() here with FIXME since I'm not perfectly sure if skipping the internal read is safe in all other targets.

[thegameg]

In D44958#1059982, @junbuml wrote:

I think I saw the case where LiveRegUnits is used to find a scratch reg. It seems possible to use it here as well. Let me take a little closer look at it. I think it should be a separate NFC patch.

Even without LiveRegUnits, I think you could replace ModifiedRegs and UsedRegs to be ModifiedRegUnits and UsedRegUnits, and use MCRegUnitIterators to check if registers alias.

I'm ok to commit this and follow-up in separate patches for the rest. Thanks!

lib/CodeGen/MachineSink.cpp

1086

Ok, sounds good to me.

1127–1129

Any particular reason why we need to store both the register and the operand idx?

[junbuml]

Even without LiveRegUnits, I think you could replace ModifiedRegs and UsedRegs to be ModifiedRegUnits and UsedRegUnits, and use MCRegUnitIterators to check if registers alias.
I'm ok to commit this and follow-up in separate patches for the rest. Thanks!

Make sense to me. I will do this in a separate patch.

lib/CodeGen/MachineSink.cpp

1127–1129

The operand idx will be used in clearKillFlags() to set the kill flag in src reg of the Copy sunk if the src reg is killed in an instruction placed between Copy and the end of the block.

[thegameg]

lib/CodeGen/MachineSink.cpp

1127–1129

Right, I was just curious why not only use the operand idx and use MI.getOperand(OpIdx).getReg() to retrieve the register.

[junbuml]

lib/CodeGen/MachineSink.cpp

1127–1129

I see what you mean. I will change it to keep only operand idx and get the register using MI.getOperand(OpIdx).getReg() .

[junbuml]

Addressed Francis' comment.

[thegameg]

LGTM, thanks!

lib/CodeGen/MachineSink.cpp

 } // namespace
 
 char PostRAMachineSinking::ID = 0;
 char &llvm::PostRAMachineSinkingID = PostRAMachineSinking::ID;
 
 INITIALIZE_PASS(PostRAMachineSinking, "postra-machine-sink",
                 "PostRA Machine Sink", false, false)
 
+static bool aliasWithRegsInLiveIn(MachineBasicBlock *SI,
+                                  SmallSet<unsigned, 8> &AliasedRegs) {
+  for (const auto LI : SI->liveins())
+    if (AliasedRegs.count(LI.PhysReg))
+      return true;
+  return false;
+}
+
 static MachineBasicBlock *
 getSingleLiveInSuccBB(MachineBasicBlock &CurBB,
                       ArrayRef<MachineBasicBlock *> SinkableBBs, unsigned Reg,
                       const TargetRegisterInfo *TRI) {
   SmallSet<unsigned, 8> AliasedRegs;
   for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
     AliasedRegs.insert(*AI);
 
getSingleLiveInSuccBB(MachineBasicBlock &CurBB,
   // Reg is not live-in to any sinkable successors.
   if (!BB)
     return nullptr;
 
   // Check if any register aliased with Reg is live-in in other successors.
   for (auto *SI : CurBB.successors()) {
     if (SI == BB)
       continue;
-    for (const auto LI : SI->liveins())
-      if (AliasedRegs.count(LI.PhysReg))
-        return nullptr;
+    if (aliasWithRegsInLiveIn(SI, AliasedRegs))
   }
   return BB;
 }
 
+static MachineBasicBlock *getSingleLiveInSuccBB(
+    MachineBasicBlock &CurBB, ArrayRef<MachineBasicBlock *> SinkableBBs,
+    ArrayRef<unsigned> DefedRegsInCopy, const TargetRegisterInfo *TRI) {
+  MachineBasicBlock *SingleBB = nullptr;
+  for (auto DefReg : DefedRegsInCopy) {
+    MachineBasicBlock *BB =
+        getSingleLiveInSuccBB(CurBB, SinkableBBs, DefReg, TRI);
+    if (!BB || (SingleBB && SingleBB != BB))
+      return nullptr;
+    SingleBB = BB;
+  }
+  return SingleBB;
+}
+
+static void
+clearKillFlags(MachineInstr *MI, MachineBasicBlock &CurBB,
+               SmallVectorImpl<std::pair<unsigned, unsigned>> &UsedRegsInCopy,
+               BitVector &UsedRegs, const TargetRegisterInfo *TRI) {
+  for (auto U : UsedRegsInCopy) {
+    unsigned OpIdx = U.first;
+    unsigned SrcReg = U.second;
+    if (UsedRegs[SrcReg]) {
+      MachineBasicBlock::iterator NI = std::next(MI->getIterator());
+      for (MachineInstr &UI : make_range(NI, CurBB.end())) {
+        if (UI.killsRegister(SrcReg, TRI)) {
+          UI.clearRegisterKills(SrcReg, TRI);
+          MI->getOperand(OpIdx).setIsKill(true);
+          break;
+        }
+      }
+    }
+  }
+}
+
+static void
+updateLiveIn(MachineBasicBlock *SuccBB,
+             SmallVectorImpl<std::pair<unsigned, unsigned>> &UsedRegsInCopy,
+             SmallVectorImpl<unsigned> &DefedRegsInCopy) {
+  for (auto DefReg : DefedRegsInCopy)
+    SuccBB->removeLiveIn(DefReg);
+  for (auto U : UsedRegsInCopy) {
+    unsigned Reg = U.second;
+    if (!SuccBB->isLiveIn(Reg))
+      SuccBB->addLiveIn(Reg);
+  }
+}
+
+static bool hasRegisterDependency(
+    MachineInstr *MI,
+    SmallVectorImpl<std::pair<unsigned, unsigned>> &UsedRegsInCopy,
+    SmallVectorImpl<unsigned> &DefedRegsInCopy, BitVector &ModifiedRegs,
+    BitVector &UsedRegs) {
+  bool HasRegDependency = false;
+  for (unsigned i = 0, e = MI->getNumOperands(); i != e; ++i) {
+    MachineOperand &MO = MI->getOperand(i);
+    if (!MO.isReg())
+      continue;
+    unsigned Reg = MO.getReg();
+    if (!Reg)
+      continue;
+    if (MO.isDef()) {
+      if (ModifiedRegs[Reg] || UsedRegs[Reg]) {
+        HasRegDependency = true;
+        break;
+      }
+      DefedRegsInCopy.push_back(Reg);
+    } else if (MO.isUse()) {
+      if (ModifiedRegs[Reg]) {
+        HasRegDependency = true;
+        break;
+      }
+      UsedRegsInCopy.push_back(std::make_pair(i, Reg));
+    }
+  }
+  return HasRegDependency;
+}
+
 bool PostRAMachineSinking::tryToSinkCopy(MachineBasicBlock &CurBB,
                                          MachineFunction &MF,
                                          const TargetRegisterInfo *TRI,
                                          const TargetInstrInfo *TII) {
   SmallVector<MachineBasicBlock *, 2> SinkableBBs;
   // FIXME: For now, we sink only to a successor which has a single predecessor
   // so that we can directly sink COPY instructions to the successor without
   // adding any new block or branch instruction.
   for (MachineBasicBlock *SI : CurBB.successors())
     if (!SI->livein_empty() && SI->pred_size() == 1)
       SinkableBBs.push_back(SI);
 
   if (SinkableBBs.empty())
     return false;
 
   bool Changed = false;
 
   // Track which registers have been modified and used between the end of the
   // block and the current instruction.
   ModifiedRegs.reset();
   UsedRegs.reset();
 
+
   for (auto I = CurBB.rbegin(), E = CurBB.rend(); I != E;) {
     MachineInstr *MI = &*I;
     ++I;
 
     // Do not move any instruction across function call.
     if (MI->isCall())
       return false;
 
     if (!MI->isCopy() || !MI->getOperand(0).isRenamable()) {
       TII->trackRegDefsUses(*MI, ModifiedRegs, UsedRegs, TRI);
       continue;
     }
 
-    unsigned DefReg = MI->getOperand(0).getReg();
-    unsigned SrcReg = MI->getOperand(1).getReg();
+    SmallVector<std::pair<unsigned, unsigned>, 2> UsedRegsInCopy;
+    SmallVector<unsigned, 2> DefedRegsInCopy;
+
     // Don't sink the COPY if it would violate a register dependency.
-    if (ModifiedRegs[DefReg] || ModifiedRegs[SrcReg] || UsedRegs[DefReg]) {
+    if (hasRegisterDependency(MI, UsedRegsInCopy, DefedRegsInCopy, ModifiedRegs,
+                              UsedRegs)) {
       TII->trackRegDefsUses(*MI, ModifiedRegs, UsedRegs, TRI);
       continue;
     }
-
+    assert((!UsedRegsInCopy.empty() && !DefedRegsInCopy.empty()) &&
+           "Unexpect SrcReg or DefReg");
     MachineBasicBlock *SuccBB =
-        getSingleLiveInSuccBB(CurBB, SinkableBBs, DefReg, TRI);
+        getSingleLiveInSuccBB(CurBB, SinkableBBs, DefedRegsInCopy, TRI);
     // Don't sink if we cannot find a single sinkable successor in which Reg
     // is live-in.
     if (!SuccBB) {
       TII->trackRegDefsUses(*MI, ModifiedRegs, UsedRegs, TRI);
       continue;
     }
     assert((SuccBB->pred_size() == 1 && *SuccBB->pred_begin() == &CurBB) &&
            "Unexpected predecessor");
 
     // Clear the kill flag if SrcReg is killed between MI and the end of the
     // block.
-    if (UsedRegs[SrcReg]) {
-      MachineBasicBlock::iterator NI = std::next(MI->getIterator());
-      for (MachineInstr &UI : make_range(NI, CurBB.end())) {
-        if (UI.killsRegister(SrcReg, TRI)) {
-          UI.clearRegisterKills(SrcReg, TRI);
-          MI->getOperand(1).setIsKill(true);
-          break;
-        }
-      }
-    }
-
+    clearKillFlags(MI, CurBB, UsedRegsInCopy, UsedRegs, TRI);
     MachineBasicBlock::iterator InsertPos = SuccBB->getFirstNonPHI();
     SuccBB->splice(InsertPos, &CurBB, MI);
-    SuccBB->removeLiveIn(DefReg);
-    if (!SuccBB->isLiveIn(SrcReg))
-      SuccBB->addLiveIn(SrcReg);
+    updateLiveIn(SuccBB, UsedRegsInCopy, DefedRegsInCopy);
 
     Changed = true;
     ++NumPostRACopySink;
   }
   return Changed;
 }
 
 bool PostRAMachineSinking::runOnMachineFunction(MachineFunction &MF) {

test/CodeGen/SystemZ/no-postra-sink.mir

+# RUN: llc -mtriple=s390x-linux-gnu -mcpu=z13 -run-pass=postra-machine-sink -verify-machineinstrs -o - %s  | FileCheck %s
+
+---
+# Don't sink COPY to bb.2 since SLLK define r13l that is aliased with r12q.
+# CHECK-LABEL: name: donot_sink_copy
+# CHECK-LABEL: bb.0:
+# CHECK: renamable $r0l = COPY renamable $r12l, implicit killed $r12q
+# CHECK-LABEL: bb.2:
+# CHECK-NOT: COPY
+name:            donot_sink_copy
+tracksRegLiveness: true
+body:             |
+  bb.0 :
+    successors: %bb.1, %bb.2
+    liveins: $r2d, $r3d, $r4d, $r5d, $r12q
+  
+    renamable $r0l = COPY renamable $r12l, implicit killed $r12q
+    renamable $r13l = SLLK renamable $r4l, $noreg, 1
+    CHIMux renamable $r3l, 0, implicit-def $cc, implicit killed $r3d
+    BRC 14, 6, %bb.2, implicit killed $cc
+    J %bb.1
+  
+  bb.1:
+    successors: 
+  
+  bb.2:
+    successors:
+    liveins: $r2d, $r4d, $r5d, $r0l, $r13l
+  
+    renamable $r0d = LGFR killed renamable $r0l
+    renamable $r11d = LGFR killed renamable $r13l
+...
+
+# Don't sink COPY to bb.2 since SLLK use r1l that is aliased with r0q.
+# CHECK-LABEL: name: donot_sink_copy2
+# CHECK-LABEL: bb.0:
+# CHECK: renamable $r0l = COPY renamable $r12l, implicit-def $r0q
+# CHECK-LABEL: bb.2:
+# CHECK-NOT: COPY
+name:            donot_sink_copy2
+tracksRegLiveness: true
+body:             |
+  bb.0 :
+    successors: %bb.1, %bb.2
+    liveins: $r2d, $r3d, $r4d, $r5d, $r12q
+
+    renamable $r0l = COPY renamable $r12l, implicit def $r0q
+    renamable $r13l = SLLK renamable $r1l, $noreg, 1
+    CHIMux renamable $r3l, 0, implicit-def $cc, implicit killed $r3d
+    BRC 14, 6, %bb.2, implicit killed $cc
+    J %bb.1
+
+  bb.1:
+    successors:
+
+  bb.2:
+    successors:
+    liveins: $r2d, $r4d, $r5d, $r0l, $r13l
+
+    renamable $r0d = LGFR killed renamable $r0l
+    renamable $r11d = LGFR killed renamable $r13l
+...
+