1 files changed, 16 insertions, 9 deletions

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 55aa1d4..c6c48cf 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -455,12 +455,14 @@ static unsigned EstimateRuntime(MachineBasicBlock::iterator I,
 // with a conditional branch to the next block, optimize by reversing the
 // test and conditionally branching to SuccMBB instead.
 static void FixTail(MachineBasicBlock *CurMBB, MachineBasicBlock *SuccBB,
-                    const TargetInstrInfo *TII) {
+                    const TargetInstrInfo *TII, const DebugLoc &BranchDL) {
   MachineFunction *MF = CurMBB->getParent();
   MachineFunction::iterator I = std::next(MachineFunction::iterator(CurMBB));
   MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
   SmallVector<MachineOperand, 4> Cond;
   DebugLoc dl = CurMBB->findBranchDebugLoc();
+  if (!dl)
+    dl = BranchDL;
   if (I != MF->end() && !TII->analyzeBranch(*CurMBB, TBB, FBB, Cond, true)) {
     MachineBasicBlock *NextBB = &*I;
     if (TBB == NextBB && !Cond.empty() && !FBB) {
@@ -686,7 +688,8 @@ unsigned BranchFolder::ComputeSameTails(unsigned CurHash,
 
 void BranchFolder::RemoveBlocksWithHash(unsigned CurHash,
                                         MachineBasicBlock *SuccBB,
-                                        MachineBasicBlock *PredBB) {
+                                        MachineBasicBlock *PredBB,
+                                        const DebugLoc &BranchDL) {
   MPIterator CurMPIter, B;
   for (CurMPIter = std::prev(MergePotentials.end()),
       B = MergePotentials.begin();
@@ -694,7 +697,7 @@ void BranchFolder::RemoveBlocksWithHash(unsigned CurHash,
     // Put the unconditional branch back, if we need one.
     MachineBasicBlock *CurMBB = CurMPIter->getBlock();
     if (SuccBB && CurMBB != PredBB)
-      FixTail(CurMBB, SuccBB, TII);
+      FixTail(CurMBB, SuccBB, TII, BranchDL);
     if (CurMPIter == B)
       break;
   }
@@ -908,6 +911,7 @@ bool BranchFolder::TryTailMergeBlocks(MachineBasicBlock *SuccBB,
   // Walk through equivalence sets looking for actual exact matches.
   while (MergePotentials.size() > 1) {
     unsigned CurHash = MergePotentials.back().getHash();
+    const DebugLoc &BranchDL = MergePotentials.back().getBranchDebugLoc();
 
     // Build SameTails, identifying the set of blocks with this hash code
     // and with the maximum number of instructions in common.
@@ -918,7 +922,7 @@ bool BranchFolder::TryTailMergeBlocks(MachineBasicBlock *SuccBB,
     // If we didn't find any pair that has at least MinCommonTailLength
     // instructions in common, remove all blocks with this hash code and retry.
     if (SameTails.empty()) {
-      RemoveBlocksWithHash(CurHash, SuccBB, PredBB);
+      RemoveBlocksWithHash(CurHash, SuccBB, PredBB, BranchDL);
       continue;
     }
 
@@ -965,7 +969,7 @@ bool BranchFolder::TryTailMergeBlocks(MachineBasicBlock *SuccBB,
       // Split a block so that one does.
       if (!CreateCommonTailOnlyBlock(PredBB, SuccBB,
                                      maxCommonTailLength, commonTailIndex)) {
-        RemoveBlocksWithHash(CurHash, SuccBB, PredBB);
+        RemoveBlocksWithHash(CurHash, SuccBB, PredBB, BranchDL);
         continue;
       }
     }
@@ -1013,7 +1017,8 @@ bool BranchFolder::TailMergeBlocks(MachineFunction &MF) {
     if (MergePotentials.size() == TailMergeThreshold)
       break;
     if (!TriedMerging.count(&MBB) && MBB.succ_empty())
-      MergePotentials.push_back(MergePotentialsElt(HashEndOfMBB(MBB), &MBB));
+      MergePotentials.push_back(MergePotentialsElt(HashEndOfMBB(MBB), &MBB,
+                                                   MBB.findBranchDebugLoc()));
   }
 
   // If this is a large problem, avoid visiting the same basic blocks
@@ -1115,8 +1120,8 @@ bool BranchFolder::TailMergeBlocks(MachineFunction &MF) {
         }
 
         // Remove the unconditional branch at the end, if any.
+        DebugLoc dl = PBB->findBranchDebugLoc();
         if (TBB && (Cond.empty() || FBB)) {
-          DebugLoc dl = PBB->findBranchDebugLoc();
           TII->removeBranch(*PBB);
           if (!Cond.empty())
             // reinsert conditional branch only, for now
@@ -1124,7 +1129,8 @@ bool BranchFolder::TailMergeBlocks(MachineFunction &MF) {
                               NewCond, dl);
         }
 
-        MergePotentials.push_back(MergePotentialsElt(HashEndOfMBB(*PBB), PBB));
+        MergePotentials.push_back(
+            MergePotentialsElt(HashEndOfMBB(*PBB), PBB, dl));
       }
     }
 
@@ -1142,7 +1148,8 @@ bool BranchFolder::TailMergeBlocks(MachineFunction &MF) {
     PredBB = &*std::prev(I); // this may have been changed in TryTailMergeBlocks
     if (MergePotentials.size() == 1 &&
         MergePotentials.begin()->getBlock() != PredBB)
-      FixTail(MergePotentials.begin()->getBlock(), IBB, TII);
+      FixTail(MergePotentials.begin()->getBlock(), IBB, TII,
+              MergePotentials.begin()->getBranchDebugLoc());
   }
 
   return MadeChange;