[PostRASink] extend the live-in check for all aliased registers

[junbuml]

Summary

Extend the live-in check for all aliased registers so that we can
allow sinking Copy instructions even when only implicit def is in successor's
live-in.

[qcolombet]

Hi,

LGTM with small nits.

Cheers,
Quentin

lib/CodeGen/MachineSink.cpp

987–988

This should be const, right?

1011–1012

Ditto

[junbuml]

Thanks Quentin for the review.

lib/CodeGen/MachineSink.cpp

static bool aliasWithRegsInLiveIn(MachineBasicBlock &MBB, unsigned Reg,
                                   const TargetRegisterInfo *TRI) {
   LiveRegUnits LiveInRegUnits(*TRI);
   LiveInRegUnits.addLiveIns(MBB);
   return !LiveInRegUnits.available(Reg);
 }
 
 static MachineBasicBlock *
 getSingleLiveInSuccBB(MachineBasicBlock &CurBB,
-                      ArrayRef<MachineBasicBlock *> SinkableBBs, unsigned Reg,
-                      const TargetRegisterInfo *TRI) {
+                      const SmallPtrSetImpl<MachineBasicBlock *> &SinkableBBs,
+                      unsigned Reg, const TargetRegisterInfo *TRI) {
   // Try to find a single sinkable successor in which Reg is live-in.
   MachineBasicBlock *BB = nullptr;
   for (auto *SI : SinkableBBs) {
-    if (SI->isLiveIn(Reg)) {
+    if (aliasWithRegsInLiveIn(*SI, Reg, TRI)) {
       // If BB is set here, Reg is live-in to at least two sinkable successors,
       // so quit.
       if (BB)
         return nullptr;
       BB = SI;
     }
   }
   // 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;
-    if (aliasWithRegsInLiveIn(*SI, Reg, TRI))
+    if (!SinkableBBs.count(SI) && aliasWithRegsInLiveIn(*SI, Reg, TRI))
       return nullptr;
   }
   return BB;
 }
 
-static MachineBasicBlock *getSingleLiveInSuccBB(
-    MachineBasicBlock &CurBB, ArrayRef<MachineBasicBlock *> SinkableBBs,
-    ArrayRef<unsigned> DefedRegsInCopy, const TargetRegisterInfo *TRI) {
+static MachineBasicBlock *
+getSingleLiveInSuccBB(MachineBasicBlock &CurBB,
+                      const SmallPtrSetImpl<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;
   }
static bool hasRegisterDependency(MachineInstr *MI,
   }
   return HasRegDependency;
 }
 
 bool PostRAMachineSinking::tryToSinkCopy(MachineBasicBlock &CurBB,
                                          MachineFunction &MF,
                                          const TargetRegisterInfo *TRI,
                                          const TargetInstrInfo *TII) {
-  SmallVector<MachineBasicBlock *, 2> SinkableBBs;
+  SmallPtrSet<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);
+      SinkableBBs.insert(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.

test/CodeGen/AArch64/post-ra-machine-sink.mir

body: |
   bb.1:
     $x0 = COPY $xzr
     RET $x0
 
   bb.2:
     liveins: $w0, $w19
     $w0 = ADDWrr $w0, $w19
     RET $x0
+...
+
 ---
+# Sink w19 to %bb.3 through %bb.2.
+# CHECK-LABEL: name: sinkcopy8
+# CHECK-LABEL: bb.0:
+# CHECK-NOT: renamable $w19 = COPY $w0, implicit-def $x19
+# CHECK-LABEL: bb.2:
+# CHECK: $w1 = ADDWrr $w1, $w0, implicit $x0
+# CHECK-LABEL: bb.3:
+# CHECK: liveins: $x1, $w0
+# CHECK: renamable $w19 = COPY killed $w0, implicit-def $x19
+name: sinkcopy8
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $w0, $x1
+    $w1 = SUBSWri $w1, 1, 0, implicit-def $nzcv
+    renamable $w19 = COPY $w0, implicit-def $x19
+    Bcc 11, %bb.2, implicit $nzcv
 
+  bb.1:
+    liveins: $x0
+    $w19 = COPY $wzr
+    RET $x0
+
+  bb.2:
+    liveins: $w0, $x1, $x19
+    $w1 = ADDWrr $w1, $w0, implicit killed $x0
+
+  bb.3:
+    liveins: $x1, $x19
+    $x0 = ADDXrr $x1, $x19
+    RET $x0
+...
+
+---
 # Don't sink w19 as w0 is defined in bb.0.
 # CHECK-LABEL: name: donotsinkcopy1
 # CHECK-LABEL: bb.0:
 # CHECK: renamable $w19 = COPY $w0
 # CHECK: $w0 = LDRWui $sp, 0
 name: donotsinkcopy1
 tracksRegLiveness: true
 body: |