[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 ↗

(On Diff #140009)

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

1127 ↗

(On Diff #140009)

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 ↗

(On Diff #140009)

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

1091 ↗

(On Diff #141676)

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

1091 ↗

(On Diff #141676)

Ok, sounds good to me.

1126 ↗

(On Diff #140009)

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

1126 ↗

(On Diff #140009)

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

1126 ↗

(On Diff #140009)

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

1126 ↗

(On Diff #140009)

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!

llvm/trunk/lib/CodeGen/MachineSink.cpp

public:
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
 
   MachineFunctionProperties getRequiredProperties() const override {
     return MachineFunctionProperties().set(
         MachineFunctionProperties::Property::NoVRegs);
   }
 
 private:
   /// Track which registers have been modified and used.
   BitVector ModifiedRegs, UsedRegs;
 
   /// Sink Copy instructions unused in the same block close to their uses in
   /// successors.
   bool tryToSinkCopy(MachineBasicBlock &BB, MachineFunction &MF,
                      const TargetRegisterInfo *TRI, const TargetInstrInfo *TII);
 };
 } // 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<unsigned> &UsedOpsInCopy,
+                           BitVector &UsedRegs, const TargetRegisterInfo *TRI) {
+  for (auto U : UsedOpsInCopy) {
+    MachineOperand &MO = MI->getOperand(U);
+    unsigned SrcReg = MO.getReg();
+    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);
+          MO.setIsKill(true);
+          break;
+        }
+      }
+    }
+  }
+}
+
+static void updateLiveIn(MachineInstr *MI, MachineBasicBlock *SuccBB,
+                         SmallVectorImpl<unsigned> &UsedOpsInCopy,
+                         SmallVectorImpl<unsigned> &DefedRegsInCopy) {
+  for (auto DefReg : DefedRegsInCopy)
+    SuccBB->removeLiveIn(DefReg);
+  for (auto U : UsedOpsInCopy) {
+    unsigned Reg = MI->getOperand(U).getReg();
+    if (!SuccBB->isLiveIn(Reg))
+      SuccBB->addLiveIn(Reg);
+  }
+}
+
+static bool hasRegisterDependency(MachineInstr *MI,
+                                  SmallVectorImpl<unsigned> &UsedOpsInCopy,
+                                  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);
+
+      // FIXME: instead of isUse(), readsReg() would be a better fix here,
+      // For example, we can ignore modifications in reg with undef. However,
+      // it's not perfectly clear if skipping the internal read is safe in all
+      // other targets.
+    } else if (MO.isUse()) {
+      if (ModifiedRegs[Reg]) {
+        HasRegDependency = true;
+        break;
+      }
+      UsedOpsInCopy.push_back(i);
+    }
+  }
+  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 (auto I = CurBB.rbegin(), E = CurBB.rend(); I != E;) {
     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();
+    // Track the operand index for use in Copy.
+    SmallVector<unsigned, 2> UsedOpsInCopy;
+    // Track the register number defed in Copy.
+    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, UsedOpsInCopy, DefedRegsInCopy, ModifiedRegs,
+                              UsedRegs)) {
       TII->trackRegDefsUses(*MI, ModifiedRegs, UsedRegs, TRI);
       continue;
     }
-
+    assert((!UsedOpsInCopy.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, UsedOpsInCopy, UsedRegs, TRI);
     MachineBasicBlock::iterator InsertPos = SuccBB->getFirstNonPHI();
     SuccBB->splice(InsertPos, &CurBB, MI);
-    SuccBB->removeLiveIn(DefReg);
-    if (!SuccBB->isLiveIn(SrcReg))
-      SuccBB->addLiveIn(SrcReg);
+    updateLiveIn(MI, SuccBB, UsedOpsInCopy, DefedRegsInCopy);
 
     Changed = true;
     ++NumPostRACopySink;
   }
   return Changed;
 }
 
 bool PostRAMachineSinking::runOnMachineFunction(MachineFunction &MF) {

llvm/trunk/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
+...
+