wip debug location tail merging

Fixes #94050

Author: alanzhao1

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,15 +688,16 @@ 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();
        CurMPIter->getHash() == CurHash; --CurMPIter) {
     // 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,16 +1120,17 @@ 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
             TII->insertBranch(*PBB, (TBB == IBB) ? FBB : TBB, nullptr,
                               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;

llvm/lib/CodeGen/BranchFolding.h

@@ -50,10 +50,11 @@ class TargetRegisterInfo;
     class MergePotentialsElt {
       unsigned Hash;
       MachineBasicBlock *Block;
+      DebugLoc BranchDebugLoc;
 
     public:
-      MergePotentialsElt(unsigned h, MachineBasicBlock *b)
-        : Hash(h), Block(b) {}
+      MergePotentialsElt(unsigned h, MachineBasicBlock *b, DebugLoc bdl)
+          : Hash(h), Block(b), BranchDebugLoc(std::move(bdl)) {}
 
       unsigned getHash() const { return Hash; }
       MachineBasicBlock *getBlock() const { return Block; }
@@ -62,6 +63,8 @@ class TargetRegisterInfo;
         Block = MBB;
       }
 
+      const DebugLoc &getBranchDebugLoc() { return BranchDebugLoc; }
+
       bool operator<(const MergePotentialsElt &) const;
     };
 
@@ -162,8 +165,9 @@ class TargetRegisterInfo;
 
     /// Remove all blocks with hash CurHash from MergePotentials, restoring
     /// branches at ends of blocks as appropriate.
-    void RemoveBlocksWithHash(unsigned CurHash, MachineBasicBlock* SuccBB,
-                                                MachineBasicBlock* PredBB);
+    void RemoveBlocksWithHash(unsigned CurHash, MachineBasicBlock *SuccBB,
+                              MachineBasicBlock *PredBB,
+                              const DebugLoc &BranchDL);
 
     /// None of the blocks to be tail-merged consist only of the common tail.
     /// Create a block that does by splitting one.

llvm/test/CodeGen/X86/tail-merging-preserve-debugloc.ll

@@ -0,0 +1,31 @@
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -enable-tail-merge=1 -stop-after=branch-folder | FileCheck %s
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-grtev4-linux-gnu"
+
+define i32 @main(i1 %0) {
+entry:
+  br i1 %0, label %1, label %2
+
+1:                                                ; preds = %entry
+  store i64 1, ptr null, align 1
+; CHECK: JMP_1 %bb.3, debug-location !3
+  br label %3, !dbg !3
+
+2:                                                ; preds = %entry
+  store i64 0, ptr null, align 1
+  br label %3
+
+3:                                                ; preds = %2, %1
+  ret i32 0
+}
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!2}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1, producer: "", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly, nameTableKind: None)
+!1 = !DIFile(filename: "foo.c", directory: "/tmp", checksumkind: CSK_MD5, checksum: "2d07c91bb9d9c2fa4eee31a1aeed20e3")
+!2 = !{i32 2, !"Debug Info Version", i32 3}
+!3 = !DILocation(line: 17, column: 3, scope: !4)
+!4 = distinct !DISubprogram(name: "main", scope: !1, file: !1, line: 6, type: !5, scopeLine: 6, flags: DIFlagPrototyped | DIFlagAllCallsDescribed, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+!5 = !DISubroutineType(types: !6)
+!6 = !{}