[BranchFolding] Fix live-in's when hoisting code

[Ka-Ka]

Summary

When the branch folder hoist code into a predecessor it adjust live-in's
in the blocks it hoist code from. However it fail to handle hoisted code
that contain a defed register that originally is live-in in the block
through a super register.

This is fixed by replacing the live-in handling code with calls to
utility functions in LivePhysRegs.

[Ka-Ka]

Adding reviewers.

The fault was found in an out of tree target but was reproduced with a SystemZ testcase. Ulrich Weigand, can you review the SyztemZ testcase?

include/llvm/CodeGen/LivePhysRegs.h

 /// Adds registers contained in \p LiveRegs to the block live-in list of \p MBB.
 /// Does not add reserved registers.
 void addLiveIns(MachineBasicBlock &MBB, const LivePhysRegs &LiveRegs);
 
 /// Convenience function combining computeLiveIns() and addLiveIns().
 void computeAndAddLiveIns(LivePhysRegs &LiveRegs,
                           MachineBasicBlock &MBB);
 
+/// Convenience function for recomputing live-in's for \p MBB.
+static inline void recomputeLiveIns(MachineBasicBlock &MBB) {
+  LivePhysRegs LPR;
+  MBB.clearLiveIns();
+  computeAndAddLiveIns(LPR, MBB);
+}
+
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_LIVEPHYSREGS_H

lib/CodeGen/BranchFolding.cpp

bool BranchFolder::HoistCommonCodeInSuccs(MachineBasicBlock *MBB) {
   // point to the end of the block.
   SmallSet<unsigned, 4> Uses, Defs;
   MachineBasicBlock::iterator Loc =
     findHoistingInsertPosAndDeps(MBB, TII, TRI, Uses, Defs);
   if (Loc == MBB->end())
     return false;
 
   bool HasDups = false;
-  SmallVector<unsigned, 4> LocalDefs, LocalKills;
   SmallSet<unsigned, 4> ActiveDefsSet, AllDefsSet;
   MachineBasicBlock::iterator TIB = TBB->begin();
   MachineBasicBlock::iterator FIB = FBB->begin();
   MachineBasicBlock::iterator TIE = TBB->end();
   MachineBasicBlock::iterator FIE = FBB->end();
   while (TIB != TIE && FIB != FIE) {
     // Skip dbg_value instructions. These do not count.
     TIB = skipDebugInstructionsForward(TIB, TIE);
while (TIB != TIE && FIB != FIE) {
     // Remove kills from ActiveDefsSet, these registers had short live ranges.
     for (const MachineOperand &MO : TIB->operands()) {
       if (!MO.isReg() || !MO.isUse() || !MO.isKill())
         continue;
       unsigned Reg = MO.getReg();
       if (!Reg)
         continue;
       if (!AllDefsSet.count(Reg)) {
-        LocalKills.push_back(Reg);
         continue;
       }
       if (TargetRegisterInfo::isPhysicalRegister(Reg)) {
         for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
           ActiveDefsSet.erase(*AI);
       } else {
         ActiveDefsSet.erase(Reg);
       }
     }
 
     // Track local defs so we can update liveins.
     for (const MachineOperand &MO : TIB->operands()) {
       if (!MO.isReg() || !MO.isDef() || MO.isDead())
         continue;
       unsigned Reg = MO.getReg();
       if (!Reg || TargetRegisterInfo::isVirtualRegister(Reg))
         continue;
-      LocalDefs.push_back(Reg);
       addRegAndItsAliases(Reg, TRI, ActiveDefsSet);
       addRegAndItsAliases(Reg, TRI, AllDefsSet);
     }
 
     HasDups = true;
     ++TIB;
     ++FIB;
   }
 
   if (!HasDups)
     return false;
 
   MBB->splice(Loc, TBB, TBB->begin(), TIB);
   FBB->erase(FBB->begin(), FIB);
 
-  // Update livein's.
-  bool ChangedLiveIns = false;
-  for (unsigned i = 0, e = LocalDefs.size(); i != e; ++i) {
-    unsigned Def = LocalDefs[i];
-    if (ActiveDefsSet.count(Def)) {
-      TBB->addLiveIn(Def);
-      FBB->addLiveIn(Def);
-      ChangedLiveIns = true;
-    }
-  }
-  for (unsigned K : LocalKills) {
-    TBB->removeLiveIn(K);
-    FBB->removeLiveIn(K);
-    ChangedLiveIns = true;
-  }
-
-  if (ChangedLiveIns) {
-    TBB->sortUniqueLiveIns();
-    FBB->sortUniqueLiveIns();
+  if (UpdateLiveIns) {
+    recomputeLiveIns(*TBB);
+    recomputeLiveIns(*FBB);
   }
 
   ++NumHoist;
   return true;
 }

test/CodeGen/SystemZ/branch-folder-hoist-livein.mir

+# RUN: llc -verify-machineinstrs -O1 -mtriple=s390x-ibm-linux -o - %s -run-pass=branch-folder | FileCheck %s
+--- |
+  target datalayout = "E-m:e-i1:8:16-i8:8:16-i64:64-f128:64-a:8:16-n32:64"
+  target triple = "s390x-ibm-linux"
+
+  @b = external global i16, align 1
+  @d = external global i16, align 1
+
+  define void @f1() {
+    ret void
+  }
+
+...
+---
+name:            f1
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    successors: %bb.2(0x7fffffff), %bb.1(0x00000001)
+    liveins:
+
+    renamable $r1d = LGRL @b :: (load 4 from got, align 8)
+    renamable $r1l = LH killed renamable $r1d, 0, $noreg, implicit-def $r1d :: (dereferenceable load 1 from @b)
+    renamable $r2l = LHI 0
+    renamable $r3d = LGRL @d :: (load 4 from got, align 8)
+    renamable $r4d = LLILL 0, implicit-def $r4q
+
+    CHI killed renamable $r2l, 0, implicit-def $cc
+    BRC 14, 6, %bb.2, implicit killed $cc
+    J %bb.1
+
+  bb.1:
+    successors:
+    liveins: $r1l, $r3d, $r4q
+
+    renamable $r4d = COPY killed renamable $r4d, implicit killed $r4q
+    STH renamable $r1l, killed renamable $r3d, 0, $noreg, implicit killed $r4d :: (store 1 into @d)
+
+  bb.2:
+    liveins: $r1l, $r3d, $r4q
+
+    renamable $r4d = COPY killed renamable $r4d, implicit killed $r4q
+    STH renamable $r1l, killed renamable $r3d, 0, $noreg, implicit killed $r4d :: (store 1 into @d)
+    Return
+
+...
+
+# CHECK: renamable $r4d = COPY killed renamable $r4d, implicit killed $r4q
+# CHECK-NEXT: CHI killed renamable $r2l, 0, implicit-def $cc
+# CHECK-NEXT: BRC 14, 6, %bb.2, implicit killed $cc
+# CHECK-NEXT: {{^  $}}
+# CHECK-NEXT: bb.1:
+# CHECK-NEXT: successors:
+# CHECK-NEXT: liveins: $r1l, $r3d, $r4d
+
+# CHECK: bb.2:
+# CHECK-NEXT: liveins: $r1l, $r3d, $r4d

test/CodeGen/X86/tail-call-conditional.mir

body:             |
 
   ; CHECK: bb.1:
   ; CHECK-NEXT: successors: %bb.2({{[^)]+}}){{$}}
   ; CHECK-NEXT: liveins: $rax, $rsi
   ; CHECK-NEXT: {{^  $}}
   ; CHECK-NEXT: $rdi = COPY $rsi
   ; CHECK-NEXT: $rsi = COPY $rax
   ; CHECK-NEXT: CMP64ri8 $rax, 9, implicit-def $eflags
-  ; CHECK-NEXT: TCRETURNdi64cc @f1, 0, 3, csr_64, implicit $rsp, implicit $eflags, implicit $ssp, implicit $rsp, implicit $rdi, implicit $rsi, implicit $rax, implicit-def $rax, implicit $hsi, implicit-def $hsi, implicit $sil, implicit-def $sil, implicit $si, implicit-def $si, implicit $esi, implicit-def $esi, implicit $rsi, implicit-def $rsi, implicit $hdi, implicit-def $hdi, implicit $dil, implicit-def $dil, implicit $di, implicit-def $di, implicit $edi, implicit-def $edi, implicit $rdi, implicit-def $rdi, implicit $hax, implicit-def $hax, implicit $ah, implicit-def $ah, implicit $al, implicit-def $al, implicit $ax, implicit-def $ax, implicit $eax, implicit-def $eax
+  ; CHECK-NEXT: TCRETURNdi64cc @f1, 0, 3, csr_64, implicit $rsp, implicit $eflags, implicit $ssp, implicit $rsp, implicit $rdi, implicit $rsi, implicit $rdi, implicit-def $rdi, implicit $hsi, implicit-def $hsi, implicit $sil, implicit-def $sil, implicit $si, implicit-def $si, implicit $esi, implicit-def $esi, implicit $rsi, implicit-def $rsi, implicit $hdi, implicit-def $hdi, implicit $dil, implicit-def $dil, implicit $di, implicit-def $di, implicit $edi, implicit-def $edi
 
   bb.1:
     successors: %bb.2, %bb.3
     liveins: $rax, $rsi
 
     CMP64ri8 $rax, 9, implicit-def $eflags
     JA_1 %bb.3, implicit $eflags
     JMP_1 %bb.2
 
   bb.2:
     liveins: $rax, $rsi
 
     $rdi = COPY $rsi
     $rsi = COPY $rax
 
     TCRETURNdi64 @f1, 0, csr_64, implicit $rsp, implicit $rdi, implicit $rsi
 
   ; CHECK: bb.2:
-  ; CHECK-NEXT: liveins: $rax, $rdi, $rsi
+  ; CHECK-NEXT: liveins: $rdi, $rsi
   ; CHECK-NEXT: {{^  $}}
   ; CHECK-NEXT: TCRETURNdi64 @f2, 0, csr_64, implicit $rsp, implicit $rdi, implicit $rsi
 
   bb.3:
     liveins: $rax, $rsi
 
     $rdi = COPY $rsi
     $rsi = COPY $rax
     TCRETURNdi64 @f2, 0, csr_64, implicit $rsp, implicit $rdi, implicit $rsi
 
   bb.4:
     dead $eax = MOV32ri64 123, implicit-def $rax
     RET 0, $rax
 
 ...