From 857c25bd343c5e8216301e689cdaac2ebd68e58f Mon Sep 17 00:00:00 2001
From: Stephen Tozer <stephen.tozer@sony.com>
Date: Fri, 27 Sep 2024 15:24:13 +0100
Subject: [PATCH] [DebugInfo][LoopUnroll] Preserve DebugLocs on optimized cond
 branches

This patch fixes a simple error where as part of loop unrolling we optimize
conditional loop-exiting branches into unconditional branches when we know that
they will or won't exit the loop, but does not propagate the source location of
the original branch to the new one.
---
 llvm/lib/Transforms/Utils/LoopUnroll.cpp      |  3 +-
 .../LoopUnroll/preserve-branch-debuglocs.ll   | 52 +++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/LoopUnroll/preserve-branch-debuglocs.ll

diff --git a/llvm/lib/Transforms/Utils/LoopUnroll.cpp b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
index b90addcef69e6..4cbbac07915fb 100644
--- a/llvm/lib/Transforms/Utils/LoopUnroll.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
@@ -883,7 +883,8 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
     DeadSucc->removePredecessor(Src, /* KeepOneInputPHIs */ true);
 
     // Replace the conditional branch with an unconditional one.
-    BranchInst::Create(Dest, Term->getIterator());
+    auto *BI = BranchInst::Create(Dest, Term->getIterator());
+    BI->setDebugLoc(Term->getDebugLoc());
     Term->eraseFromParent();
 
     DTUpdates.emplace_back(DominatorTree::Delete, Src, DeadSucc);
diff --git a/llvm/test/Transforms/LoopUnroll/preserve-branch-debuglocs.ll b/llvm/test/Transforms/LoopUnroll/preserve-branch-debuglocs.ll
new file mode 100644
index 0000000000000..623fb68c784fe
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/preserve-branch-debuglocs.ll
@@ -0,0 +1,52 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 5
+; RUN: opt < %s -passes=loop-unroll -S | FileCheck %s
+; RUN: opt < %s -passes=loop-unroll-full -S | FileCheck %s
+
+;; Loop Unrolling should preserve the DebugLocs of conditional branches that get
+;; optimized into unconditional branches.
+
+define i32 @_ZeqRK4int3S1_() {
+; CHECK-LABEL: define i32 @_ZeqRK4int3S1_() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    br label %[[FOR_BODY:.*]]
+; CHECK:       [[FOR_COND:.*]]:
+; CHECK-NEXT:    br label %[[CLEANUP:.*]], !dbg [[DBG3:![0-9]+]]
+; CHECK:       [[FOR_BODY]]:
+; CHECK-NEXT:    br i1 false, label %[[FOR_COND]], label %[[CLEANUP]]
+; CHECK:       [[CLEANUP]]:
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  br label %for.body
+
+for.cond:                                         ; preds = %for.body
+  br i1 false, label %cleanup, label %for.body, !dbg !3
+
+for.body:                                         ; preds = %for.cond, %entry
+  br i1 false, label %for.cond, label %cleanup
+
+cleanup:                                          ; preds = %for.body, %for.cond
+  ret i32 0
+}
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!2}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus, file: !1, producer: "clang version 20.0.0git", isOptimized: true, runtimeVersion: 0, emissionKind: NoDebug)
+!1 = !DIFile(filename: "test.cpp", directory: "/tmp")
+!2 = !{i32 2, !"Debug Info Version", i32 3}
+!3 = !DILocation(line: 304, column: 2, scope: !4)
+!4 = distinct !DILexicalBlock(scope: !5, file: !1, line: 304, column: 2)
+!5 = distinct !DISubprogram(name: "operator==", linkageName: "_ZeqRK4int3S1_", scope: !1, file: !1, line: 302, type: !6, scopeLine: 303, spFlags: DISPFlagDefinition, unit: !0)
+!6 = distinct !DISubroutineType(types: !7)
+!7 = !{}
+;.
+; CHECK: [[META0:![0-9]+]] = distinct !DICompileUnit(language: DW_LANG_C_plus_plus, file: [[META1:![0-9]+]], producer: "{{.*}}clang version {{.*}}", isOptimized: true, runtimeVersion: 0, emissionKind: NoDebug)
+; CHECK: [[META1]] = !DIFile(filename: "test.cpp", directory: {{.*}})
+; CHECK: [[META2:![0-9]+]] = !{i32 2, !"Debug Info Version", i32 3}
+; CHECK: [[DBG3]] = !DILocation(line: 304, column: 2, scope: [[META4:![0-9]+]])
+; CHECK: [[META4]] = distinct !DILexicalBlock(scope: [[META5:![0-9]+]], file: [[META1]], line: 304, column: 2)
+; CHECK: [[META5]] = distinct !DISubprogram(name: "operator==", linkageName: "_ZeqRK4int3S1_", scope: [[META1]], file: [[META1]], line: 302, type: [[META6:![0-9]+]], scopeLine: 303, spFlags: DISPFlagDefinition, unit: [[META0]])
+; CHECK: [[META6]] = distinct !DISubroutineType(types: [[META7:![0-9]+]])
+; CHECK: [[META7]] = !{}
+;.